impl

app.py

@@ -1,7 +1,307 @@
+"""
+streamlit run app.py --server.address 0.0.0.0
+"""
+
+from __future__ import annotations
+
 import streamlit as st
 import os
 
-# 環境変数全てを整形して表示
-all_env_dict = os.environ
-for key in all_env_dict.keys():
-    st.write(f"{key}: {all_env_dict[key]}")
+import faiss
+from sentence_transformers import SentenceTransformer
+import torch
+from openai import OpenAI
+import streamlit as st
+import pandas as pd
+import os
+from time import time
+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"
+WIKIPEDIA_JA_EMB_DS = "hotchpotch/wikipedia-passages-jawiki-embeddings"
+
+EMB_MODEL_PQ = {
+    "intfloat/multilingual-e5-small": 96,
+    "intfloat/multilingual-e5-base": 192,
+    "intfloat/multilingual-e5-large": 256,
+    "cl-nagoya/sup-simcse-ja-base": 192,
+    "pkshatech/GLuCoSE-base-ja": 192,
+}
+
+EMB_MODEL_NAMES = list(EMB_MODEL_PQ.keys())
+
+OPENAI_MODEL_NAMES = [
+    "gpt-3.5-turbo-1106",
+    "gpt-4-1106-preview",
+]
+
+E5_QUERY_TYPES = [
+    "passage",
+    "query",
+]
+
+DEFAULT_QA_PROMPT = """
+## Instruction
+
+Prepare an explanatory statement for the question, including as much detailed explanation as possible.
+Avoid speculations or information not contained in the contexts. Heavily favor knowledge provided in the documents before falling back to baseline knowledge or other contexts. If searching the contexts didn"t yield any answer, just say that.
+
+Responses must be given in Japanese.
+
+## Contexts
+
+{contexts}
+
+## Question
+
+{question}
+""".strip()
+
+
+if os.getenv("SPACE_ID"):
+    USE_HF_SPACE = True
+    os.environ["HF_HOME"] = "/data/.huggingface"
+else:
+    USE_HF_SPACE = False
+
+# for tokenizer
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
+
+
+@st.cache_data
+def get_model(name: str, max_seq_length=512):
+    device = "cpu"
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif torch.backends.mps.is_available():
+        device = "mps"
+    model = SentenceTransformer(name, device=device)
+    model.max_seq_length = max_seq_length
+    return model
+
+
+@st.cache_data
+def get_wikija_ds(name: str = WIKIPEDIA_JS_DS_NAME):
+    ds = load_dataset(path=WIKIPEDIA_JA_DS, name=name, split="train")
+    return ds
+
+
+@st.cache_data
+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(
+        f"https://huggingface.co/datasets/{ja_emb_ds}/resolve/main/{target_path}"
+    )
+    index = faiss.read_index(index_local_path)
+    index.nprobe = 128
+    return index
+
+
+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
+):
+    start_time = time()
+    emb = text_to_emb(emb_model, question, search_text_prefix)
+    emb_exec_time = time() - start_time
+    scores, indexes = faiss_index.search(emb, top_k)
+    faiss_seartch_time = time() - emb_exec_time - start_time
+    scores = scores[0]
+    indexes = indexes[0]
+    results = []
+    for idx, score in zip(indexes, scores):  # type: ignore
+        idx = int(idx)
+        passage = ds[idx]
+        results.append((score, passage))
+    return results, emb_exec_time, faiss_seartch_time
+
+
+def to_contexts(passages):
+    contexts = ""
+    for passage in passages:
+        title = passage["title"]
+        text = passage["text"]
+        # section = passage["section"]
+        contexts += f"- {title}: {text}\n"
+    return contexts
+
+
+def qa(
+    question: str,
+    passages: list,
+    model_name: str,
+    temperature: int,
+    qa_prompt: str,
+    max_tokens=2000,
+):
+    client = OpenAI()
+    contexts = to_contexts(passages)
+    prompt = qa_prompt.format(contexts=contexts, question=question)
+    response = client.chat.completions.create(
+        model=model_name,
+        messages=[
+            {"role": "user", "content": prompt},
+        ],
+        stream=True,
+        temperature=temperature,
+        max_tokens=max_tokens,
+        seed=42,
+    )
+    for chunk in response:
+        delta = chunk.choices[0].delta
+        yield delta.content or ""
+
+
+def generate_answer(
+    buf, question, passages, model_name, temperature, qa_prompt, max_tokens
+):
+    buf.write("⏳回答の生成中...")
+    texts = ""
+    for char in qa(
+        question=question,
+        passages=passages,
+        model_name=model_name,
+        temperature=temperature,
+        qa_prompt=qa_prompt,
+    ):
+        texts += char
+        buf.write(texts)
+
+
+def to_df(scores, passages):
+    df = pd.DataFrame(passages)
+    df["text"] = df["text"]
+    df["score"] = scores
+    df_rows = ["score", "title", "text", "section"]
+    df = df[df_rows]
+    return df
+
+
+def app():
+    st.title("Wikipedia 日本語 RAG 検索")
+    st.subheader("⭐️大元へのリンクを貼る")
+
+    st.text_area(
+        "Question",
+        key="question",
+        value="1975年に『アザミ嬢のララバイ』でデビューした女性歌手で、『わかれうた』『地上の星』などの曲を出しているのは誰?",
+    )
+    if not OPENAI_API_KEY:
+        st.text_input(
+            "OpenAI API Key",
+            key="openai_api_key",
+            type="password",
+            placeholder="※ API_KEY 未入力時は、回答生成せずに検索のみ",
+        )
+    else:
+        st.session_state.openai_api_key = OPENAI_API_KEY
+
+    with st.expander("オプション"):
+        option_cols_main = st.columns(2)
+        with option_cols_main[0]:
+            st.selectbox("Emb Model", EMB_MODEL_NAMES, index=0, key="emb_model_name")
+        with option_cols_main[1]:
+            st.selectbox(
+                "OpenAI Model", OPENAI_MODEL_NAMES, index=0, key="openai_model_name"
+            )
+        emb_model_name = st.session_state.emb_model_name
+        option_cols_sub = st.columns(2)
+        with option_cols_sub[0]:
+            st.number_input("Top K", value=5, key="top_k", min_value=1, max_value=20)
+        with option_cols_sub[1]:
+            if "-e5-" in emb_model_name:
+                st.radio(
+                    "Passage or Query (only e5)",
+                    E5_QUERY_TYPES,
+                    index=0,
+                    key="e5_query_or_passage",
+                    horizontal=True,
+                )
+                e5_query_or_passage = st.session_state.e5_query_or_passage
+                index_emb_model_name = (
+                    f"{emb_model_name.split('/')[-1]}-{e5_query_or_passage}"
+                )
+                search_text_prefix = f"{e5_query_or_passage}: "
+            else:
+                index_emb_model_name = emb_model_name.split("/")[-1]
+                search_text_prefix = ""
+        option_cols = st.columns(3)
+        with option_cols[0]:
+            st.slider("Temperature", 0.0, 1.0, value=0.8, key="temperature")
+        with option_cols[1]:
+            st.slider("nprobe", 16, 1024, value=128, key="nprobe")
+        with option_cols[2]:
+            st.number_input(
+                "max_tokens", value=2000, key="max_tokens", min_value=1, max_value=16000
+            )
+        st.text_area("QA Prompt", value=DEFAULT_QA_PROMPT, key="qa_prompt")
+
+    loading_placeholder = st.empty()
+    loading_placeholder.text("⏳ Loading - Embedding Model...")
+    emb_model = get_model(st.session_state.emb_model_name)
+    loading_placeholder.text("⏳ Loading - Faiss Index...")
+    emb_model_pq = EMB_MODEL_PQ[emb_model_name]
+    index_name = f"{index_emb_model_name}/index_IVF2048_PQ{emb_model_pq}.faiss"
+    faiss_index = get_faiss_index(index_name=index_name)
+    faiss_index.nprobe = st.session_state.nprobe
+    loading_placeholder.text("⏳ Loading - Huggingface Dataset...")
+    ds = get_wikija_ds()
+    loading_placeholder.empty()
+
+    if st.button("Search"):
+        answer_header = st.empty()
+        answer_text_buffer = st.empty()
+
+        question = st.session_state.question
+        top_k = st.session_state.top_k
+        scores = []
+        passages = []
+        search_results, emb_exec_time, faiss_seartch_time = search(
+            faiss_index,
+            emb_model,
+            ds,
+            question,
+            search_text_prefix=search_text_prefix,
+            top_k=top_k,
+        )
+        st.subheader("Search Results: ")
+        st.write(
+            f"⏱️ generate embedding: {emb_exec_time*1000:.2f}ms /  faiss search: {faiss_seartch_time*1000:.2f}ms"
+        )
+        for score, passage in search_results:
+            scores.append(score)
+            passages.append(passage)
+        df = to_df(scores, passages)
+        st.dataframe(df, hide_index=True)
+
+        openai_api_key = st.session_state.openai_api_key
+        if openai_api_key:
+            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
+            generate_answer(
+                buf=answer_text_buffer,
+                question=question,
+                passages=passages,
+                model_name=openai_model_name,
+                temperature=temperature,
+                qa_prompt=qa_prompt,
+                max_tokens=max_tokens,
+            )
+
+
+if __name__ == "__main__":
+    app()