import streamlit as st
import torch as torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoConfig
from transformers import pipeline
from PIL import Image

model_name = "peace4ever/roberta-large-finetuned-mongolian_v3"  
pipeline = pipeline(task="sentiment-analysis", model=model_name)

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

if text is not None:
    
    col1, col2 = st.columns(2)

    predictions = pipeline(text)

    col2.header("Probabilities")
    for p in predictions:
        col2.subheader(f"{ p['label'] }: { round(p['score'] * 100, 1)}%")


# tokenizer = AutoTokenizer.from_pretrained(model_name)
# model = AutoModelForSequenceClassification.from_pretrained(model_name)

# 
# encoded_input = tokenizer(text, return_tensors="pt")
# output = model(**encoded_input)

# label_map = {"positive": 0, "negative": 1, "neutral": 2}

# # Update the model configuration with custom labels
# config = AutoConfig.from_pretrained(model_name)
# config.label2id = {"positive": 0, "negative": 1, "neutral": 2}
# config.id2label = {0: "positive", 1: "negative", 2: "neutral"}
# config.save_pretrained(model_name)

# predicted_label_id = torch.argmax(output.logits, dim=1).item()
# id2label = model.config.id2label
# predicted_label = id2label[predicted_label_id]

# print("Predicted Class:", predicted_label)

# st.json(predicted_label)