()
| 181 | |
| 182 | @st.cache_resource |
| 183 | def load_model(): |
| 184 | model = (AutoModelForCausalLM.from_pretrained( |
| 185 | '/share/new_models/Shanghai_AI_Laboratory/internlm2-chat-1_8b', |
| 186 | trust_remote_code=True).to(torch.bfloat16).cuda()) |
| 187 | tokenizer = AutoTokenizer.from_pretrained( |
| 188 | '/share/new_models/Shanghai_AI_Laboratory/internlm2-chat-1_8b', |
| 189 | trust_remote_code=True) |
| 190 | return model, tokenizer |
| 191 | |
| 192 | |
| 193 | def prepare_generation_config(): |