import streamlit as st
from transformers import pipeline

# Initialize the POS tagger with from_pt=True for PyTorch model
pos_tagger = pipeline(
    model="projecte-aina/roberta-base-ca-cased-pos",
    aggregation_strategy="simple",
 
)

# Function to get color for each POS tag
def get_color(pos_tag):
    colors = {
        "PRON": "blue",
        "VERB": "green",
        "ADP": "red",
        "PROPN": "orange",
        # Add more POS tags and colors as needed
    }
    return colors.get(pos_tag, "grey")  # Default color

# Streamlit app
st.title("Part-of-Speech Highlighter")

# Text input
text = st.text_area("Enter your text here:", "I am an artist and I live in Dublin")

# Process text and highlight POS
if text:
    pos_result = pos_tagger(text)
    st.markdown(pos_result)