()
| 98 | |
| 99 | |
| 100 | def load_model(): |
| 101 | llm_model_path = load_config('llm', 'llm_model_path') |
| 102 | base_model_type = load_config('llm', 'base_model_type') |
| 103 | print(f"base model type:{base_model_type}") |
| 104 | max_length = 32768 |
| 105 | if base_model_type == 'internlm-chat-7b': |
| 106 | generation_config = GenerationConfig( |
| 107 | max_length=max_length) # InternLM1 |
| 108 | elif base_model_type == 'internlm2-chat-1.8b': |
| 109 | generation_config = GenerationConfig( |
| 110 | max_length=max_length, top_p=0.8, temperature=0.8, repetition_penalty=1.17) # InternLM2 1.8b need 惩罚参数 |
| 111 | else: |
| 112 | generation_config = GenerationConfig( |
| 113 | max_length=max_length, top_p=0.8, temperature=0.8, repetition_penalty=1.002) # InternLM2 2 need 惩罚参数 |
| 114 | # int4 量化加载 |
| 115 | quantization_config = BitsAndBytesConfig( |
| 116 | load_in_4bit=True, |
| 117 | bnb_4bit_compute_dtype=torch.float16, |
| 118 | bnb_4bit_quant_type="nf4", |
| 119 | bnb_4bit_use_double_quant=True, |
| 120 | ) |
| 121 | print("正在从本地加载模型...") |
| 122 | model = AutoModelForCausalLM.from_pretrained(llm_model_path, trust_remote_code=True, torch_dtype=torch.float16, |
| 123 | device_map="auto", |
| 124 | quantization_config=quantization_config).eval() |
| 125 | tokenizer = AutoTokenizer.from_pretrained(llm_model_path, trust_remote_code=True) |
| 126 | llm = CookMasterLLM(model, tokenizer) |
| 127 | model.generation_config.max_length = generation_config.max_length |
| 128 | model.generation_config.top_p = generation_config.top_p |
| 129 | model.generation_config.temperature = generation_config.temperature |
| 130 | model.generation_config.repetition_penalty = generation_config.repetition_penalty |
| 131 | print(model.generation_config) |
| 132 | print("完成本地模型的加载") |
| 133 | return model, tokenizer, llm |
no test coverage detected