Update app.py

app.py

@@ -1,32 +1,26 @@
 import gradio as gr
-from transformers import BertTokenizer, BertForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 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)
+# 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
 
-# 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)
+# 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 probabilities[0].tolist()
-
+    return {label: prob.item() for label, prob in zip(model.config.id2label.values(), probabilities[0])}
 
+# Create the Gradio interface
 iface = gr.Interface(
-    fn=classify_text,
-    inputs=gr.inputs.Textbox(),
-    outputs=gr.outputs.Label(num_top_classes=2),
-    live=True,
-    interpretation="default"
+    fn=classify_laptop,
+    inputs=gr.inputs.Textboxbox(),
+    outputs=gr.outputs.Label(num_top_classes=5),
+    live=True
 )
 
+# Launch the Gradio interface
 iface.launch()