Add "Search Only" to OpenAI model options and make OpenAI API key input optional

app.py

@@ -1,23 +1,15 @@
-"""
-streamlit run app.py --server.address 0.0.0.0
-"""
-
 from __future__ import annotations
 
-import streamlit as st
 import os
-
-import faiss
-from sentence_transformers import SentenceTransformer
+import pandas as pd
 import torch
-from openai import OpenAI
+import faiss
 import streamlit as st
-import pandas as pd
-import os
 from time import time
+from openai import OpenAI
+from sentence_transformers import SentenceTransformer
+from datasets import load_dataset
 from datasets.download import DownloadManager
-from datasets import load_dataset  # type: ignore
-
 
 WIKIPEDIA_JA_DS = "singletongue/wikipedia-utils"
 WIKIPEDIA_JS_DS_NAME = "passages-c400-jawiki-20230403"
@@ -36,6 +28,7 @@ EMB_MODEL_NAMES = list(EMB_MODEL_PQ.keys())
 OPENAI_MODEL_NAMES = [
     "gpt-3.5-turbo-1106",
     "gpt-4-1106-preview",
+    "Search Only",
 ]
 
 E5_QUERY_TYPES = [
@@ -60,7 +53,6 @@ Responses must be given in Japanese.
 {question}
 """.strip()
 
-
 if os.getenv("SPACE_ID"):
     USE_HF_SPACE = True
     os.environ["HF_HOME"] = "/data/.huggingface"
@@ -68,9 +60,7 @@ if os.getenv("SPACE_ID"):
 else:
     USE_HF_SPACE = False
 
-# for tokenizer
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
-
 OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
 
 
@@ -81,6 +71,7 @@ def get_model(name: str, max_seq_length=512):
         device = "cuda"
     elif torch.backends.mps.is_available():
         device = "mps"
+
     model = SentenceTransformer(name, device=device)
     model.max_seq_length = max_seq_length
     return model
@@ -93,9 +84,7 @@ def get_wikija_ds(name: str = WIKIPEDIA_JS_DS_NAME):
 
 
 @st.cache_resource
-def get_faiss_index(
-    index_name: str, ja_emb_ds: str = WIKIPEDIA_JA_EMB_DS, name=WIKIPEDIA_JS_DS_NAME
-):
+def get_faiss_index(index_name: str, ja_emb_ds: str = WIKIPEDIA_JA_EMB_DS, name=WIKIPEDIA_JS_DS_NAME):
     target_path = f"faiss_indexes/{name}/{index_name}"
     dm = DownloadManager()
     index_local_path = dm.download(
@@ -110,9 +99,7 @@ def text_to_emb(model, text: str, prefix: str):
     return model.encode([prefix + text], normalize_embeddings=True)
 
 
-def search(
-    faiss_index, emb_model, ds, question: str, search_text_prefix: str, top_k: int
-):
+def search(faiss_index, emb_model, ds, question: str, search_text_prefix: str, top_k: int):
     start_time = time()
     emb = text_to_emb(emb_model, question, search_text_prefix)
     emb_exec_time = time() - start_time
@@ -121,7 +108,7 @@ def search(
     scores = scores[0]
     indexes = indexes[0]
     results = []
-    for idx, score in zip(indexes, scores):  # type: ignore
+    for idx, score in zip(indexes, scores):
         idx = int(idx)
         passage = ds[idx]
         results.append((score, passage))
@@ -133,7 +120,6 @@ def to_contexts(passages):
     for passage in passages:
         title = passage["title"]
         text = passage["text"]
-        # section = passage["section"]
         contexts += f"- {title}: {text}\n"
     return contexts
 
@@ -211,15 +197,13 @@ def app():
         key="question",
         value="楽曲『約束はいらない』でデビューした、声優は誰?",
     )
-    if not OPENAI_API_KEY:
-        st.text_input(
-            "OpenAI API Key",
-            key="openai_api_key",
-            type="password",
-            placeholder="※ OpenAI API Key 未入力時は回答を生成せずに、検索のみ実行します",
-        )
-    else:
-        st.session_state.openai_api_key = OPENAI_API_KEY
+    st.text_input(
+        "OpenAI API Key",
+        key="openai_api_key",
+        type="password",
+        value=OPENAI_API_KEY if OPENAI_API_KEY else "",
+        placeholder="※ OpenAI API Key 未入力時は回答を生成せずに、検索のみ実行します",
+    )
 
     with st.expander("オプション"):
         option_cols_main = st.columns(2)
@@ -229,6 +213,8 @@ def app():
             st.selectbox(
                 "OpenAI Model", OPENAI_MODEL_NAMES, index=0, key="openai_model_name"
             )
+        if "emb_model_name" not in st.session_state:
+            st.session_state.emb_model_name = EMB_MODEL_NAMES[0]  # replace with the actual default value you want to use
         emb_model_name = st.session_state.emb_model_name
         option_cols_sub = st.columns(2)
         with option_cols_sub[0]:
@@ -300,10 +286,10 @@ def app():
         st.dataframe(df, hide_index=True)
 
         openai_api_key = st.session_state.openai_api_key
-        if openai_api_key:
+        openai_model_name = st.session_state.openai_model_name
+        if openai_api_key and openai_model_name != "Search Only":
             openai_api_key = openai_api_key.strip()
             answer_header.subheader("Answer: ")
-            openai_model_name = st.session_state.openai_model_name
             temperature = st.session_state.temperature
             qa_prompt = st.session_state.qa_prompt
             max_tokens = st.session_state.max_tokens
@@ -320,4 +306,4 @@ def app():
 
 
 if __name__ == "__main__":
-    app()
+    app()