Update app.py

app.py

@@ -17,29 +17,65 @@ ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer)
 
 # กำหนดสีของ Entity แต่ละประเภท
 ENTITY_COLORS = {
-    "B-HSKILL": "#FFD700",  # สีทอง
-    "I-HSKILL": "#FFA500",  # สีส้ม
-    "B-SSKILL": "#FFD700",  # สีทอง
-    "I-SSKILL": "#FFA500",  # สีส้ม
-    "O": "#D3D3D3"  # สีเทา (Default)
+    "hskill": "#FFD700",  # สีทอง (Hard Skill)
+    "sskill": "#87CEFA",  # สีฟ้าอ่อน (Soft Skill)
 }
 
+# ฟังก์ชันรวม Entity (B-* และ I-*)
+def merge_entities(entities):
+    merged = []
+    current_entity = None
+
+    for entity in entities:
+        word, label = entity["word"], entity["entity"]
+
+        if label.startswith("B-"):  # ถ้าเป็น Entity ตัวแรก
+            if current_entity:  # บันทึก Entity ก่อนหน้า
+                merged.append(current_entity)
+            current_entity = {"word": word, "entity": label[2:]}  # ตัด B- ออก
+        elif label.startswith("I-") and current_entity:  # ถ้าเป็น Entity ที่ต่อเนื่อง
+            current_entity["word"] += " " + word  # รวมคำเข้าไป
+        else:
+            if current_entity:
+                merged.append(current_entity)
+            current_entity = {"word": word, "entity": label}
+
+    if current_entity:
+        merged.append(current_entity)
+
+    return merged
+
+# ฟังก์ชันแยกเป็น Hard Skill (hskill) และ Soft Skill (sskill)
+def categorize_entities(entities):
+    hskill = []
+    sskill = []
+    
+    for entity in entities:
+        if entity["entity"] in ["SKILL", "TOOL"]:  # รวม B-* และ I-* เป็นหมวดเดียว
+            hskill.append(entity)
+        else:
+            sskill.append(entity)
+    
+    return hskill, sskill
+
 # ฟังก์ชันสร้าง HTML สำหรับแสดงผล Entity ที่ถูก Highlight พร้อม Tag
-def highlight_entities(text, entities):
+def highlight_entities(text, hskill, sskill):
+    all_entities = hskill + sskill  # รวมทั้งหมดเพื่อไฮไลต์
     highlighted_text = text
-    for entity in sorted(entities, key=lambda e: e["start"], reverse=True):
+
+    for entity in sorted(all_entities, key=lambda e: text.find(e["word"]), reverse=True):
         entity_word = entity["word"]
-        entity_label = entity["entity"]
-        entity_color = ENTITY_COLORS.get(entity_label, "#D3D3D3")  # ใช้สีที่กำหนดไว้
+        entity_label = "hskill" if entity in hskill else "sskill"
+        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"]:]
+        highlighted_text = highlighted_text.replace(
+            entity_word,
+            f'<span style="background-color: {entity_color}; padding: 2px 5px; border-radius: 5px;">[{entity_word}] ({entity_label.upper()})</span>'
         )
 
     return f'<div style="font-size:16px; line-height:1.6;">{highlighted_text}</div>'
+
     
 # UI ด้วย Streamlit
 col1, col2, col3 = st.columns([4, 4, 4])
@@ -103,19 +139,28 @@ with col3:
         ner_results = ner_pipeline(text)
 
         if ner_results:
-            entity_list = [
-                {"word": entity['word'], "entity": entity['entity'], "start": entity['start'], "end": entity['end']}
+            # รวม Entity ที่มี B- และ I- ให้เป็นคำเดียวกัน
+            merged_entities = merge_entities([
+                {"word": entity['word'], "entity": entity['entity']}
                 for entity in ner_results
-            ]
+            ])
+
+            # แยกเป็น Hard Skill และ Soft Skill
+            hskill, sskill = categorize_entities(merged_entities)
 
             # แสดงผล Entity Highlighted Text พร้อม Tag
             st.markdown("### Annotated Text (Tagged)")
-            highlighted_html = highlight_entities(text, entity_list)
+            highlighted_html = highlight_entities(text, hskill, sskill)
             st.markdown(highlighted_html, unsafe_allow_html=True)
 
-            # แสดงข้อมูล Entity เป็นตาราง
-            st.markdown("### Extracted Entities")
-            st.table(entity_list)
+            # แสดงข้อมูล Hard Skill
+            st.markdown("### Hard Skills (hskill)")
+            st.table(hskill)
+
+            # แสดงข้อมูล Soft Skill
+            st.markdown("### Soft Skills (sskill)")
+            st.table(sskill)
+
         else:
             st.write("No entities found.")