Update app.py

app.py

@@ -1,4 +1,6 @@
 import streamlit as st
+import matplotlib.pyplot as plt
+import numpy as np
 from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
 
 # โหลด Tokenizer และ Model
@@ -9,6 +11,41 @@ model = AutoModelForTokenClassification.from_pretrained(model_name)
 # สร้าง NER Pipeline
 ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer)
 
+# กำหนดสีของ Entity แต่ละประเภท
+ENTITY_COLORS = {
+    "B-SKILL": "#FFD700",  # สีทอง
+    "I-SKILL": "#FFA500",  # สีส้ม
+    "B-TOOL": "#87CEFA",  # สีฟ้าอ่อน
+    "I-TOOL": "#1E90FF",  # สีฟ้าเข้ม
+    "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
+    
 # UI ด้วย Streamlit
 col1, col2, col3 = st.columns([4, 4, 4])
 
@@ -69,23 +106,19 @@ with col3:
         
     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)
-            entity_start = entity["start"]
-            entity_end = entity["end"]
-            annotated_entities.append({"Entity": entity_text, "Label": entity_label, "Start": entity_start, "End": entity_end})
-        
+            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:
-            start_pos = text.find(new_entity_text)
-            if start_pos != -1:
-                annotated_entities.append({"Entity": new_entity_text, "Label": new_entity_label, "Start": start_pos, "End": start_pos + len(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)