Update app.py

app.py

@@ -1,10 +1,4 @@
 import streamlit as st
-
-import matplotlib
-matplotlib.use('Agg')  # ใช้ Agg เพื่อรองรับ Headless Environment
-import matplotlib.pyplot as plt
-
-import numpy as np
 from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
 
 # โหลด Tokenizer และ Model
@@ -17,8 +11,8 @@ ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer)
 
 # กำหนดสีของ Entity แต่ละประเภท
 ENTITY_COLORS = {
-    "hskill": "#FFD700",  # สีทอง (Hard Skill)
-    "sskill": "#87CEFA",  # สีฟ้าอ่อน (Soft Skill)
+    "HSKILL": "#FFD700",  # สีทอง
+    "SKILL": "#87CEFA",  # สีฟ้าอ่อน
 }
 
 # ฟังก์ชันรวม Entity (B-* และ I-*)
@@ -27,47 +21,35 @@ def merge_entities(entities):
     current_entity = None
 
     for entity in entities:
-        word, label = entity["word"], entity["entity"]
+        word, label, start, end = entity["word"], entity["entity"].split("-")[-1], entity["start"], entity["end"]
 
-        if label.startswith("B-"):  # ถ้าเป็น Entity ตัวแรก
-            if current_entity:  # บันทึก Entity ก่อนหน้า
+        if entity["entity"].startswith("B-"):  # ถ้าเป็น Entity ตัวแรก
+            if current_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  # รวมคำเข้าไป
+            current_entity = {"word": word, "entity": label, "start": start, "end": end}
+        elif entity["entity"].startswith("I-") and current_entity:  # ถ้าเป็น Entity ที่ต่อเนื่อง
+            current_entity["word"] += " " + word
+            current_entity["end"] = end  # ปรับตำแหน่ง end ใหม่
         else:
             if current_entity:
                 merged.append(current_entity)
-            current_entity = {"word": word, "entity": label}
+            current_entity = {"word": word, "entity": label, "start": start, "end": end}
 
     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 ["HSKILL"]:  # รวม B-* และ I-* เป็นหมวดเดียว
-            hskill.append(entity)
-        else:
-            sskill.append(entity)
-    
-    return hskill, sskill
-
-def highlight_entities(text, hskill, sskill):
 # ฟังก์ชันสร้าง Named Entity Annotation แสดงแบบ Tagging
-    all_entities = sorted(hskill + sskill, key=lambda e: e["start"])  # เรียงตามตำแหน่งเริ่มต้น
+def generate_ner_html(text, entities):
+    all_entities = sorted(entities, key=lambda e: e["start"])  # เรียงตามตำแหน่งเริ่มต้น
 
     last_idx = 0
     annotated_text = ""
 
     for entity in all_entities:
         start, end, word, entity_type = entity["start"], entity["end"], entity["word"], entity["entity"]
-        color = ENTITY_COLORS["hskill"] if entity in hskill else ENTITY_COLORS["sskill"]
+        color = ENTITY_COLORS.get(entity_type, ENTITY_COLORS["OTHER"])
 
         # เพิ่มข้อความก่อน Entity
         annotated_text += text[last_idx:start]
@@ -75,7 +57,7 @@ def highlight_entities(text, hskill, sskill):
         # เพิ่ม Entity แบบมีไฮไลต์และ Tagging
         annotated_text += f'''
             <span style="background-color: {color}; padding: 3px 6px; border-radius: 5px; font-weight: bold;">
-                {word} <span style="background: white; color: black; font-size: 12px; padding: 1px 4px; border-radius: 3px;">{entity_type.upper()}</span>
+                {word} <span style="background: white; color: black; font-size: 12px; padding: 1px 4px; border-radius: 3px;">{entity_type}</span>
             </span>
         '''
         last_idx = end  # อัปเดตตำแหน่งล่าสุด
@@ -85,91 +67,50 @@ def highlight_entities(text, hskill, sskill):
 
     return f'<div style="font-size:16px; line-height:1.6;">{annotated_text}</div>'
 
-    
 # UI ด้วย Streamlit
-col1, col2, col3 = st.columns([4, 4, 4])
+col1, col2 = st.columns([1, 1])
 
 with col1:
     st.header("Input")
-    default_text="""Job Description:
+    default_text = """Job Description:
 We are seeking a talented Software Engineer to join our dynamic team at Tech Innovations Inc. You will be responsible for designing, developing, and maintaining software applications that meet the needs of our clients.
-
 Key Responsibilities:
-
 Develop high-quality software design and architecture
 Identify, prioritize, and execute tasks in the software development life cycle
 Review and debug code
 Collaborate with other developers and engineers to ensure software quality
 Required Qualifications:
-
 Bachelor’s degree in Computer Science or related field
 Proven experience as a Software Engineer or similar role
 Familiarity with Agile development methodologies
 Proficiency in programming languages such as Java, Python, or C#
 Strong problem-solving skills and the ability to work in a team
 Preferred Qualifications:
-
-
-                """
-    text = st.text_area("Enter text for NER analysis:", value=default_text, height=400, max_chars=None, key=None, help=None, placeholder=None)
+"""
+    text = st.text_area("Enter text for NER analysis:", value=default_text, height=400)
     analyze_button = st.button("Analyze")
 
-    st.write("""**Example Inputs:**
-- Experience with cloud services (AWS, Azure)
-- Knowledge of databases (SQL, NoSQL)
-- Familiarity with front-end technologies (HTML, CSS, JavaScript)""")
-    
 with col2:
     st.header("Result")
 
-    # ใช้ st.markdown กับ CSS เพื่อปรับขนาดฟอนต์
-    st.markdown("<span style='font-size: 14px;'>Press button [Analyze]</span>", unsafe_allow_html=True)
-
-    if analyze_button:
-        ner_results = ner_pipeline(text)
-        
-        
-        # Display results in a structured output block
-        if ner_results:
-            output_data = [{"Entity": entity['word'], "Label": entity['entity'], "Score": f"{entity['score']:.4f}"} for entity in ner_results]
-            st.table(output_data)  # Display as a table
-        else:
-            st.write("No entities found.")
-
-        # ใช้ st.markdown กับ CSS เพื่อปรับขนาดฟอนต์
-        st.markdown("<span style='font-size: 14px;'>JSON</span>", unsafe_allow_html=True)
-
-        st.write(ner_results)
-        
-with col3:
-    st.header("Annotation")
-        
     if analyze_button:
         ner_results = ner_pipeline(text)
 
         if ner_results:
             # รวม Entity ที่มี B- และ I- ให้เป็นคำเดียวกัน
             merged_entities = merge_entities([
-                {"word": entity['word'], "entity": entity['entity']}
+                {"word": entity['word'], "entity": entity['entity'], "start": entity['start'], "end": entity['end']}
                 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, hskill, sskill)
-            st.markdown(highlighted_html, unsafe_allow_html=True)
+            # แสดงผล Named Entity Annotation
+            st.markdown("### Named Entity Annotation (Tagging)")
+            ner_html = generate_ner_html(text, merged_entities)
+            st.markdown(ner_html, unsafe_allow_html=True)
 
-            # แสดงข้อมูล Hard Skill
-            st.markdown("### Hard Skills (hskill)")
-            st.table(hskill)
-
-            # แสดงข้อมูล Soft Skill
-            st.markdown("### Soft Skills (sskill)")
-            st.table(sskill)
+            # แสดงข้อมูล Entity ในรูปแบบตาราง
+            st.markdown("### Extracted Entities")
+            st.table(merged_entities)
 
         else:
             st.write("No entities found.")
-