import gradio as gr
from transformers import pipeline
import re

text_pipe = pipeline("text-classification", model="karalif/myTestModel", return_all_scores=True)

def predict(text):
    greeting_pattern = r"^(Halló|Hæ|Sæl|Góða|Kær|Daginn|Kvöldið|Ágæt|Elsku)"
    
    greeting_feedback = ""
    
    results = text_pipe(text)
    all_scores = results[0]
    response = ""
    for result in all_scores:
        label = result['label']
        score = result['score']
        if label in ["Politeness"]:
            response += f"<span style='background-color:#b8e994; color:black;'>{label}</span>: {score:.3f}<br>"  # Light Green
        if label in ["Sentiment"]:
            response += f"<span style='background-color:#fff3cd; color:black;'>{label}</span>: {score:.3f}<br>"  # Light Yellow
        if label in ["Formality"]:
            response += f"<span style='background-color:#bee5eb; color:black;'>{label}</span>: {score:.3f}<br>"  # Light Blue
        if label in ["Toxicity"]:
            response += f"<span style='background-color:#f8d7da; color:black;'>{label}</span>: {score:.3f}<br>"  # Light Red
        else:
            response += f"{label}: {score:.3f}<br>"
    
    if not re.match(greeting_pattern, text, re.IGNORECASE):
        greeting_feedback = "Heilsaðu dóninn þinn<br>"

    response += greeting_feedback
    
    return response

description_html = """
    <center>
        <img src='http://www.ru.is/media/HR_logo_vinstri_transparent.png' width='250' height='auto'>
    </center>
    """
textbox = gr.Textbox()

demo = gr.Interface(
    fn=predict,
    inputs=gr.TextArea(label="Enter text here:"),
    outputs=gr.HTML(label="Leiðrétt"),
    description=description_html,
    examples=[
        ["Það voru vitni að árásinni sem tilkynntu málið til lögreglu sem kom skjótt á vettvang."],
        ["Ég held þetta sé ekki góður tími fara heimsókn."],
        ["Sæl og blessuð Kristín, hvað er að frella af þér gamla??"],
        ["Hver á þenan bússtað? já eða nei."],
        ["Hafi þau svo látið gólfið þorna vel og síðan flotað það til lagfæringar eftir motturnar."],
    ],
    theme=gr.themes.Default(primary_hue="red", secondary_hue="pink")
)

demo.launch()