Rishi-19's picture
Rename main.py to app.py
c04acaa
raw
history blame contribute delete
614 Bytes
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}")