(model_name, cache_dir="llm_weights")
| 25 | } |
| 26 | |
| 27 | def get_llm(model_name, cache_dir="llm_weights"): |
| 28 | model = AutoModelForCausalLM.from_pretrained( |
| 29 | modeltype2path[model_name], |
| 30 | torch_dtype=torch.bfloat16, |
| 31 | cache_dir=cache_dir, |
| 32 | low_cpu_mem_usage=True, |
| 33 | device_map="auto" |
| 34 | ) |
| 35 | |
| 36 | model.seqlen = model.config.max_position_embeddings |
| 37 | return model |
| 38 | |
| 39 | def main(): |
| 40 | parser = argparse.ArgumentParser() |