pmelnechuk commited on
Commit
579ba27
verified
1 Parent(s): 3571f6c

Update src/model_load.py

Browse files
Files changed (1) hide show
  1. src/model_load.py +2 -1
src/model_load.py CHANGED
@@ -14,7 +14,8 @@ def load_model():
14
  tokenizer = AutoTokenizer.from_pretrained(model_name)
15
  model = AutoModelForCausalLM.from_pretrained(model_name,
16
  torch_dtype=torch.float16,
17
- device_map="auto")
 
18
 
19
  # Crear pipeline de generaci贸n de texto
20
  text_generation_pipeline = pipeline(
 
14
  tokenizer = AutoTokenizer.from_pretrained(model_name)
15
  model = AutoModelForCausalLM.from_pretrained(model_name,
16
  torch_dtype=torch.float16,
17
+ device_map="auto",
18
+ quantization_config=BitsAndBytesConfig(load_in_8bit=True))
19
 
20
  # Crear pipeline de generaci贸n de texto
21
  text_generation_pipeline = pipeline(