import streamlit as st
import altair as alt
import torch
from transformers import AlbertTokenizer, AlbertForSequenceClassification
import sentencepiece as spm

# Load pre-trained model and tokenizer
model_name = "albert-base-v2"
tokenizer = AlbertTokenizer.from_pretrained(model_name)
model = AlbertForSequenceClassification.from_pretrained(model_name)

# Define function to classify input text
def classify_text(text):
    inputs = tokenizer(text, padding=True, truncation=True, return_tensors="pt")
    outputs = model(**inputs)
    logits = outputs.logits.detach().numpy()[0]
    probabilities = torch.softmax(torch.tensor(logits), dim=0).tolist()
    return probabilities

# Set up Streamlit app
st.title("ALBERT Text Classification App")

# Create input box for user to enter text
text_input = st.text_area("Enter text to classify", height=200)

# Classify input text and display results
if st.button("Classify"):
    if text_input:
        probabilities = classify_text(text_input)
        df = pd.DataFrame({
            'Label': ['Negative', 'Positive'],
            'Probability': probabilities
        })
        chart = alt.Chart(df).mark_bar().encode(
            x='Probability',
            y=alt.Y('Label', sort=['Negative', 'Positive'])
        )
        st.write(chart)
    else:
        st.write("Please enter some text to classify.")