import streamlit as st
import torch
import pickle

# Load the saved model on the CPU
model = torch.load('saved_model.pth', map_location=torch.device('cpu'))

# Load the saved tokenizer
with open('tokenizer.pkl', 'rb') as f:
    tokenizer = pickle.load(f)

st.title("Text Classification Streamlit App")

input_text = st.text_input("Enter text:")
if st.button("Predict"):
    with torch.no_grad():
        inputs = tokenizer(input_text, return_tensors="pt")
        logits = model(**inputs).logits
        predicted_class = torch.argmax(logits, dim=1).item()

        st.write(f"Predicted Class: {predicted_class}")