import gradio as gr
from transformers import BertTokenizer, BertForSequenceClassification
import torch

# Load the tokenizer and model
model_path = "laptop_data.pkl"  # Replace with the actual path
tokenizer = BertTokenizer.from_pretrained(model_path)
model = BertForSequenceClassification.from_pretrained(model_path)

# Set the model to evaluation mode
model.eval()


def classify_text(text):
    inputs = tokenizer(text, return_tensors="pt",
                       padding=True, truncation=True)
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    probabilities = torch.softmax(logits, dim=1)
    return probabilities[0].tolist()


iface = gr.Interface(
    fn=classify_text,
    inputs=gr.inputs.Textbox(),
    outputs=gr.outputs.Label(num_top_classes=2),
    live=True,
    interpretation="default"
)

iface.launch()