7jimmy's picture
Create app.py
bc7757c
import streamlit as st
import transformers
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import numpy as np
# Load the pre-trained text classification model from Hugging Face
model_name = "bert-base-uncased"
num_labels = 2
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=num_labels)
tokenizer = AutoTokenizer.from_pretrained(model_name)
def classify_text(text):
# Preprocess the text input
encoded_text = tokenizer(text, truncation=True, padding=True, return_tensors="pt")
# Make predictions using the pre-trained model
with torch.no_grad():
outputs = model(**encoded_text)
logits = outputs.logits
predictions = np.argmax(logits, axis=1)
# Convert predictions to class labels
class_labels = ["positive", "negative"]
predicted_labels = [class_labels[i] for i in predictions]
# Return the predicted labels
return predicted_labels
# Initialize the Streamlit app
st.title("Text Classification Demo")
# Create the text input field
input_text = st.text_input("Enter text to classify:", "")
# Make predictions and display the results
if input_text:
predicted_labels = classify_text(input_text)
st.write("Predicted labels:", predicted_labels)