Update README.md
Browse files
README.md
CHANGED
@@ -45,7 +45,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
|
|
45 |
class TextGenerationAssistant:
|
46 |
def __init__(self, model_id:str):
|
47 |
self.tokenizer = AutoTokenizer.from_pretrained(model_id)
|
48 |
-
self.model = AutoModelForCausalLM.from_pretrained(model_id, device_map='auto',load_in_8bit=
|
49 |
self.pipe = pipeline("text-generation",
|
50 |
model=self.model,
|
51 |
tokenizer=self.tokenizer,
|
|
|
45 |
class TextGenerationAssistant:
|
46 |
def __init__(self, model_id:str):
|
47 |
self.tokenizer = AutoTokenizer.from_pretrained(model_id)
|
48 |
+
self.model = AutoModelForCausalLM.from_pretrained(model_id, device_map='auto',load_in_8bit=True,load_in_4bit=False)
|
49 |
self.pipe = pipeline("text-generation",
|
50 |
model=self.model,
|
51 |
tokenizer=self.tokenizer,
|