modify demo layout

app.py

@@ -4,7 +4,7 @@ from langchain.chains import RetrievalQA
 from langchain.text_splitter import SpacyTextSplitter
 from langchain_community.chat_models import ChatZhipuAI, ChatGooglePalm
 from langchain_community.document_loaders import PyPDFLoader
-from langchain_community.embeddings import HuggingFaceInferenceAPIEmbeddings
+from langchain_community.embeddings import HuggingFaceInferenceAPIEmbeddings, GooglePalmEmbeddings
 from langchain_community.vectorstores import Chroma
 from langchain_core.prompts import PromptTemplate
 
@@ -44,6 +44,8 @@ class RAGDemo(object):
             return
         if 'glm' in embedding_model_name:
             gradio.Error("GLM is not supported yet.")
+        elif 'gemini' in embedding_model_name:
+            self.embedding = GooglePalmEmbeddings(google_api_key=api_key, show_progress_bar=True)
         else:
             self.embedding = HuggingFaceInferenceAPIEmbeddings(
                 api_key=api_key, model_name=embedding_model_name
@@ -53,7 +55,6 @@ class RAGDemo(object):
         if not file_path:
             gradio.Error("Please enter vector database file path.")
             return
-        gr.Info("Building vector database...")
         loader = PyPDFLoader(file_path)
         pages = loader.load()
 
@@ -63,7 +64,11 @@ class RAGDemo(object):
         self.vector_db = Chroma.from_documents(
             documents=docs, embedding=self.embedding
         )
-        gr.Info("Vector database built.")
+
+    def _init_settings(self, model_name, api_key, embedding_model, embedding_api_key, data_file):
+        self._init_chat_model(model_name, api_key)
+        self._init_embedding(embedding_model, embedding_api_key)
+        self._build_vector_db(data_file)
 
     def _retrieval_qa(self, input_text):
         basic_qa = RetrievalQA.from_chain_type(
@@ -77,37 +82,37 @@ class RAGDemo(object):
         with gr.Blocks() as demo:
             gr.Markdown("# RAG Demo\n\nbase on the [RAG learning note](https://www.jianshu.com/p/9792f1e6c3f9) and "
                         "[rag-practice](https://github.com/hiwei93/rag-practice/tree/main)")
-            with gr.Row():
-                with gr.Column():
-                    input_text = gr.Textbox(placeholder="input your question...", label="input")
-                    submit_btn = gr.Button("submit")
-                    with gr.Accordion("model settings"):
-                        api_key = gr.Textbox(placeholder="your api key", label="api key")
+            with gr.Tab("Settings"):
+                with gr.Row():
+                    with gr.Column():
                         model_name = gr.Dropdown(
                             choices=['glm-3-turbo', 'gemini-1.0-pro'],
                             value='glm-3-turbo',
                             label="model"
                         )
-                    with gr.Accordion("knowledge base settigns"):
-                        embedding_api_key = gr.Textbox(placeholder="your api key", label="embedding api key")
+                        api_key = gr.Textbox(placeholder="your api key for LLM", label="api key")
                         embedding_model = gr.Dropdown(
-                            choices=['glm-embedding-2', 'sentence-transformers/all-MiniLM-L6-v2',
+                            choices=['glm-embedding-2', 'gemini-embedding', 'sentence-transformers/all-MiniLM-L6-v2',
                                      'intfloat/multilingual-e5-large'],
                             value="sentence-transformers/all-MiniLM-L6-v2",
                             label="embedding model"
                         )
-                        data_file = gr.File(file_count='single', label="data pdf file")
-                with gr.Column():
-                    output = gr.TextArea(label="answer")
-            model_name.select(
-                self._init_chat_model,
-                inputs=[model_name, api_key]
-            )
-            embedding_model.select(
-                self._init_embedding,
-                inputs=[embedding_model, embedding_api_key]
+                        embedding_api_key = gr.Textbox(placeholder="your api key for embedding", label="embedding api key")
+                    with gr.Column():
+                        data_file = gr.File(file_count='single', label="pdf file")
+                        initial_btn = gr.Button("submit")
+            with gr.Tab("RAG"):
+                with gr.Row():
+                    with gr.Column():
+                        input_text = gr.Textbox(placeholder="input your question...", label="input")
+                        submit_btn = gr.Button("submit")
+                    with gr.Column():
+                        output = gr.TextArea(label="answer")
+            initial_btn.click(
+                self._init_settings,
+                inputs=[model_name, api_key, embedding_model, embedding_api_key, data_file]
             )
-            data_file.upload(self._build_vector_db, inputs=data_file)
+
             submit_btn.click(
                 self._retrieval_qa,
                 inputs=input_text,
@@ -117,4 +122,4 @@ class RAGDemo(object):
 
 
 app = RAGDemo()
-app().launch()
+app().launch(debug=True)