import streamlit as st
import spacy
from transformers import pipeline
from PIL import Image

# Check if the model is installed, if not, download it
try:
    nlp = spacy.load("en_core_web_sm")
except OSError:
    st.write("Downloading spaCy model...")
    subprocess.run(["python", "-m", "spacy", "download", "en_core_web_sm"])
    nlp = spacy.load("en_core_web_sm")

st.title("SpaGAN Demo")
st.write("Enter a text, and the system will highlight the geo-entities within it.")

user_input = st.text_area("Input Text", height=200)

# Process the text when the button is clicked
if st.button("Highlight Geo-Entities"):
    if user_input.strip():
        # Process the text using spaCy
        doc = nlp(user_input)

        # Highlight geo-entities
        highlighted_text = user_input
        for ent in reversed(doc.ents):
            if ent.label_ in ["GPE", "LOC"]:  # GPE = Geopolitical Entity, LOC = Location
                highlighted_text = (
                    highlighted_text[:ent.start_char] +
                    f"**:green[{ent.text}]**" + 
                    highlighted_text[ent.end_char:]
                )

        # Display the highlighted text
        st.markdown(highlighted_text)
    else:
        st.error("Please enter some text.")

st.write("Note: The model identifies and highlights geo-entities.")