import streamlit as st

from transformers import pipeline

model_name = "peace4ever/roberta-large-finetuned-mongolian_v4"
pipeline = pipeline(task="sentiment-analysis", model=model_name)  # Load pre-trained pipeline

st.title("Эерэг? Сөрөг эсвэл аль нь ч биш?")
text = st.text_area("Өгүүлбэр оруулна уу?")

if text is not None:
    predictions = pipeline(text)
    label = predictions[0]["label"]
    probability = predictions[0]["score"]
    col1, col2 = st.columns(2)
    col1.header("Sentiment")
    col2.header("Probability")

    if label == "entailment":
        sentiment = "Negative"
    elif label == "contradiction":
        sentiment = "Neutral"
    elif label == "neutral":
        sentiment = "Positive" 

    col1.write(sentiment)
    col2.write(f"{probability:.2f}") 
    # label = predictions[0]["label"]
    # probability = predictions[0]["score"]
    # col1.write(label)
    # col2.write(f"{probability:.2f}")