NuchaITSkillNER / app.py
Nucha's picture
Update app.py
7abeff1 verified
raw
history blame
1.45 kB
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")
# โหลดโมเดล NER จาก Hugging Face
model_name = "Nucha/Nucha_SkillNER_BERT"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForTokenClassification.from_pretrained(model_name)
# สร้าง pipeline สำหรับ NER
ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple")
# UI ด้วย Streamlit
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:
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)