Update app.py

app.py

@@ -1,18 +1,12 @@
-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
+from sparknlp_display import Display
+from pyspark.sql import SparkSession
 
-# ตรวจสอบว่ามีการติดตั้งโมเดล spacy หรือไม่
-try:
-    nlp = spacy.load("en_core_web_sm")
-except OSError:
-    import os
-    os.system("python -m spacy download en_core_web_sm")
-    nlp = spacy.load("en_core_web_sm")
+# สร้าง Spark Session
+spark = SparkSession.builder \
+    .appName("NER Analysis") \
+    .getOrCreate()
 
 # โหลดโมเดล NER จาก Hugging Face
 model_name = "Nucha/Nucha_SkillNER_BERT"
@@ -23,23 +17,17 @@ model = AutoModelForTokenClassification.from_pretrained(model_name)
 ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple")
 
 # UI ด้วย Streamlit
-st.title("NER Analysis with Nucha SkillNER BERT and spacy displacy")
+st.title("NER Analysis with Nucha SkillNER BERT and Spark NLP Display")
 
 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)
+
+    # สร้าง DataFrame สำหรับผลลัพธ์
+    data = [{"word": entity['word'], "start": entity['start'], "end": entity['end'], "label": entity['entity']} for entity in ner_results]
+    ner_df = spark.createDataFrame(data)
+
+    # แสดงผลด้วย sparknlp_display
+    display = Display()
+    st.write(display.display(ner_df, "word", "label"))