app.py

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
model_name = "Nucha/Nucha_ITSkillNER_BERT"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForTokenClassification.from_pretrained(model_name)

# สร้าง NER Pipeline
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)
}

# ฟังก์ชันสร้าง 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])

with col1:
    st.header("Input")
    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)
    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_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.")