import streamlit as st
import torch
from transformers import DistilBertForSequenceClassification, DistilBertTokenizer
from torch.nn.functional import softmax

# Load the model and tokenizer
model = DistilBertForSequenceClassification.from_pretrained('./fine_tuned_distilbert')
tokenizer = DistilBertTokenizer.from_pretrained('./fine_tuned_distilbert')

# Device setup
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

mapping = {"Remembering": 0, "Understanding": 1, "Applying": 2, "Analyzing": 3, "Evaluating": 4, "Creating": 5}

# Reverse the mapping to get the class name from the index
reverse_mapping = {v: k for k, v in mapping.items()}

def predict_with_loaded_model(text):
    # Tokenize the input text
    inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=512)
    input_ids = inputs['input_ids'].to(device)

    model.eval()
    with torch.no_grad():
        # Get the raw logits from the model
        outputs = model(input_ids)
        logits = outputs.logits
        
        # Apply softmax to get probabilities
        probabilities = softmax(logits, dim=-1)
    
    # Convert probabilities to a list or dictionary of class probabilities
    probabilities = probabilities.squeeze().cpu().numpy()
    
    # Map the probabilities to the class labels using the reverse mapping
    class_probabilities = {reverse_mapping[i]: prob for i, prob in enumerate(probabilities)}
    
    return class_probabilities

# Streamlit App
st.title("Question Bloom Score Prediction")

# Create an input box for the user to enter a question
question = st.text_area("Enter a question:")

# If a question is entered, make the prediction
if question:
    class_probabilities = predict_with_loaded_model(question)

    # Display the probabilities for each class label
    st.write("**Class Probabilities (Bloom Scores)**")
    for class_label, prob in class_probabilities.items():
        st.write(f"{class_label}: {prob:.4f}")