Update app.py

app.py

@@ -14,13 +14,15 @@ from langchain.schema import Document
 from langchain.embeddings import OpenAIEmbeddings
 from langchain.embeddings.huggingface import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
-from langchain.chains import ConversationalRetrievalChain,RetrievalQA
+from langchain.chains import ConversationalRetrievalChain,RetrievalQA,LLMChain
 from langchain.prompts import PromptTemplate
 from langchain.prompts.prompt import PromptTemplate
 from langchain.chat_models import ChatOpenAI
+from langchain import OpenAI,VectorDBQA
 
 def load_model():
     tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
+    # gpu:.half().cuda()
     model = AutoModel.from_pretrained("THUDM/chatglm-6b-int4", trust_remote_code=True).quantize(bits=4, compile_parallel_kernel=True, parallel_num=2).float()
     model = model.eval()
     return tokenizer,model
@@ -32,11 +34,11 @@ def chat_glm(input, history=None):
 
     tokenizer,model = load_model()
     response, history = model.chat(tokenizer, input, history)
-    logger.info("chatglm:", input,response)
+    logger.debug("chatglm:", input,response)
     return history, history
 
 def search_web(query):
-    logger.info("searchweb:", query)
+    logger.debug("searchweb:", query)
     results = ddg(query)
     web_content = ''
     if results:
@@ -44,55 +46,46 @@ def search_web(query):
             web_content += result['body']
     return web_content
 
-def chat_gpt(input, use_web, history=None):
-    if history is None:
-        history = []
-    history = [] # 4097 tokens limit
+def search_vec(query):
+    logger.debug("searchvec:", query)
     embedding_model_name = 'GanymedeNil/text2vec-large-chinese'
     vec_path = 'cache'
     embeddings = HuggingFaceEmbeddings(model_name=embedding_model_name)
-
-    if use_web:
-        web_content = search_web(input)
-    else:
-        web_content = None
-    web_content = None  # 4097 tokens limit
-    if web_content:
-        prompt_template = f"""基于以下已知信息，简洁和专业的来回答用户的问题。
-                            如果无法从中得到答案，请说 "根据已知信息无法回答该问题" 或 "没有提供足够的相关信息"，不允许在答案中添加编造成分，答案请使用中文。
-                            已知网络检索内容：{web_content}""" + """
-                            已知内容:
-                            {context}
-                            问题:
-                            {question}"""
-    else:
-        prompt_template = """基于以下已知信息，请简洁并专业地回答用户的问题。
-            如果无法从中得到答案，请说 "根据已知信息无法回答该问题" 或 "没有提供足够的相关信息"。不允许在答案中添加编造成分。另外，答案请使用中文。
-            已知内容:
-            {context}
-            问题:
-            {question}"""
-            
-    prompt = PromptTemplate(template=prompt_template,input_variables=["context", "question"])
     vector_store = FAISS.load_local(vec_path,embeddings)
+    
+    qa = VectorDBQA.from_chain_type(llm=OpenAI(), chain_type="stuff", vectorstore=vector_store,return_source_documents=True)
+    result = qa({"query": query})
+    return result
 
-    qa = RetrievalQA.from_llm(
-        llm = ChatOpenAI(temperature=0.7,model_name='gpt-3.5-turbo'),
-        retriever = vector_store.as_retriever(search_kwargs={"k": 3}),
-        prompt = prompt,
-        return_source_documents=True
-    )
+def chat_gpt(input, use_web, use_vec, history=None):
+    if history is None:
+        history = []
+    # history = [] # 4097 tokens limit
     
-    result = qa({"query": input, "chat_history": history})
-    logger.info("chatgpt:", input,result)
-    return result["result"]
+    context = "无"
+    if use_vec:
+        context = search_vec(input)['result']
+    prompt_template = f"""基于以下已知信息，请简洁并专业地回答用户的问题。
+        如果无法从中得到答案，请说 "根据已知信息无法回答该问题" 或 "没有提供足够的相关信息"。若答案中存在编造成分，请在该部分开头添加“据我推测”。另外，答案请使用中文。
+        已知内容:
+        {context}"""+"""
+        问题:
+        {question}"""
+            
+    prompt = PromptTemplate(template=prompt_template,input_variables=["question"])
+    
+    llm = OpenAI(temperature = 0.2)
+    chain = LLMChain(llm=llm, prompt=prompt)
+    result = chain.run(text)
+    return result
 
 def predict(input,
             large_language_model,
             use_web,
+            use_vec,
             openai_key,
             history=None):
-    logger.info("predict..",large_language_model,use_web)
+    logger.debug("predict..",large_language_model,use_web)
     if openai_key is not None:
         os.environ['OPENAI_API_KEY'] = openai_key
     else:
@@ -101,11 +94,14 @@ def predict(input,
         history = []
 
     if large_language_model == "GPT-3.5-turbo":
-        resp  = chat_gpt(input, use_web, history)
+        resp  = chat_gpt(input, use_web, use_vec, history)
     elif large_language_model == "ChatGLM-6B-int4":
-        resp = chat_glm(input, history)
+        _,resp = chat_glm(input, history)
+        resp = resp[-1][1]
     elif large_language_model == "Search Web":
         resp = search_web(input)
+    elif large_language_model == "Search VectorStore":
+        resp = search_vec(input)
 
     history.append((input, resp))
     return '', history, history
@@ -125,12 +121,15 @@ with block as demo:
             model_choose = gr.Accordion("模型选择")
             with model_choose:
                 large_language_model = gr.Dropdown(
-                    ["ChatGLM-6B-int4","GPT-3.5-turbo","Search Web"],
+                    ["ChatGLM-6B-int4","GPT-3.5-turbo","Search Web","Search VectorStore"],
                     label="large language model",
                     value="ChatGLM-6B-int4")
             use_web = gr.Radio(["True", "False"],
                     label="Web Search",
                     value="False")
+            use_vec = gr.Radio(["True", "False"],
+                    label="VectorStore Search",
+                    value="False")
             openai_key = gr.Textbox(label="请输入OpenAI API key", type="password")
         with gr.Column(scale=4):
             chatbot = gr.Chatbot(label='ChatLLM').style(height=600)
@@ -143,7 +142,7 @@ with block as demo:
 
                 send.click(predict,
                            inputs=[
-                               message, large_language_model, use_web, openai_key, state
+                               message, large_language_model, use_web, use_vec, openai_key, state
                            ],
                            outputs=[message, chatbot, state])
                 clear_history.click(fn=clear_session,
@@ -153,7 +152,7 @@ with block as demo:
 
                 message.submit(predict,
                                inputs=[
-                                   message, large_language_model, use_web, openai_key, state
+                                   message, large_language_model, use_web, use_vec, openai_key, state
                                ],
                                outputs=[message, chatbot, state])
     gr.Markdown("""提醒：<br>