Merge pull request #60 from kuun993/main

修复关于qwen的一些代码错误

llm/LLM.py

@@ -26,15 +26,15 @@ class LLM:
         elif model_name == 'ChatGPT':
             llm = ChatGPT(model_path, api_key=api_key)
         elif model_name == 'Qwen':
-            llm = Qwen(self.mode, model_path)
+            llm = Qwen(model_path=model_path, api_key=api_key, api_base=proxy_url)
         elif model_name == 'VllmGPT':
             llm = VllmGPT()
         return llm
 
 
-    def test_Qwen(self, question="如何应对压力？", model_path="Qwen/Qwen-1_8B-Chat"):
-        llm = Qwen(self.mode, model_path)
-        answer = llm.generate(question)
+    def test_Qwen(self, question="如何应对压力？", model_path="Qwen/Qwen-1_8B-Chat", api_key=None, proxy_url=None):
+        llm = Qwen(model_path=model_path, api_key=api_key, api_base=proxy_url)
+        answer = llm.chat(question)
         print(answer)
 
     def test_Gemini(self, question="如何应对压力？", model_path='gemini-pro', api_key=None, proxy_url=None):
@@ -43,10 +43,12 @@ class LLM:
         print(answer)
 
 if __name__ == '__main__':
-    # llm = LLM()
+    llm = LLM()
     # llm.test_Gemini(api_key='你的API Key', proxy_url=None)
     # llm = LLM().init_model('Gemini', model_path= 'gemini-pro',api_key='AIzaSyBWAWfT8zsyAZcRIXLS5Vzlw8KKCN9qsAg', proxy_url='http://172.31.71.58:7890')
     # response = llm.chat("如何应对压力？")
-    llm = LLM().init_model('VllmGPT', model_path= 'THUDM/chatglm3-6b')
-    response = llm.chat("如何应对压力？")
+    # llm = LLM().init_model('VllmGPT', model_path= 'THUDM/chatglm3-6b')
+    # response = llm.chat("如何应对压力？")
     # print(response)
+
+    llm.test_Qwen(api_key="none", proxy_url="http://10.1.1.113:18000/v1")

llm/Qwen.py

@@ -1,15 +1,33 @@
 import os
-import torch
-import requests
-from transformers import AutoModelForCausalLM, AutoTokenizer
+import openai
+
+'''
+`huggingface`连接不上可以使用 `modelscope`
+`pip install modelscope`
+'''
+from modelscope import AutoModelForCausalLM, AutoTokenizer
+#from transformers import AutoModelForCausalLM, AutoTokenizer
+
 os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
 
-class Qwen:
-    def __init__(self,  model_path="Qwen/Qwen-1_8B-Chat") -> None:
-        '''暂时不写api版本,与Linly-api相类似,感兴趣可以实现一下'''
-        self.model, self.tokenizer = self.init_model(model_path)
 
-    def init_model(self, path = "Qwen/Qwen-1_8B-Chat"):
+class Qwen:
+    def __init__(self, model_path="Qwen/Qwen-1_8B-Chat", api_base=None, api_key=None) -> None:
+        '''暂时不写api版本,与Linly-api相类似,感兴趣可以实现一下'''
+        # 默认本地推理
+        self.local = True
+
+        # api_base和api_key不为空时使用openapi的方式
+        if api_key is not None and api_base is not None:
+            openai.api_base = api_base
+            openai.api_key = api_key
+            self.local = False
+            return
+
+        self.model, self.tokenizer = self.init_model(model_path)
+        self.data = {}
+
+    def init_model(self, path="Qwen/Qwen-1_8B-Chat"):
         model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-1_8B-Chat",
                                                      device_map="auto",
                                                      trust_remote_code=True).eval()
@@ -18,8 +36,21 @@ class Qwen:
         return model, tokenizer
 
     def chat(self, question):
+        # 优先调用qwen openapi的方式
+        if not self.local:
+            # 不使用流式回复的请求
+            response = openai.ChatCompletion.create(
+                model="Qwen",
+                messages=[
+                    {"role": "user", "content": question}
+                ],
+                stream=False,
+                stop=[]
+            )
+            return response.choices[0].message.content
 
-        self.data["question"] = f"{self.prompt} ### Instruction:{question}  ### Response:"
+        # 默认本地推理
+        self.data["question"] = f"{question} ### Instruction:{question}  ### Response:"
         try:
             response, history = self.model.chat(self.tokenizer, self.data["question"], history=None)
             print(history)
@@ -30,8 +61,9 @@ class Qwen:
 
 def test():
     llm = Qwen(model_path="Qwen/Qwen-1_8B-Chat")
-    answer = llm.generate("如何应对压力？")
+    answer = llm.chat(question="如何应对压力？")
     print(answer)
 
+
 if __name__ == '__main__':
     test()