加载预训练模型和分词器。 Args: generation_config：模型配置参数。 Returns: model (Transformers模型): 预训练模型。 tokenizer (Transformers分词器): 分词器。 llm (CookMasterLLM): langchain封装的大模型。
(generation_config)
| 61 | |
| 62 | @st.cache_resource |
| 63 | def load_model(generation_config): |
| 64 | """ |
| 65 | 加载预训练模型和分词器。 |
| 66 | |
| 67 | Args: |
| 68 | generation_config：模型配置参数。 |
| 69 | |
| 70 | Returns: |
| 71 | model (Transformers模型): 预训练模型。 |
| 72 | tokenizer (Transformers分词器): 分词器。 |
| 73 | llm (CookMasterLLM): langchain封装的大模型。 |
| 74 | """ |
| 75 | |
| 76 | if load_4bit == False: |
| 77 | |
| 78 | model = ( |
| 79 | AutoModelForCausalLM.from_pretrained(llm_model_path, trust_remote_code=True) |
| 80 | .to(torch.bfloat16) |
| 81 | .cuda() |
| 82 | ) |
| 83 | tokenizer = AutoTokenizer.from_pretrained(llm_model_path, trust_remote_code=True) |
| 84 | |
| 85 | else: |
| 86 | # int4 量化加载 |
| 87 | quantization_config = BitsAndBytesConfig( |
| 88 | load_in_4bit=True, |
| 89 | bnb_4bit_compute_dtype=torch.float16, |
| 90 | bnb_4bit_quant_type="nf4", |
| 91 | bnb_4bit_use_double_quant=True, |
| 92 | ) |
| 93 | print("正在从本地加载模型...") |
| 94 | model = AutoModelForCausalLM.from_pretrained(llm_model_path, trust_remote_code=True, torch_dtype=torch.float16, |
| 95 | device_map="auto", |
| 96 | quantization_config=quantization_config) |
| 97 | tokenizer = AutoTokenizer.from_pretrained(llm_model_path, trust_remote_code=True) |
| 98 | |
| 99 | llm = CookMasterLLM(model, tokenizer) |
| 100 | print("完成本地模型的加载") |
| 101 | model.generation_config.max_length = generation_config.max_length |
| 102 | model.generation_config.top_p = generation_config.top_p |
| 103 | model.generation_config.temperature = generation_config.temperature |
| 104 | model.generation_config.repetition_penalty = generation_config.repetition_penalty |
| 105 | # print(model.generation_config) |
| 106 | return model, tokenizer, llm |
| 107 | |
| 108 | |
| 109 | def combine_history(prompt): |