dev/streamlit-ui (#12)

- modify to sdk streamlit (06d784cd04f9d79f73d733703dc18fcff332a3c0)

README.md

@@ -3,8 +3,8 @@ title: NVDA 日本語版ガイドブックQA
 emoji: 👀
 colorFrom: green
 colorTo: yellow
-sdk: gradio
-sdk_version: 3.38.0
+sdk: streamlit
+sdk_version: 1.25.0
 app_file: app.py
 pinned: false
 license: cc0-1.0

app.py

@@ -1,9 +1,13 @@
 from time import time
-import gradio as gr
+from typing import Iterable
+
+# import gradio as gr
+import streamlit as st
 from langchain.chains import RetrievalQA
 from langchain.embeddings import OpenAIEmbeddings
 from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.prompts import PromptTemplate
+
+# from langchain.prompts import PromptTemplate
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from langchain.llms import HuggingFacePipeline
@@ -25,7 +29,7 @@ E5_EMBEDDINGS = HuggingFaceEmbeddings(
     encode_kwargs=E5_ENCODE_KWARGS,
 )
 
-if torch.cuda.is_available():
+if False and torch.cuda.is_available():  # TODO: for local debug
     RINNA_MODEL_NAME = "rinna/bilingual-gpt-neox-4b-instruction-ppo"
     RINNA_TOKENIZER = AutoTokenizer.from_pretrained(RINNA_MODEL_NAME, use_fast=False)
     RINNA_MODEL = AutoModelForCausalLM.from_pretrained(
@@ -86,17 +90,6 @@ def _get_llm_model(
     return llm
 
 
-# prompt_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.
-
-# {context}
-
-# Question: {question}
-# Answer in Japanese:"""
-# PROMPT = PromptTemplate(
-#     template=prompt_template, input_variables=["context", "question"]
-# )
-
-
 def get_retrieval_qa(
     collection_name: str | None,
     model_name: str | None,
@@ -122,7 +115,6 @@ def get_retrieval_qa(
     llm = _get_llm_model(model_name, temperature)
 
     # chain_type_kwargs = {"prompt": PROMPT}
-
     result = RetrievalQA.from_chain_type(
         llm=llm,
         chain_type="stuff",
@@ -146,11 +138,8 @@ def get_related_url(metadata):
         yield f'<p>URL: <a href="{url}">{url}</a> (category: {category})</p>'
 
 
-def main(
-    query: str, collection_name: str, model_name: str, option: str, temperature: float
-):
+def run_qa(query: str, qa: RetrievalQA) -> tuple[str, str]:
     now = time()
-    qa = get_retrieval_qa(collection_name, model_name, temperature, option)
     try:
         result = qa(query)
     except InvalidRequestError as e:
@@ -163,29 +152,62 @@ def main(
     return result["result"], html
 
 
+def main(
+    query: str,
+    collection_name: str | None,
+    model_name: str | None,
+    option: str | None,
+    temperature: float,
+    e5_option: list[str],
+) -> Iterable[tuple[str, tuple[str, str]]]:
+    qa = get_retrieval_qa(collection_name, model_name, temperature, option)
+    if collection_name == "E5":
+        for option in e5_option:
+            if option == "No":
+                yield "E5 No", run_qa(query, qa)
+            elif option == "Query":
+                yield "E5 Query", run_qa("query: " + query, qa)
+            elif option == "Passage":
+                yield "E5 Passage", run_qa("passage: " + query, qa)
+            else:
+                raise ValueError("Unknow option")
+    else:
+        yield "OpenAI", run_qa(query, qa)
+
+
 AVAILABLE_LLMS = ["GPT-3.5", "GPT-4"]
 
 if RINNA_MODEL is not None:
     AVAILABLE_LLMS.append("rinna")
 
-nvdajp_book_qa = gr.Interface(
-    fn=main,
-    inputs=[
-        gr.Textbox(label="query"),
-        gr.Radio(["E5", "OpenAI"], value="E5", label="Embedding"),
-        gr.Radio(
-            AVAILABLE_LLMS, value="GPT-3.5", label="Model", info="GPU環境だとrinnaが選択可能"
-        ),
-        gr.Radio(
-            ["All", "ja-book", "ja-nvda-user-guide", "en-nvda-user-guide"],
-            value="All",
-            label="絞り込み",
-            info="ドキュメント制限する?",
-        ),
-        gr.Slider(0, 2),
-    ],
-    outputs=[gr.Textbox(label="answer"), gr.outputs.HTML()],
-)
+with st.form("my_form"):
+    query = st.text_input(label="query")
+    collection_name = st.radio(options=["E5", "OpenAI"], label="Embedding")
 
+    # if collection_name == "E5":  # TODO : 選択肢で選べるようにする
+    e5_option = st.multiselect("E5 option", ["No", "Query", "Passage"], default="No")
 
-nvdajp_book_qa.launch()
+    model_name = st.radio(
+        options=AVAILABLE_LLMS,
+        label="Model",
+        help="GPU環境だとrinnaが選択可能",
+    )
+    option = st.radio(
+        options=["All", "ja-book", "ja-nvda-user-guide", "en-nvda-user-guide"],
+        label="絞り込み",
+        help="ドキュメント制限する?",
+    )
+    temperature = st.slider(label="temperature", min_value=0, max_value=2)
+
+    submitted = st.form_submit_button("Submit")
+    if submitted:
+        with st.spinner("Searching..."):
+            results = main(
+                query, collection_name, model_name, option, temperature, e5_option
+            )
+            for type_, (answer, html) in results:
+                with st.container():
+                    st.header(type_)
+                    st.write(answer)
+                    st.markdown(html, unsafe_allow_html=True)
+                    st.divider()