(model_name, cache_dir="llm_weights")
| 279 | } |
| 280 | |
| 281 | def get_llm(model_name, cache_dir="llm_weights"): |
| 282 | model = AutoModelForCausalLM.from_pretrained( |
| 283 | modeltype2path[model_name], |
| 284 | torch_dtype=torch.bfloat16, |
| 285 | cache_dir=cache_dir, |
| 286 | low_cpu_mem_usage=True, |
| 287 | device_map="cuda", |
| 288 | ) |
| 289 | |
| 290 | model.seqlen = model.config.max_position_embeddings |
| 291 | return model |
| 292 | |
| 293 | tokenizer = AutoTokenizer.from_pretrained(modeltype2path["llama2-7b-hf"]) |
| 294 | inputs = tokenizer( |