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

# Load the model and tokenizer from Hugging Face
model_name = "fajjos/keyword_variable_detection_v1"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Streamlit UI
st.title("Keyword and Variable Extraction")
st.write("Instruction: Extract keywords and variables from the prompt.")

# Input for user text
user_input = st.text_area("Input Text", "University of Oxford timeshigher ranking")

if st.button("Predict"):
    # Tokenize the input
    inputs = tokenizer(user_input, return_tensors="pt")

    # Run inference
    with torch.no_grad():
        outputs = model(**inputs)
    
    # Process the outputs
    predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
    predicted_class = predictions.argmax().item()
    
    st.write(f"Predicted class: {predicted_class}, Confidence: {predictions[0][predicted_class]:.2f}")