|
@@ -7,25 +7,25 @@ import torch
|
|
|
|
|
|
|
|
|
|
class ChatGLM_LLM(LLM):
|
|
class ChatGLM_LLM(LLM):
|
|
- # ���ڱ��� InternLM �Զ��� LLM ��
|
|
|
|
|
|
+ # 基于本地 InternLM 自定义 LLM 类
|
|
tokenizer: AutoTokenizer = None
|
|
tokenizer: AutoTokenizer = None
|
|
model: AutoModelForCausalLM = None
|
|
model: AutoModelForCausalLM = None
|
|
|
|
|
|
def __init__(self, model_path: str):
|
|
def __init__(self, model_path: str):
|
|
- # model_path: InternLM ģ��·��
|
|
|
|
- # �ӱ��س�ʼ��ģ��
|
|
|
|
|
|
+ # model_path: InternLM 模型路径
|
|
|
|
+ # 从本地初始化模型
|
|
super().__init__()
|
|
super().__init__()
|
|
- print("���ڴӱ��ؼ���ģ��...")
|
|
|
|
|
|
+ print("正在从本地加载模型...")
|
|
self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
|
self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
|
self.model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True).to(torch.bfloat16).cuda(
|
|
self.model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True).to(torch.bfloat16).cuda(
|
|
device=1)
|
|
device=1)
|
|
self.model = self.model.eval()
|
|
self.model = self.model.eval()
|
|
- print("��ɱ���ģ�͵ļ���")
|
|
|
|
|
|
+ print("完成本地模型的加载")
|
|
|
|
|
|
def _call(self, prompt: str, stop: Optional[List[str]] = None,
|
|
def _call(self, prompt: str, stop: Optional[List[str]] = None,
|
|
run_manager: Optional[CallbackManagerForLLMRun] = None,
|
|
run_manager: Optional[CallbackManagerForLLMRun] = None,
|
|
**kwargs: Any):
|
|
**kwargs: Any):
|
|
- # ��д���ú���
|
|
|
|
|
|
+ # 重写调用函数
|
|
response, history = self.model.chat(self.tokenizer, prompt, history=[], do_sample=False)
|
|
response, history = self.model.chat(self.tokenizer, prompt, history=[], do_sample=False)
|
|
return response
|
|
return response
|
|
|
|
|