外觀更新

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)
+        embeddings = torch.load(embeddings_path, weights_only=True)
         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 = 'sentence-transformers/all-MiniLM-L6-v2'
+model_name = 'sickcell/cti-semantic-search-minilm'
 data_file = 'labeled_cti_data.json'
 output_dir = '.'
 
@@ -55,7 +55,8 @@ 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'])
+            'entity_groups': get_entity_groups(data[idx]['entities']),
+            'entities': data[idx]['entities']
         })
     return results
 
@@ -63,24 +64,63 @@ 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"
-        formatted_results += f"   命名實體: {', '.join(result['entity_groups'])}\n\n"
+        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>"
+
     return formatted_results
 
-def transcribe_audio(audio):
+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 轉錄音檔"""
     try:
-        # 將音頻文件上傳到Whisper API
-        with open(audio, "rb") as audio_file:
+        with open(audio_path, "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 = [
@@ -98,40 +138,47 @@ example_queries = [
 
 # 自定義 CSS
 custom_css = """
-.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);}
+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;}
 """
 
 # 創建Gradio界面
 with gr.Blocks(css=custom_css) as iface:
     gr.Markdown("# AskCTI")
-    gr.Markdown("輸入查詢或使用語音輸入以查詢相關情資威脅情報，將顯示前3個最相關的結果。")
-    
-    with gr.Row(equal_height=True):
-        with gr.Column(scale=1, min_width=300):
-            query_input = gr.Textbox(lines=3, label="文字查詢")
+    gr.Markdown("輸入查詢或使用語音輸入以搜索相關威脅情報，將顯示前3個最相關的結果，包括實體標註。")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            query_input = gr.Textbox(lines=2, label="查詢", placeholder="輸入你的查詢...")
             with gr.Row():
-                submit_btn = gr.Button("查詢")
-                audio_input = gr.Audio(type="filepath", label="語音輸入")
-            
+                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")
+
             gr.Markdown("### 範例查詢")
-            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="語音轉錄結果")
+            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="結果")
 
     submit_btn.click(search_and_format, inputs=[query_input], outputs=[output])
-    audio_input.change(audio_to_search, inputs=[audio_input], outputs=[output, transcription_output, query_input])
+    audio_input.change(
+        fn=audio_to_text,  # 直接呼叫 audio_to_text 函數
+        inputs=[audio_input], 
+        outputs=[query_input]  # 將轉錄結果輸出到 query_input
+    )
 
 # 啟動Gradio界面
 iface.launch()