Update app.py

app.py

@@ -17,38 +17,29 @@ ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer)
 
 # กำหนดสีของ Entity แต่ละประเภท
 ENTITY_COLORS = {
-    "B-SKILL": "#FFD700",  # สีทอง
-    "I-SKILL": "#FFA500",  # สีส้ม
-    "B-TOOL": "#87CEFA",  # สีฟ้าอ่อน
-    "I-TOOL": "#1E90FF",  # สีฟ้าเข้ม
+    "B-HSKILL": "#FFD700",  # สีทอง
+    "I-HSKILL": "#FFA500",  # สีส้ม
+    "B-SSKILL": "#FFD700",  # สีทอง
+    "I-SSKILL": "#FFA500",  # สีส้ม
     "O": "#D3D3D3"  # สีเทา (Default)
 }
 
-# ฟังก์ชันวาด Entity Annotation เป็นภาพ
-def draw_annotation(text, entities):
-    fig, ax = plt.subplots(figsize=(12, len(text.split("\n")) * 0.5))
-    ax.set_xlim(0, 1)
-    ax.set_ylim(0, 1)
-    ax.axis("off")
-
-    # แยกบรรทัดของข้อความ
-    lines = text.split("\n")
-    y = 0.9  # ตำแหน่งเริ่มต้น
-
-    for line in lines:
-        words = line.split(" ")
-        x = 0.05  # ระยะห่างซ้ายสุด
-        for word in words:
-            color = "#FFFFFF"  # สีพื้นหลังปกติ
-            for entity in entities:
-                if entity["Entity"] == word:
-                    color = ENTITY_COLORS.get(entity["Label"], "#D3D3D3")  # เลือกสีตาม Label
-            
-            ax.text(x, y, word, fontsize=12, bbox=dict(facecolor=color, edgecolor="black", boxstyle="round,pad=0.3"))
-            x += (len(word) / 80) + 0.02  # ปรับระยะห่างของแต่ละคำ
-        y -= 0.05  # ลดระดับบรรทัดลง
-
-    return fig
+# ฟังก์ชันสร้าง HTML สำหรับแสดงผล Entity ที่ถูก Highlight พร้อม Tag
+def highlight_entities(text, entities):
+    highlighted_text = text
+    for entity in sorted(entities, key=lambda e: e["start"], reverse=True):
+        entity_word = entity["word"]
+        entity_label = entity["entity"]
+        entity_color = ENTITY_COLORS.get(entity_label, "#D3D3D3")  # ใช้สีที่กำหนดไว้
+
+        # ใช้ HTML + CSS ในการไฮไลต์ Entity และใส่ TAG
+        highlighted_text = (
+            highlighted_text[:entity["start"]]
+            + f'<span style="background-color: {entity_color}; padding: 2px 5px; border-radius: 5px;">[{entity_word}] ({entity_label})</span>'
+            + highlighted_text[entity["end"]:]
+        )
+
+    return f'<div style="font-size:16px; line-height:1.6;">{highlighted_text}</div>'
     
 # UI ด้วย Streamlit
 col1, col2, col3 = st.columns([4, 4, 4])
@@ -108,22 +99,23 @@ with col2:
 with col3:
     st.header("Annotation")
         
-    if analyze_button and ner_results:
-        st.write("Edit or Add Entities:")
-
-        annotated_entities = []
-        for i, entity in enumerate(ner_results):
-            entity_text = st.text_input(f"Entity {i+1}", value=entity['word'])
-            entity_label = st.selectbox(f"Label {i+1}", ["O", "B-SKILL", "I-SKILL", "B-TOOL", "I-TOOL"], index=0)
-            annotated_entities.append({"Entity": entity_text, "Label": entity_label})
-
-        # เพิ่ม Entity ใหม่
-        new_entity_text = st.text_input("New Entity")
-        new_entity_label = st.selectbox("New Label", ["O", "B-SKILL", "I-SKILL", "B-TOOL", "I-TOOL"], index=0)
-        if st.button("Add New Entity") and new_entity_text:
-            annotated_entities.append({"Entity": new_entity_text, "Label": new_entity_label})
-
-        if st.button("Save Annotation"):
-            st.write("Saved Annotations:")
-            st.json(annotated_entities)
+    if analyze_button:
+        ner_results = ner_pipeline(text)
+
+        if ner_results:
+            entity_list = [
+                {"word": entity['word'], "entity": entity['entity'], "start": entity['start'], "end": entity['end']}
+                for entity in ner_results
+            ]
+
+            # แสดงผล Entity Highlighted Text พร้อม Tag
+            st.markdown("### Annotated Text (Tagged)")
+            highlighted_html = highlight_entities(text, entity_list)
+            st.markdown(highlighted_html, unsafe_allow_html=True)
+
+            # แสดงข้อมูล Entity เป็นตาราง
+            st.markdown("### Extracted Entities")
+            st.table(entity_list)
+        else:
+            st.write("No entities found.")