import streamlit as st
import pandas as pd
import torch
from transformers import RobertaTokenizer, RobertaForSequenceClassification

# Load model
model_name = "syedkhalid076/RoBERTa-Sentimental-Analysis-v2"
tokenizer = RobertaTokenizer.from_pretrained(model_name)
model = RobertaForSequenceClassification.from_pretrained(model_name)
model.eval()

# Define sentiment labels
sentiment_labels = {0: "Negative", 1: "Neutral", 2: "Positive"}

# Function to predict sentiment for a single sentence
def predict_sentiment(sentence):
    inputs = tokenizer(sentence, return_tensors="pt", max_length=512, truncation=True)
    outputs = model(**inputs)
    logits = outputs.logits.detach().cpu()
    predicted_class = torch.argmax(logits, dim=-1).item()
    sentiment = sentiment_labels[predicted_class]
    return sentiment

# Function to process CSV file and predict sentiment for each row
def process_csv(file):
    df = pd.read_csv(file)
    if 'Text' not in df.columns:
        st.error("CSV file must have a 'Text' column with sentences for analysis.")
        return None
    df['sentiment'] = df['text'].apply(predict_sentiment)
    return df

# Streamlit app
def main():
    st.title("Sentiment Analysis App")
    st.write("Analyze text sentiment as Negative, Neutral, or Positive.")
    st.write("NOTE: If uploading a CSV file, ensure the column containing text is named 'text' (case-sensitive).")

    option = st.radio("Choose input type:", ("Write a sentence", "Upload a CSV file"))

    if option == "Write a sentence":
        sentence = st.text_input("Enter a sentence:")
        if st.button("Analyze"):
            if sentence.strip():
                sentiment = predict_sentiment(sentence)
                st.write("Sentiment:", sentiment)
            else:
                st.warning("Please enter a valid sentence.")

    elif option == "Upload a CSV file":
        file = st.file_uploader("Upload CSV file", type=['csv'])
        if file is not None:
            df = process_csv(file)
            if df is not None:
                st.write(df)

if __name__ == '__main__':
    main()