(model_name, cache_dir="llm_weights")
| 332 | } |
| 333 | |
| 334 | def get_llm(model_name, cache_dir="llm_weights"): |
| 335 | model = AutoModelForCausalLM.from_pretrained( |
| 336 | modeltype2path[model_name], |
| 337 | torch_dtype=torch.bfloat16, |
| 338 | cache_dir=cache_dir, |
| 339 | low_cpu_mem_usage=True, |
| 340 | device_map="cuda", |
| 341 | ) |
| 342 | |
| 343 | model.seqlen = model.config.max_position_embeddings |
| 344 | return model |
| 345 | |
| 346 | if args.model == "llama2-7b-chat-hf": |
| 347 | tokenizer = AutoTokenizer.from_pretrained(modeltype2path["llama2-7b-chat-hf"]) |