pmelnechuk commited on
Commit
881820d
verified
1 Parent(s): 1fffd44

Update src/model_load.py

Browse files
Files changed (1) hide show
  1. src/model_load.py +2 -2
src/model_load.py CHANGED
@@ -9,13 +9,13 @@ from langchain.chains import LLMChain
9
  def load_model():
10
 
11
  model_name="tiiuae/Falcon3-10B-Instruct"
12
- max_memory = {0: "23GB", "cpu": "30GB"}
13
  # Cargar tokenizer y modelo de Hugging Face
14
  tokenizer = AutoTokenizer.from_pretrained(model_name)
15
  model = AutoModelForCausalLM.from_pretrained(model_name,
16
  device_map="auto",
17
  quantization_config=BitsAndBytesConfig(load_in_8bit=True),
18
- torch_dtype=torch.float8,
19
  max_memory=max_memory)
20
 
21
  # Crear pipeline de generaci贸n de texto
 
9
  def load_model():
10
 
11
  model_name="tiiuae/Falcon3-10B-Instruct"
12
+ max_memory = {0: "24GB", "cpu": "30GB"}
13
  # Cargar tokenizer y modelo de Hugging Face
14
  tokenizer = AutoTokenizer.from_pretrained(model_name)
15
  model = AutoModelForCausalLM.from_pretrained(model_name,
16
  device_map="auto",
17
  quantization_config=BitsAndBytesConfig(load_in_8bit=True),
18
+ torch_dtype=torch.float16,
19
  max_memory=max_memory)
20
 
21
  # Crear pipeline de generaci贸n de texto