add model selection

app.py

@@ -3,7 +3,7 @@ import os
 import json
 import requests
 from langchain import FAISS
-from langchain.embeddings import CohereEmbeddings
+from langchain.embeddings import CohereEmbeddings, OpenAIEmbeddings
 from langchain import VectorDBQA
 from langchain.chat_models import ChatOpenAI
 from prompts import MyTemplate
@@ -15,7 +15,7 @@ from langchain.prompts.chat import (
 
 # Streaming endpoint
 API_URL = "https://api.openai.com/v1/chat/completions"  # os.getenv("API_URL") + "/generate_stream"
-embeddings_key = '5IRbILAbjTI0VcqTsktBfKsr13Lych9iBAFbLpkj'
+cohere_key = '5IRbILAbjTI0VcqTsktBfKsr13Lych9iBAFbLpkj'
 faiss_store = './indexer'
 
 def gen_conversation(conversations):
@@ -32,8 +32,9 @@ def gen_conversation(conversations):
     return messages
 
 
-def predict(inputs, top_p, temperature, openai_api_key, enable_index, max_tokens,
-            chat_counter, chatbot=[], history=[]):  # repetition_penalty, top_k
+def predict(inputs, top_p, temperature, openai_api_key, enable_index, max_tokens, model,
+            chat_counter, chatbot=[], history=[]):
+    model = model[0]
     headers = {
         "Content-Type": "application/json",
         "Authorization": f"Bearer {openai_api_key}"
@@ -43,7 +44,11 @@ def predict(inputs, top_p, temperature, openai_api_key, enable_index, max_tokens
     #Debugging
     if enable_index:
         # Faiss 检索最近的embedding
-        docsearch = FAISS.load_local(faiss_store, CohereEmbeddings(cohere_api_key=embeddings_key))
+        if model =='openai':
+            docsearch = FAISS.load_local(faiss_store, OpenAIEmbeddings(openai_api_key=openai_api_key))
+        else:
+            docsearch = FAISS.load_local(faiss_store, CohereEmbeddings(cohere_api_key=cohere_key ))
+        # 构建模板
         llm = ChatOpenAI(openai_api_key=openai_api_key, max_tokens=max_tokens)
         messages_combine = [
             SystemMessagePromptTemplate.from_template(MyTemplate['chat_combine_template']),
@@ -133,7 +138,7 @@ with gr.Blocks(css="""#col_container {width: 1000px; margin-left: auto; margin-r
     gr.HTML("""<h1 align="center">🚀Finance ChatBot🚀</h1>""")
     with gr.Column(elem_id="col_container"):
         openai_api_key = gr.Textbox(type='password', label="输入OPEN API Key")
-        chatbot = gr.Chatbot(elem_id='chatbot')  # c
+        chatbot = gr.Chatbot(elem_id='chatbot')
         inputs = gr.Textbox(placeholder="您有什么问题可以问我", label="输入数字经济，两会，硅谷银行相关的提问")
         state = gr.State([])
 
@@ -148,16 +153,15 @@ with gr.Blocks(css="""#col_container {width: 1000px; margin-left: auto; margin-r
                               label="Max Tokens", )
             temperature = gr.Slider(minimum=-0, maximum=5.0, value=1.0, step=0.1, interactive=True,
                                     label="Temperature", )
-            # top_k = gr.Slider( minimum=1, maximum=50, value=4, step=1, interactive=True, label="Top-k",)
-            # repetition_penalty = gr.Slider( minimum=0.1, maximum=3.0, value=1.03, step=0.01, interactive=True, label="Repetition Penalty", )
+            model = gr.CheckboxGroup(["cohere", "openai", "mpnet"])
             chat_counter = gr.Number(value=0, precision=0)
             enable_index = gr.Checkbox(label='是', info='是否使用研报等金融数据')
             # 后续考虑加入搜索结果
             enable_search = gr.Checkbox(label='是', info='是否使用搜索结果')
 
-    inputs.submit(predict, [inputs, top_p, temperature, openai_api_key, enable_index, max_tokens, chat_counter, chatbot, state],
+    inputs.submit(predict, [inputs, top_p, temperature, openai_api_key, enable_index, max_tokens, model, chat_counter, chatbot, state],
                   [chatbot, state, chat_counter], )
-    run.click(predict, [inputs, top_p, temperature, openai_api_key, enable_index, max_tokens, chat_counter, chatbot, state],
+    run.click(predict, [inputs, top_p, temperature, openai_api_key, enable_index, max_tokens, model, chat_counter, chatbot, state],
               [chatbot, state, chat_counter], )
 
     # 每次对话结束都重置对话