Upload Linly.py

LLM/Linly.py

@@ -0,0 +1,89 @@
+import os
+import torch
+import requests
+import json
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from configs import ip, api_port, model_path
+os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
+
+class Linly:
+    def __init__(self, mode='api', model_path="Linly-AI/Chinese-LLaMA-2-7B-hf", prefix_prompt = '''请用少于25个字回答以下问题\n\n'''):
+        # mode = api need
+        # 定义设置的api的服务器,首先记得运行Linly-api-fast.py 填入ip地址和端口号
+        self.url = f"http://{ip}:{api_port}" # local server: http://ip:port
+        self.headers = {
+            "Content-Type": "application/json"
+        }
+        self.data = {
+            "question": "北京有什么好玩的地方？"
+        }
+        # 全局设定的prompt
+        self.prefix_prompt = prefix_prompt
+        self.mode = mode
+        if mode != 'api':
+            self.model, self.tokenizer = self.init_model(model_path)
+        self.history = []
+    
+    def init_model(self, path = "Linly-AI/Chinese-LLaMA-2-7B-hf"):
+        model = AutoModelForCausalLM.from_pretrained(path, device_map="cuda:0",
+                                                    torch_dtype=torch.bfloat16, trust_remote_code=True)
+        tokenizer = AutoTokenizer.from_pretrained(path, use_fast=False, trust_remote_code=True)
+        return model, tokenizer   
+    
+    def generate(self, question, system_prompt=""):
+        if self.mode != 'api':
+            self.data["question"] = self.message_to_prompt(question, system_prompt)
+            inputs = self.tokenizer(self.data["question"], return_tensors="pt").to("cuda:0")
+            try:
+                generate_ids = self.model.generate(inputs.input_ids, 
+                                                   max_new_tokens=2048, 
+                                                   do_sample=True, 
+                                                   top_k=20, 
+                                                   top_p=0.84,
+                                                   temperature=1, 
+                                                   repetition_penalty=1.15, 
+                                                   eos_token_id=2, 
+                                                   bos_token_id=1,
+                                                   pad_token_id=0)
+                response = self.tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+                response = response.split("### Response:")[-1]
+                return response
+            except:
+                return "对不起，你的请求出错了，请再次尝试。\nSorry, your request has encountered an error. Please try again.\n"
+        elif self.mode == 'api':
+            return self.predict_api(question)
+    
+    def message_to_prompt(self, message, system_prompt=""):
+        system_prompt = self.prefix_prompt + system_prompt
+        for interaction in self.history:
+            user_prompt, bot_prompt = str(interaction[0]).strip(' '), str(interaction[1]).strip(' ')
+            system_prompt = f"{system_prompt} User: {user_prompt} Bot: {bot_prompt}"
+        prompt = f"{system_prompt} ### Instruction:{message.strip()}  ### Response:"
+        return prompt
+        
+    def predict_api(self, question):
+        # FastAPI Predict 调用API来进行预测
+        self.data["question"] = question
+        headers = {'Content-Type': 'application/json'}
+        data = {"prompt": question}
+        response = requests.post(url=self.url, headers=headers, data=json.dumps(data))
+        return response.json()['response']
+            
+    def chat(self, system_prompt, message, history):
+        self.history = history
+        prompt = self.message_to_prompt(message, system_prompt)
+        response = self.generate(prompt)
+        self.history.append([message, response])
+        return response, self.history
+    
+    def clear_history(self):
+        # 清空历史记录
+        self.history = []
+
+def test():
+    llm = Linly(mode='offline',model_path='../Linly-AI/Chinese-LLaMA-2-7B-hf')
+    answer = llm.generate("如何应对压力？")
+    print(answer)
+
+if __name__ == '__main__':
+    test()