app.py

import streamlit as st
from transformers import pipeline

# Initialize the POS tagger
pos_tagger = pipeline(
    model="vblagoje/bert-english-uncased-finetuned-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)
    highlighted_text = ""
    for word_info in pos_result:
        color = get_color(word_info["entity_group"])
        highlighted_text += f'<span style="color:{color};">{word_info["word"]}</span> '

    st.markdown(highlighted_text, unsafe_allow_html=True)