(model_name, cache_dir="llm_weights")
| 532 | } |
| 533 | |
| 534 | def get_llm(model_name, cache_dir="llm_weights"): |
| 535 | model = AutoModelForCausalLM.from_pretrained( |
| 536 | modeltype2path[model_name], |
| 537 | torch_dtype=torch.bfloat16, |
| 538 | cache_dir=cache_dir, |
| 539 | low_cpu_mem_usage=True, |
| 540 | device_map="cuda", |
| 541 | ) |
| 542 | |
| 543 | model.seqlen = model.config.max_position_embeddings |
| 544 | return model |
| 545 | |
| 546 | tokenizer = AutoTokenizer.from_pretrained(modeltype2path["llama2-7b-hf"]) |
| 547 | inputs = tokenizer( |