Update app.py

app.py

@@ -16,7 +16,7 @@ def load_or_create_model_and_embeddings(model_name, data_file, output_dir):
     if os.path.exists(model_path) and os.path.exists(embeddings_path):
         print("載入已保存的模型和嵌入...")
         model = SentenceTransformer(model_path)
-        embeddings = torch.load(embeddings_path, weights_only=True)
+        embeddings = torch.load(embeddings_path)
         with open(data_file, 'r', encoding='utf-8') as f:
             data = json.load(f)
     else:
@@ -32,7 +32,7 @@ def load_or_create_model_and_embeddings(model_name, data_file, output_dir):
     return model, embeddings, data
 
 # 設置參數
-model_name = 'sickcell/cti-semantic-search-minilm'
+model_name = 'sentence-transformers/all-MiniLM-L6-v2'
 data_file = 'labeled_cti_data.json'
 output_dir = '.'
 
@@ -55,8 +55,7 @@ def semantic_search(query, top_k=3):
         results.append({
             'text': data[idx]['text'],
             'similarity_score': 1 - distances[0][i] / 2,
-            'entity_groups': get_entity_groups(data[idx]['entities']),
-            'entities': data[idx]['entities']
+            'entity_groups': get_entity_groups(data[idx]['entities'])
         })
     return results
 
@@ -64,63 +63,24 @@ def search_and_format(query):
     results = semantic_search(query)
     formatted_results = ""
     for i, result in enumerate(results, 1):
-        formatted_results += f"<h3>結果 {i}:</h3>"
-        formatted_results += "<h4>NER 定義</h4>"
-
-        words = result['text'].split()
-
-        color_map = {
-            'PERSON': 'lightpink',
-            'ORG': 'lightblue',
-            'PLACE': 'lightyellow',
-            'TECHNOLOGY': 'lightgreen',
-            'MALWARE': 'plum',
-            'ATTACK': 'peachpuff'
-        }
-
-        formatted_text = []
-        for word in words:
-            found = False
-            for entity in result['entities']:
-                if word in entity['word']:
-                    color = color_map.get(entity['entity_group'], 'lightgray')
-                    formatted_word = f'<span style="background-color: {color};">{word} <sup>{entity["entity_group"]}</sup></span>'
-                    formatted_text.append(formatted_word)
-                    found = True
-                    break
-            if not found:
-                formatted_text.append(word)
-
-        formatted_results += ' '.join(formatted_text) + "<br><br>"
-        formatted_results += f"<strong>相似度分數:</strong> {result['similarity_score']:.4f}<br><br>"
-
+        formatted_results += f"{i}. 相似度分數: {result['similarity_score']:.4f}\n"
+        formatted_results += f"   情資: {result['text']}\n"
+        formatted_results += f"   命名實體: {', '.join(result['entity_groups'])}\n\n"
     return formatted_results
 
-def audio_to_text(audio_data):
-    """將音檔資料轉錄為文字"""
-    # 顯示載入動畫
-    query_input.update(value="正在轉錄中...")
-
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_audio:
-        temp_audio.write(audio_data)
-        temp_audio_path = temp_audio.name
-
-    transcription = transcribe_audio(temp_audio_path)
-
-    os.remove(temp_audio_path)
-
-    # 更新查詢框
-    query_input.update(value=transcription)
-
-def transcribe_audio(audio_path):
-    """使用 OpenAI Whisper API 轉錄音檔"""
+def transcribe_audio(audio):
     try:
-        with open(audio_path, "rb") as audio_file:
+        # 將音頻文件上傳到Whisper API
+        with open(audio, "rb") as audio_file:
             transcript = openai.Audio.transcribe("whisper-1", audio_file)
         return transcript.text
     except Exception as e:
         return f"轉錄時發生錯誤: {str(e)}"
 
+def audio_to_search(audio):
+    transcription = transcribe_audio(audio)
+    search_results = search_and_format(transcription)
+    return search_results, transcription, transcription
 
 # 示例問題
 example_queries = [
@@ -138,47 +98,40 @@ example_queries = [
 
 # 自定義 CSS
 custom_css = """
-body {font-family: Arial, sans-serif;}
-.container {max-width: 1200px; margin: auto;}
-.input-row {display: flex; gap: 10px; margin-bottom: 20px; align-items: flex-end;}
-.query-input {flex-grow: 1;}
-.output-area {border: 1px solid #ddd; padding: 15px; border-radius: 5px;}
-.examples-grid {display: grid; grid-template-columns: repeat(auto-fill, minmax(200px, 1fr)); gap: 10px; margin-top: 20px;}
-.example-button {width: 100%;}
-span sup {font-size: 0.7em; font-weight: bold;}
-/* 新增的樣式 */
-.small-button {padding: 5px 10px; font-size: 0.9em;}
+.container {display: flex; flex-direction: row;}
+.input-column {flex: 1; padding-right: 20px;}
+.output-column {flex: 2;}
+.examples-list {display: flex; flex-wrap: wrap; gap: 10px;}
+.examples-list > * {flex-basis: calc(50% - 5px);}
 """
 
 # 創建Gradio界面
 with gr.Blocks(css=custom_css) as iface:
     gr.Markdown("# AskCTI")
-    gr.Markdown("輸入查詢或使用語音輸入以搜索相關威脅情報，將顯示前3個最相關的結果，包括實體標註。")
-
-    with gr.Row():
-        with gr.Column(scale=1):
-            query_input = gr.Textbox(lines=2, label="查詢", placeholder="輸入你的查詢...")
+    gr.Markdown("輸入查詢或使用語音輸入以查詢相關情資威脅情報，將顯示前3個最相關的結果。")
+    
+    with gr.Row(equal_height=True):
+        with gr.Column(scale=1, min_width=300):
+            query_input = gr.Textbox(lines=3, label="文字查詢")
             with gr.Row():
-                submit_btn = gr.Button("查詢", variant="primary", elem_classes="small-button")
-                #audio_input = gr.Audio(source="microphone", label="錄音", elem_classes="small-button")
-                audio_input = gr.Audio(sources="microphone", label="錄音", elem_classes="small-button")
-
+                submit_btn = gr.Button("查詢")
+                audio_input = gr.Audio(type="filepath", label="語音輸入")
+            
             gr.Markdown("### 範例查詢")
-            example_buttons = []
-            for query in example_queries:
-                btn = gr.Button(query)
-                btn.click(lambda x: x, inputs=[gr.Textbox(value=query, visible=False)], outputs=[query_input])
-                example_buttons.append(btn)
-
-        with gr.Column(scale=1):
-            output = gr.HTML(label="結果")
+            for i in range(0, len(example_queries), 2):
+                with gr.Row():
+                    for j in range(2):
+                        if i + j < len(example_queries):
+                            gr.Button(example_queries[i+j]).click(
+                                lambda x: x, inputs=[gr.Textbox(value=example_queries[i+j], visible=False)], outputs=[query_input]
+                            )
+        
+        with gr.Column(scale=2):
+            output = gr.Textbox(lines=20, label="查詢結果")
+            transcription_output = gr.Textbox(lines=3, label="語音轉錄結果")
 
     submit_btn.click(search_and_format, inputs=[query_input], outputs=[output])
-    audio_input.change(
-        fn=audio_to_text,  # 直接呼叫 audio_to_text 函數
-        inputs=[audio_input], 
-        outputs=[query_input]  # 將轉錄結果輸出到 query_input
-    )
+    audio_input.change(audio_to_search, inputs=[audio_input], outputs=[output, transcription_output, query_input])
 
 # 啟動Gradio界面
 iface.launch()