import streamlit as st
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

# Set up the device (GPU or CPU)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Function to perform sentiment analysis
def perform_sentiment_analysis(text):
    inputs = tokenizer(text, padding=True, truncation=True, return_tensors="pt")
    inputs = inputs.to(device)
    outputs = model(**inputs)
    logits = outputs.logits
    probabilities = torch.softmax(logits, dim=1).detach().cpu().numpy()[0]
    sentiment_label = "Positive" if probabilities[1] > probabilities[0] else "Negative"
    return sentiment_label, probabilities

# Streamlit app
def main():
    st.title("Sentiment Analysis App")
    st.write("Enter a text and select a pretrained model to perform sentiment analysis.")

    text = st.text_area("Enter text", value="")

    model_options = {
        "distilbert-base-uncased-finetuned-sst-2-english": "DistilBERT (SST-2)",
        "distilbert-base-uncased": "DistilBERT Uncased",
        "bert-base-uncased": "BERT Uncased",
        "bert-base-cased": "BERT Cased",
        "bert-large-uncased": "BERT Large Uncased",
        "bert-large-cased": "BERT Large Cased",
        "roberta-base": "RoBERTa Base",
        "roberta-large": "RoBERTa Large",
        "albert-base-v2": "ALBERT Base v2",
        "albert-large-v2": "ALBERT Large v2",
        "google/electra-base-discriminator": "Electra Base Discriminator",
        "google/electra-large-discriminator": "Electra Large Discriminator",
        "xlnet-base-cased": "XLNet Base Cased",
        "xlnet-large-cased": "XLNet Large Cased",
        "gpt2": "GPT2",
        "gpt2-medium": "GPT2 Medium",
        "gpt2-large": "GPT2 Large",
        "gpt2-xl": "GPT2 XL",
        # Add more models here if desired
    }

    selected_model = st.selectbox("Select a pretrained model", list(model_options.keys()))

    # Load the pretrained model and tokenizer
    model_name = selected_model
    model = AutoModelForSequenceClassification.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    if st.button("Analyze"):
        sentiment_label, probabilities = perform_sentiment_analysis(text)
        st.write(f"Sentiment: {sentiment_label}")
        st.write(f"Positive probability: {probabilities[1]}")
        st.write(f"Negative probability: {probabilities[0]}")

if __name__ == "__main__":
    main()