File size: 1,392 Bytes
dc47816 b0e39c2 dc47816 b0e39c2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch
# Model and tokenizer loading
model_id = "cheberle/autotrain-35swc-b4r9z"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSequenceClassification.from_pretrained(model_id)
# Move model to GPU if available
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
def predict(text):
# Tokenize input
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
# Move inputs to same device as model
inputs = {k: v.to(device) for k, v in inputs.items()}
# Get prediction
with torch.no_grad():
outputs = model(**inputs)
predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
# Get prediction probabilities and labels
probs = predictions[0].tolist()
labels = model.config.id2label
# Create formatted output
results = {labels[i]: float(probs[i]) for i in range(len(probs))}
return results
# Create Gradio interface
iface = gr.Interface(
fn=predict,
inputs=gr.Textbox(label="Input Text"),
outputs=gr.Label(label="Prediction"),
title="Model Prediction Interface",
description=f"Enter text to get predictions from {model_id}",
examples=["Example text to try"]
)
# Launch the interface
iface.launch() |