import streamlit as st
import torch
from peft import (
    PeftModel,
    PeftConfig,
)

from transformers import AutoModelForSequenceClassification, AutoTokenizer
from Preprocessor import preprocess


peft_model = "VRLLab/TurkishBERTweet-Lora-SA"
peft_config = PeftConfig.from_pretrained(peft_model)

# loading Tokenizer
padding_side = "right"
tokenizer = AutoTokenizer.from_pretrained(
    peft_config.base_model_name_or_path, padding_side=padding_side
)
if getattr(tokenizer, "pad_token_id") is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id

id2label_sa = {0: "negative", 2: "positive", 1: "neutral"}
turkishBERTweet_sa = AutoModelForSequenceClassification.from_pretrained(
    peft_config.base_model_name_or_path,
    return_dict=True,
    num_labels=len(id2label_sa),
    id2label=id2label_sa,
)
turkishBERTweet_sa = PeftModel.from_pretrained(turkishBERTweet_sa, peft_model)


st.title("TurkishBERTweet-Lora-SA")
st.write("Enter a sentence to analyze its sentiment:")


user_input = st.text_input("")
if user_input:
    with torch.no_grad():
        ids = tokenizer.encode_plus(preprocess(user_input), return_tensors="pt")
        logits = turkishBERTweet_sa(**ids).logits
        label_id = logits.argmax(-1).item()
        confidence = logits.softmax(-1)[0, label_id].item()
        st.write(f"Sentiment: {id2label_sa[label_id]}")
        st.write(f"Confidence: {confidence:.2f}")