Update app.py

app.py

@@ -1,39 +1,66 @@
-import gradio as gr
-import pandas as pd
-from sentence_transformers import SentenceTransformer, util
+import json
+import os
 import torch
+from sentence_transformers import SentenceTransformer
+import faiss
+import numpy as np
+import gradio as gr
 
-# 載入語義搜索模型
-model_checkpoint = "sickcell69/cti-semantic-search-minilm"
-model = SentenceTransformer(model_checkpoint)
+def load_or_create_model_and_embeddings(model_name, data_file):
+    model_path = os.path.join(output_dir, 'saved_model')
+    embeddings_path = os.path.join(output_dir, 'corpus_embeddings.pt')
+    if os.path.exists(model_path) and os.path.exists(embeddings_path):
+        print("載入已保存的模型和嵌入...")
+        model = SentenceTransformer(model_path)
+        embeddings = torch.load(embeddings_path)
+    else:
+        model = SentenceTransformer(model_name)
+        with open(data_file, 'r', encoding='utf-8') as f:
+            data = json.load(f)
+        texts = [item['text'] for item in data]
+        embeddings = model.encode(texts, convert_to_tensor=True)
+    return model, embeddings
 
-# 載入數據
-data_path = 'labeled_cti_data.json'
-data = pd.read_json(data_path)
+# 設置參數
+model_name = 'sentence-transformers/all-MiniLM-L6-v2'
+data_file = 'labeled_cti_data.json'
+output_dir = '.'
 
-# 載入嵌入文件
-embeddings_path = 'corpus_embeddings.pt'
-corpus_embeddings = torch.load(embeddings_path, map_location=torch.device('cpu'))
+# 載入或創建模型和嵌入
+model, embeddings= load_or_create_model_and_embeddings(model_name, data_file)
 
-def semantic_search(query):
-    query_embedding = model.encode(query, convert_to_tensor=True)
-    search_hits = util.semantic_search(query_embedding, corpus_embeddings, top_k=5)
-    
+# 創建 Faiss 索引
+dimension = embeddings.shape[1]
+index = faiss.IndexFlatL2(dimension)
+index.add(embeddings.cpu().numpy().astype('float32'))
+
+def semantic_search(query, top_k=3):
+    query_vector = model.encode([query], convert_to_tensor=True)
+    distances, indices = index.search(query_vector.cpu().numpy().astype('float32'), top_k)
     results = []
-    for hit in search_hits[0]:
-        text = " ".join(data.iloc[hit['corpus_id']]['tokens'])
-        results.append(f"Score: {hit['score']:.4f} - Text: {text}")
-    
-    return "\n".join(results)
+    for i, idx in enumerate(indices[0]):
+        results.append({
+            'text': texts[idx],
+            'similarity_score': 1 - distances[0][i] / 2
+        })
+    return results
+
+def search_and_format(query):
+    results = semantic_search(query)
+    formatted_results = ""
+    for i, result in enumerate(results, 1):
+        formatted_results += f"{i}. 相似度分數: {result['similarity_score']:.4f}\n"
+        formatted_results += f"   情一: {result['text']}\n\n"
+    return formatted_results
 
+# 創建Gradio界面
 iface = gr.Interface(
-    fn=semantic_search,
-    inputs="text",
-    outputs="text",
-    title="語義搜索應用",
-    description="輸入一個查詢，然後模型將返回最相似的結果。"
+    fn=search_and_format,
+    inputs=gr.Textbox(lines=2, placeholder="輸入您的搜索查詢..."),
+    outputs=gr.Textbox(lines=10),
+    title="語義搜索",
+    description="輸入查詢以搜索相關文本。將顯示前3個最相關的結果。"
 )
 
-if __name__ == "__main__":
-    #iface.launch()
-    iface.launch(share=True) #網頁跑不出來
+# 啟動Gradio界面
+iface.launch()