mkoot007's picture
Create app.py
aa289c2
raw
history blame
No virus
788 Bytes
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)
def classify_text(text):
encoded_text = tokenizer(text, truncation=True, padding='max_length', max_length=512, return_tensors='pt')
predictions = model(**encoded_text)
predicted_label = predictions.logits.argmax(-1).item()
predicted_class = model.config.id2label[predicted_label]
return predicted_class
interface = gr.Interface(
fn=classify_text,
inputs=[gr.Textbox(label="Input Text")],
outputs=[gr.Textbox(label="Predicted Class")],
title="Text Classification App"
)
interface.launch(share=True)