Update app.py

app.py

@@ -1,52 +1,91 @@
-import os
-os.system('pip install openai')
-os.system('pip install python-docx')
+# 长文本总结
+from langchain.chains.summarize import load_summarize_chain
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain import OpenAI
+from langchain import PromptTemplate
+from langchain.docstore.document import Document as LangDoc
+
+
+text_splitter = RecursiveCharacterTextSplitter(chunk_size=1500, chunk_overlap=100)
+
+
 import openai
 from openai.error import RateLimitError
+from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain.vectorstores import Chroma
+from langchain.prompts.chat import (
+    ChatPromptTemplate,
+    SystemMessagePromptTemplate,
+    HumanMessagePromptTemplate
+)
+from langchain.chains import ConversationalRetrievalChain
+from langchain.chat_models import ChatOpenAI
+from langchain.chains import RetrievalQA
+from langchain.output_parsers import RegexParser
 
 
 openai.api_key = ""
-prompt = """
-我需要你帮助我完成会议记录的总结和梳理。我会给你提供一长段会议记录，其中包含多个参会者的对话记录。你需要完成如下的任务：
-1. 提取出会议主要讨论的几个主题或者要点。
-2. 针对会议的每一个主题，请你找到所有相关的会议记录，并将它们梳理总结，为我生成有条理的摘要。
-
-在输出时，请注意以下几点：
-1. 输出内容中避免口语化内容
-2. 输出格式是层级列表
+
+summary_prompt = (
+"总结以下会议记录中所探讨的主要话题，忽略细节\n"
+"会议记录：{text}\n"
+"在输出时，请注意以下几点：\n"
+"1. 输出内容中避免口语化内容\n"
+"2. 每个话题用序号标注\n"
+"3. 不输出无关信息"
+)
+
+qa_prompt = """
+结合下面的信息，用中文回答最后的问题。如果你不知道答案，说“我不知道”，不可以编造答案。
+除了回答问题外，还需要输出一个分数，表示你对这个问题的回答的自信程度。分数越高，你越自信。按照以下的格式输出：
+
+回答：[回答内容]
+分数：[0到100间的数字]
+
+开始回答：
+{context}
+
+问题：{question}
 """
 
+
 def get_chatgpt_reply(query, context=[]):
     context += [query]
-    try:
-        completion = openai.ChatCompletion.create(
-            model="gpt-3.5-turbo", 
-            messages=[{"role": "user", "content": '\n\n'.join(context)[:4096]}]
-        )
-        response = completion.choices[0].message["content"]
-    except RateLimitError:
-        response = "服务器请求次数过多，请稍后重新尝试"
-    context += [response]
+
+    llm_chat = ChatOpenAI(model_name="gpt-3.5-turbo", max_tokens=2000, temperature=0.3)
+    embeddings = OpenAIEmbeddings()
+    docsearch = Chroma(persist_directory="./VectorDB", embedding_function=embeddings)
+
+    output_parser = RegexParser(
+        regex=r"(.*)\n*分数：([0-9]*).*",
+        output_keys=["answer", "score"],
+    )
+    PROMPT = PromptTemplate(
+        template=qa_prompt, input_variables=["context", "question"], output_parser=output_parser
+    )
+    chain_type_kwargs = {"prompt": PROMPT}
+
+    qa = RetrievalQA.from_chain_type(llm_chat, chain_type="map_rerank", retriever=docsearch.as_retriever(), chain_type_kwargs=chain_type_kwargs)
+    result = qa.run(query)
+
+    context += [result]
     responses = [(u,b) for u,b in zip(context[::2], context[1::2])]
     return responses, context
 
 
 def get_chatgpt_summary(content):
-    try:
-        completion = openai.ChatCompletion.create(
-            model="gpt-3.5-turbo", 
-            messages=[{"role": "system", "content": prompt}, 
-                      {"role": "user", "content": content}], 
-            temperature=0.5
-        )
-        response = completion.choices[0].message["content"]
-    except RateLimitError:
-        response = "服务器请求次数过多，请稍后重新尝试"
-    return response
+    texts = text_splitter.split_text(content)
+    docs = [LangDoc(page_content=t) for t in texts]
+    llm_summary = OpenAI(model_name="gpt-3.5-turbo", max_tokens=300, temperature=0.2)
+    each_round_template = PromptTemplate(input_variables=["text"], template=summary_prompt)
+    chain_summary = load_summarize_chain(llm_summary, chain_type="stuff", prompt=each_round_template)
+    summary = "\n*******\n".join([chain_summary.run([doc]) for doc in docs])
+    return summary
 
 
 import gradio as gr
 from docx import Document
+import os
 
 
 def upload_file(file):
@@ -55,16 +94,17 @@ def upload_file(file):
     for para in doc.paragraphs:
         content += para.text
         content += '\n'
+    texts = text_splitter.split_text(content)
+    docs = [LangDoc(page_content=t) for t in texts]
+    embeddings = OpenAIEmbeddings()
+    docsearch = Chroma.from_documents(docs, embeddings, persist_directory="./VectorDB")
+    docsearch.persist()
     return content
 
-def download_file(content):
-    doc = Document()
-    doc.add_paragraph(content)
-    doc.save("会议纪要.docx")
-    return None
 
 def set_api_key(api_key):
     openai.api_key = api_key
+    os.environ["OPENAI_API_KEY"] = api_key
     return None
 
 
@@ -75,40 +115,37 @@ with gr.Blocks(theme=gr.themes.Default(text_size='lg', radius_size='sm')) as dem
         desc = gr.Markdown("<center>让AI帮你整理会议纪要\n\n支持.docx文件</center>")
     
     with gr.Column():
-        # 文件上传
-        file_input = gr.File(file_types=[".docx"], label="原始文稿", interactive=True)
-        upload_btn = gr.Button(value="上传")
-
+        # api key
+        api_input = gr.Textbox(label="API Key", placeholder="请输入API Key", type="password")
+        api_btn = gr.Button(value="设置")
+        api_btn.click(fn=set_api_key, inputs=api_input, outputs=None)
     with gr.Row():
         with gr.Column():
-            # api key
-            api_input = gr.Textbox(label="API Key", placeholder="请输入API Key", type="password")
-            api_btn = gr.Button(value="设置")
-            api_btn.click(fn=set_api_key, inputs=api_input, outputs=None)
+            # 文件上传
+            file_input = gr.File(file_types=[".docx"], label="原始文稿", interactive=True)
+            upload_btn = gr.Button(value="上传")
             # 文字展示
             with gr.Tab("原文"):
                 # 原文
                 content_box = gr.Textbox(label="文稿内容")
-                download_btn = gr.Button(value="下载")
             with gr.Tab("总结"):
                 # 总结
                 summary_box = gr.Textbox(label="总结内容")
-                download_btn = gr.Button(value="下载")
         with gr.Column():
             # 对话交互
-            chatbot = gr.Chatbot(label="对话内容").style(height=400)
+            chatbot = gr.Chatbot(label="对话内容").style(height=500)
             state = gr.State([])
             txt = gr.Textbox(label="用户", placeholder="请输入内容")
             with gr.Row():
-                summary = gr.Button(value="一键梳理")
+                summary = gr.Button(value="一键总结")
                 clear = gr.Button(value="清空")
             
             summary.click(fn=get_chatgpt_summary, inputs=content_box, outputs=summary_box)
             txt.submit(get_chatgpt_reply, [txt, state], [chatbot, state])
             clear.click(lambda: None, None, chatbot, queue=False)
 
+    # 上传文件，langchain解析
     upload_btn.click(fn=upload_file, inputs=file_input, outputs=content_box)
-    download_btn.click(fn=download_file, inputs=content_box, outputs=None)
 
 
 demo.launch()