Update app.py

app.py

@@ -51,8 +51,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
 
@@ -60,42 +59,10 @@ 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>"
-        
-        # 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>"
-    
+        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
-
 # 示例問題
 example_queries = [
     "Tell me about recent cyber attacks from Russia",
@@ -112,35 +79,38 @@ example_queries = [
 
 # 自定義 CSS
 custom_css = """
-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;}
+.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(class_name="input-row"):
-        query_input = gr.Textbox(lines=3, label="", placeholder="輸入你的查詢...", class_name="query-input")
-        submit_btn = gr.Button("查詢", variant="primary")
+    gr.Markdown("輸入查詢以搜索相關威脅情報，將顯示前3個最相關的結果，包括實體組。")
     
-    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]
-            )
+    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="")
 
     submit_btn.click(search_and_format, inputs=[query_input], outputs=[output])
+    #clear_btn.click(lambda: "", outputs=[query_input])
 
 # 啟動Gradio界面
 iface.launch()