feat: add paramters

app.py

@@ -11,7 +11,11 @@ def load_index():
         allow_dangerous_deserialization=True
     )
 
-def search_documents(input_text, courseOfStudy, index):
+default_k = 4
+default_fetch_k = 500
+default_threshold = 0.5
+
+def search_documents(input_text, courseOfStudy, index, k = default_k, fetch_k = default_fetch_k, score_threshold = default_threshold):
     if not input_text:
         return []
     
@@ -20,7 +24,13 @@ def search_documents(input_text, courseOfStudy, index):
         metadata["学校種別"] = courseOfStudy
     
     try:
-        docs_and_scores = index.similarity_search_with_score(input_text, filter=metadata)
+        docs_and_scores = index.similarity_search_with_score(
+            input_text, 
+            filter=metadata,
+            k=k,
+            fetch_k=fetch_k,
+            score_threshold=score_threshold
+            )
     except Exception as e:
         print(f"Error during search: {e}")
         return []
@@ -61,6 +71,10 @@ def setup_gradio_interface():
         datatype=["html", "str", "str", "number", "str"],
         label="検索結果"
     )
+
+    k = gr.Number(value=default_k, label="検索結果の数", minimum=1, maximum=200)
+    fetch_k = gr.Number(value=default_fetch_k, label="検索対象の数", minimum=1, maximum=2000)
+    score_threshold = gr.Slider(value=default_threshold, label="スコアの閾値", minimum=0.0, maximum=1.0, step=0.01)
     
     index = load_index()
 
@@ -72,8 +86,8 @@ def setup_gradio_interface():
     ]
     
     interface = gr.Interface(
-        fn=lambda text, courseOfStudy: search_documents(text, courseOfStudy, index),
-        inputs=[text_input, metadata_selector], 
+        fn=lambda text, courseOfStudy, k, fetch_k, score_threshold: search_documents(text, courseOfStudy, index, k=k, fetch_k=fetch_k, score_threshold=score_threshold),
+        inputs=[text_input, metadata_selector, k, fetch_k, score_threshold], 
         outputs=[output_table, gr.JSON(label="JSON")], 
         title="Japanese Course of Study Predictor",
         description="入力したテキストに類似するテキストを持つ学習指導要領コードを検索します。",

predict.ipynb

@@ -0,0 +1,156 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain_community.vectorstores.faiss import FAISS\n",
+    "from langchain_community.embeddings.huggingface import HuggingFaceEmbeddings\n",
+    "from pprint import pprint"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def load_index():\n",
+    "    index_path = \"./storage\"\n",
+    "    embedding_model = HuggingFaceEmbeddings(model_name=\"intfloat/multilingual-e5-large\")\n",
+    "    return FAISS.load_local(\n",
+    "        folder_path=index_path, \n",
+    "        embeddings=embedding_model,\n",
+    "        allow_dangerous_deserialization=True\n",
+    "    )\n",
+    "    \n",
+    "index = load_index()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 41,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def search_documents(input_text, courseOfStudy, index, k = 4, fetch_k = 20, score_threshold=0.5):\n",
+    "    if not input_text:\n",
+    "        return []\n",
+    "    \n",
+    "    metadata = {}\n",
+    "    if courseOfStudy:\n",
+    "        metadata[\"学校種別\"] = courseOfStudy\n",
+    "    \n",
+    "    try:\n",
+    "        docs_and_scores = index.similarity_search_with_score(\n",
+    "            input_text, \n",
+    "            filter=metadata,\n",
+    "            fetch_k=fetch_k,\n",
+    "            k = k,\n",
+    "            score_threshold=score_threshold\n",
+    "            )\n",
+    "    except Exception as e:\n",
+    "        print(f\"Error during search: {e}\")\n",
+    "        return []\n",
+    "    \n",
+    "    rows = [\n",
+    "        [\n",
+    "            f\"<a href='https://w3id.org/jp-cos/{doc.metadata['id']}' target='_blank'>{doc.metadata['id']}</a>\",      \n",
+    "            doc.metadata[\"学校種別\"],\n",
+    "            doc.metadata[\"教科等\"],\n",
+    "            round(float(score), 3),\n",
+    "            doc.page_content,\n",
+    "        ]\n",
+    "        for doc, score in docs_and_scores\n",
+    "    ]\n",
+    "\n",
+    "    json_data = [\n",
+    "        {\n",
+    "            \"dcterms:identifier\": doc.metadata['id'],\n",
+    "            \"jp-cos:courseOfStudy\": doc.metadata[\"学校種別\"],\n",
+    "            \"jp-cos:subjectArea\": doc.metadata[\"教科等\"],\n",
+    "            \"score\": round(float(score), 3),\n",
+    "            \"jp-cos:sectionText\": doc.page_content,\n",
+    "        }\n",
+    "        for doc, score in docs_and_scores\n",
+    "    ]\n",
+    "\n",
+    "    return [\n",
+    "        rows,\n",
+    "        json_data\n",
+    "    ]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 42,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[{'dcterms:identifier': '8362225540000000',\n",
+      "  'jp-cos:courseOfStudy': '中学校',\n",
+      "  'jp-cos:sectionText': 'アの(ウ)の㋑については，各器官の働きを中心に扱うこと。',\n",
+      "  'jp-cos:subjectArea': '理科',\n",
+      "  'score': 0.416},\n",
+      " {'dcterms:identifier': '8361235820000000',\n",
+      "  'jp-cos:courseOfStudy': '中学校',\n",
+      "  'jp-cos:sectionText': 'アの(ア)の㋑については，pHにも触れること。',\n",
+      "  'jp-cos:subjectArea': '理科',\n",
+      "  'score': 0.417},\n",
+      " {'dcterms:identifier': '8362235720000000',\n",
+      "  'jp-cos:courseOfStudy': '中学校',\n",
+      "  'jp-cos:sectionText': 'アの(ア)の㋑については，有性生殖の仕組みを減数分裂と関連付けて扱うこと。「無性生殖」については，単細胞生物の分裂や栄養生殖にも触れること。',\n",
+      "  'jp-cos:subjectArea': '理科',\n",
+      "  'score': 0.419},\n",
+      " {'dcterms:identifier': '8361225630000000',\n",
+      "  'jp-cos:courseOfStudy': '中学校',\n",
+      "  'jp-cos:sectionText': 'アの(イ)の㋑の「酸化や還元」については，簡単なものを扱うこと。',\n",
+      "  'jp-cos:subjectArea': '理科',\n",
+      "  'score': 0.422}]\n"
+     ]
+    }
+   ],
+   "source": [
+    "input_text = \"小さい体でジャンプするトビムシは、菌糸類を食べて糞にする。\"\n",
+    "\n",
+    "grade = \"中学校\"\n",
+    "\n",
+    "k = 4\n",
+    "\n",
+    "fetch_k = 500\n",
+    "\n",
+    "threshold = 0.5\n",
+    "\n",
+    "result = search_documents(input_text, grade, index, k, fetch_k, threshold)\n",
+    "\n",
+    "pprint(result[1])"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}