Function
get_llm
(model_name, cache_dir="llm_weights")
Source from the content-addressed store, hash-verified
| 40 | |
| 41 | |
| 42 | def get_llm(model_name, cache_dir="llm_weights"): |
| 43 | if model_name in [ |
| 44 | "llama2-7b-chat-hf", |
| 45 | "llama2-13b-chat-hf", |
| 46 | "llama2-7b-hf", |
| 47 | "llama2-13b-hf", |
| 48 | ]: |
| 49 | model = AutoModelForCausalLM.from_pretrained( |
| 50 | modeltype2path[model_name], |
| 51 | torch_dtype=torch.bfloat16, |
| 52 | cache_dir=cache_dir, |
| 53 | low_cpu_mem_usage=True, |
| 54 | device_map="auto", |
| 55 | ) |
| 56 | |
| 57 | model.seqlen = model.config.max_position_embeddings |
| 58 | return model |
| 59 | |
| 60 | |
| 61 | def main(): |
Tested by
no test coverage detected