Update app.py

app.py

@@ -1,24 +1,40 @@
 import os
 os.system("python -m spacy download en_core_web_sm")
 
-
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
+import spacy
+from spacy import displacy
+
+# โหลด Spacy Model
+nlp = spacy.load("en_core_web_sm")
 
-# โหลด Tokenizer และ Model
-model_name = "Nucha/NuchaBertBaseUncased"
+# โหลดโมเดล NER จาก Hugging Face
+model_name = "Nucha/Nucha_SkillNER_BERT"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForTokenClassification.from_pretrained(model_name)
 
-# สร้าง NER Pipeline
-ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer)
+# สร้าง pipeline สำหรับ NER
+ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple")
 
 # UI ด้วย Streamlit
-st.title("NER Analysis App")
+st.title("NER Analysis with Nucha SkillNER BERT and displacy")
+
 text = st.text_area("Enter text for NER analysis:")
 
 if st.button("Analyze"):
     ner_results = ner_pipeline(text)
+    
+    # เตรียมข้อมูลสำหรับ displacy
+    ents = []
     for entity in ner_results:
-        if entity['entity']!="LABEL_0":
-            st.write(f"Entity: {entity['word']}, Label: {entity['entity']}, Score: {entity['score']:.4f}")
+        ents.append({
+            "start": entity['start'],
+            "end": entity['end'],
+            "label": entity['entity'],
+        })
+    
+    # แสดงผล displacy ผ่าน Streamlit
+    options = {"colors": {"SKILL": "lightblue"}}  # เพิ่มสีให้แต่ละ label ถ้าต้องการ
+    html = displacy.render({"text": text, "ents": ents}, style="ent", manual=True, options=options)
+    st.write(html, unsafe_allow_html=True)