laptop_guru / app.py
tdnathmlenthusiast's picture
Update app.py
40a098b
raw
history blame
No virus
892 Bytes
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
# Initialize the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("laptop_data.pkl")
model = AutoModelForSequenceClassification.from_pretrained("laptop_data.pkl")
# Define the function for classifying laptops
def classify_laptop(description):
inputs = tokenizer(description, return_tensors="pt", padding=True, truncation=True)
outputs = model(**inputs)
logits = outputs.logits
probabilities = torch.softmax(logits, dim=1)
return {label: prob.item() for label, prob in zip(model.config.id2label.values(), probabilities[0])}
# Create the Gradio interface
iface = gr.Interface(
fn=classify_laptop,
inputs=gr.inputs.Textboxbox(),
outputs=gr.outputs.Label(num_top_classes=5),
live=True
)
# Launch the Gradio interface
iface.launch()