Spaces:
Sleeping
Sleeping
import streamlit as st | |
from transformers import BertTokenizer, BertForSequenceClassification | |
import torch | |
# Load the pre-trained BERT model and tokenizer | |
model_name = "bert-base-uncased" # Replace with your specific model if needed | |
tokenizer = BertTokenizer.from_pretrained(model_name) | |
model = BertForSequenceClassification.from_pretrained(model_name) | |
# Streamlit UI | |
st.title("BERT Text Classification") | |
text_input = st.text_input("Enter text for classification:") | |
if text_input: | |
# Tokenize input | |
inputs = tokenizer(text_input, return_tensors="pt", truncation=True, padding=True, max_length=512) | |
with torch.no_grad(): | |
logits = model(**inputs).logits | |
predicted_class = torch.argmax(logits, dim=1).item() | |
st.write(f"Predicted Class: {predicted_class}") |