Fariddvp commited on
Commit
a4cb129
·
verified ·
1 Parent(s): 352e637

Update models.py

Browse files
Files changed (1) hide show
  1. models.py +4 -2
models.py CHANGED
@@ -133,12 +133,14 @@ login(token=hf_api_key)
133
 
134
 
135
 
136
-
137
  model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
 
 
 
138
 
139
  pipeline = transformers.pipeline(
140
  "text-generation",
141
- model=model_id,
142
  model_kwargs={"torch_dtype": torch.bfloat16},
143
  device_map="auto",
144
  )
 
133
 
134
 
135
 
 
136
  model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
137
+ model_config = transformers.LlamaConfig.from_pretrained(model_id)
138
+ model_config.rope_scaling = {"type": "linear", "factor": 1.0}
139
+ model = transformers.AutoModelForCausalLM.from_pretrained(model_id, config=model_config)
140
 
141
  pipeline = transformers.pipeline(
142
  "text-generation",
143
+ model=model,
144
  model_kwargs={"torch_dtype": torch.bfloat16},
145
  device_map="auto",
146
  )