change output format

app.py

@@ -51,7 +51,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
 
@@ -59,10 +60,42 @@ 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>"
+        
+        # Split the text into words
+        words = result['text'].split()
+        
+        # Create a dictionary to map entity groups to colors
+        color_map = {
+            'PERSON': 'lightpink',
+            'ORG': 'lightblue',
+            'PLACE': 'lightyellow',
+            'TECHNOLOGY': 'lightgreen',
+            'MALWARE': 'plum',
+            'ATTACK': 'peachpuff'
+        }
+        
+        # Format each word based on its entity group
+        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)
+        
+        # Join the formatted words back into a sentence
+        formatted_results += ' '.join(formatted_text) + "<br><br>"
+        formatted_results += f"<strong>相似度分數:</strong> {result['similarity_score']:.4f}<br><br>"
+    
     return formatted_results
+
 # 示例問題
 example_queries = [
     "Tell me about recent cyber attacks from Russia",
@@ -79,38 +112,35 @@ 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: 900px; margin: auto;}
+.input-row {display: flex; gap: 10px; margin-bottom: 20px;}
+.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;}
 """
 
 # 創建Gradio界面
 with gr.Blocks(css=custom_css) as iface:
     gr.Markdown("# AskCTI")
-    gr.Markdown("輸入查詢以搜索相關威脅情報，將顯示前3個最相關的結果，包括實體組。")
+    gr.Markdown("輸入查詢以搜索相關威脅情報，將顯示前3個最相關的結果，包括實體標註。")
     
-    with gr.Row(equal_height=True):
-        with gr.Column(scale=1, min_width=300):
-            query_input = gr.Textbox(lines=3, label="")
-            submit_btn = gr.Button("查詢")
-            #clear_btn = gr.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="")
+    with gr.Row(class_name="input-row"):
+        query_input = gr.Textbox(lines=3, label="", placeholder="輸入你的查詢...", class_name="query-input")
+        submit_btn = gr.Button("查詢", variant="primary")
+    
+    output = gr.HTML(class_name="output-area")
+    
+    gr.Markdown("### 範例查詢")
+    with gr.Row(class_name="examples-grid"):
+        for query in example_queries:
+            gr.Button(query, class_name="example-button").click(
+                lambda x: x, inputs=[gr.Textbox(value=query, visible=False)], outputs=[query_input]
+            )
 
     submit_btn.click(search_and_format, inputs=[query_input], outputs=[output])
-    #clear_btn.click(lambda: "", outputs=[query_input])
 
 # 啟動Gradio界面
 iface.launch()