import streamlit as st

import torch as torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoConfig


# import requests
# from transformers import pipeline

# # Load the model using its ID or name
# model_id_or_name = "peace4ever/roberta-large-finetuned-mongolian_v4"  # Update with your model ID or name
# classifier = pipeline("sentiment-analysis", model=model_id_or_name)

# # Example usage
# result = classifier("I loved Star Wars so much!")
# print(result)


model_name = "peace4ever/roberta-large-finetuned-mongolian_v3"  
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

text = st.text_area("Өгүүлбэр оруулна уу?")
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)