import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# Initialize the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("laptop_data.pkl")  # Replace with your model name or path
model = AutoModelForSequenceClassification.from_pretrained("laptop_data.pkl")  # Replace with your model name or path

# 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()