Update app.py

app.py

@@ -30,52 +30,61 @@ config = DEFAULT_CONFIG.copy()
 config.update(cli_args)
 config['vector_size'] = int(config['vector_size'])
 
-
-#加载数据
-from datasets import load_dataset
-from datasets import  concatenate_datasets
-ds1 = load_dataset("mangopy/ToolRet-Tools", "code")
-ds2 = load_dataset("mangopy/ToolRet-Tools", "customized")
-ds3 = load_dataset("mangopy/ToolRet-Tools", "web")
-ds = concatenate_datasets([ds1['tools'], ds2['tools'], ds3['tools']])
-ds = ds.rename_columns({'id':'tool'})
-
-#merge
-
-# 随便建立一个pd.DataFrame, 有两列，一列是id，一列是text
-import pandas as pd
+# ---------------------------
+# 缓存数据集加载函数（避免每次运行时重复下载数据）
+# ---------------------------
+@st.cache_data
+def load_tools_datasets():
+    from datasets import load_dataset, concatenate_datasets
+    ds1 = load_dataset("mangopy/ToolRet-Tools", "code")
+    ds2 = load_dataset("mangopy/ToolRet-Tools", "customized")
+    ds3 = load_dataset("mangopy/ToolRet-Tools", "web")
+    ds = concatenate_datasets([ds1['tools'], ds2['tools'], ds3['tools']])
+    # 重命名'id'字段为'tool'
+    ds = ds.rename_columns({'id': 'tool'})
+    return ds
+
+ds = load_tools_datasets()
 df2 = ds.to_pandas()
+# 如果数据量较大，可以通过设置索引加速后续的合并操作
+df2.set_index('tool', inplace=True)
 
-
-
+# ---------------------------
+# 缓存模型加载函数
+# ---------------------------
 @st.cache_resource
 def get_model(model_path: str = config['model_path']):
     return SentenceModel(model_path)
 
+# 缓存检索器创建函数
 @st.cache_resource
 def create_retriever(vector_sz: int, dataset_path: str, embedding_field: str, id_field: str, _model):
     retriever = JSONLIndexer(vector_sz=vector_sz, model=_model)
     retriever.load_jsonl(dataset_path, embedding_field=embedding_field, id_field=id_field)
     return retriever
 
-# 在侧边栏中添加模型配置标题
+# ---------------------------
+# 侧边栏配置
+# ---------------------------
 st.sidebar.markdown("<div style='text-align: center;'><h3>📄 Model Configuration</h3></div>", unsafe_allow_html=True)
-
-
-# 添加模型选项下拉框，目前只有一个模型可选
 model_options = ["BAAI/bge-base-en-v1.5"]
 selected_model = st.sidebar.selectbox("Select Model", model_options)
 st.sidebar.write("Selected model:", selected_model)
 st.sidebar.write("Embedding length: 768")
 
-# 使用选中的模型加载
+# 使用下拉框选中的模型（避免重复加载）
 model = get_model(selected_model)
-
-
-model = get_model(config['model_path'])
-retriever = create_retriever(config['vector_size'], config['dataset_path'], config['embedding_field'], config['id_field'], _model=model)
-
-# 美化界面
+retriever = create_retriever(
+    config['vector_size'],
+    config['dataset_path'],
+    config['embedding_field'],
+    config['id_field'],
+    _model=model
+)
+
+# ---------------------------
+# 界面样式设置
+# ---------------------------
 st.markdown("""
     <style>
     .search-container {
@@ -102,7 +111,9 @@ st.markdown("""
 
 st.markdown("<h1 style='text-align: center;'>🔍 Tool Retrieval</h1>", unsafe_allow_html=True)
 
-
+# ---------------------------
+# 主体检索区域
+# ---------------------------
 col1, col2 = st.columns([4, 1])
 with col1:
     query = st.text_input("", placeholder="Enter your search query...", key="search_query", label_visibility="collapsed")
@@ -111,15 +122,13 @@ with col2:
 
 top_k = st.slider("Top-K tools", 1, 100, 50, help="Choose the number of results to display")
 
-styled_results = None
 if search_clicked and query:
     rec_ids, scores = retriever.search_return_id(query, top_k)
-    df1 = pd.DataFrame({ "relevance": scores, "tool": rec_ids})
-    # print(df1)
-    # merge两个DataFrame
-    results_df = pd.merge(df1, df2, on='tool', how = 'left')
-
-    # results_df["interface"] = "asdasdadasdasdasdasdasdasdasasdasdasdasdasdasdasdasdasdasdasdasdasdasdasdasdasdassdasdasdasdasdasabababbabasdbabsdbasbdadabdbasdbasbdbasdbasdbasdb"
+    # 构建检索结果 DataFrame
+    df1 = pd.DataFrame({"relevance": scores, "tool": rec_ids})
+    # 使用 join 加速合并（前提是 df2 已设置好索引）
+    results_df = df1.join(df2, on='tool', how='left').reset_index(drop=False)
+    
     st.subheader("🗂️ Retrieval results")
     
     styled_results = results_df.style.apply(
@@ -129,7 +138,7 @@ if search_clicked and query:
         ],
         axis=0,
     ).format({"relevance": "{:.4f}"})
-
+    
     st.dataframe(
         styled_results,
         column_config={