Create app.py

app.py

@@ -0,0 +1,120 @@
+import gradio
+import gradio as gr
+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.vectorstores import Chroma
+from langchain_core.prompts import PromptTemplate
+
+template = """Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer. Use three sentences maximum. Keep the answer as concise as possible. Always say "thanks for asking!" at the end of the answer.
+Tips: Make sure to cite your sources, and use the exact words from the context.
+{context}
+Question: {question}
+Helpful Answer:"""
+QA_CHAIN_PROMPT = PromptTemplate.from_template(template)
+
+
+class RAGDemo(object):
+    def __init__(self):
+        self.embedding = None
+        self.vector_db = None
+        self.chat_model = None
+
+    def _init_chat_model(self, model_name, api_key):
+        if not api_key:
+            gradio.Error("Please enter model API key.")
+            return
+        if 'glm' in model_name:
+            self.chat_model = ChatZhipuAI(
+                temperature=0.5,
+                api_key=api_key,
+                model="glm-3-turbo",
+            )
+        elif 'gemini' in model_name:
+            self.chat_model = ChatGooglePalm(
+                google_api_key=api_key,
+                model_name='gemini-pro'
+            )
+
+    def _init_embedding(self, embedding_model_name, api_key):
+        if not api_key:
+            gradio.Error("Please enter embedding API key.")
+            return
+        if 'glm' in embedding_model_name:
+            gradio.Error("GLM is not supported yet.")
+        else:
+            self.embedding = HuggingFaceInferenceAPIEmbeddings(
+                api_key=api_key, model_name=embedding_model_name
+            )
+
+    def _build_vector_db(self, file_path):
+        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()
+
+        text_splitter = SpacyTextSplitter(chunk_size=500, chunk_overlap=50)
+        docs = text_splitter.split_documents(pages)
+
+        self.vector_db = Chroma.from_documents(
+            documents=docs, embedding=self.embedding
+        )
+        gr.Info("Vector database built.")
+
+    def _retrieval_qa(self, input_text):
+        basic_qa = RetrievalQA.from_chain_type(
+            self.chat_model,
+            retriever=self.vector_db.as_retriever(),
+            chain_type_kwargs={"prompt": QA_CHAIN_PROMPT}
+        )
+        return basic_qa.invoke(input_text)
+
+    def __call__(self):
+        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")
+                        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")
+                        embedding_model = gr.Dropdown(
+                            choices=['glm-embedding-2', '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]
+            )
+            data_file.upload(self._build_vector_db, inputs=data_file)
+            submit_btn.click(
+                self._retrieval_qa,
+                inputs=input_text,
+                outputs=output,
+            )
+            return demo
+
+
+app = RAGDemo()
+app().launch()