MMuzamilAI commited on
Commit
d80cee7
·
verified ·
1 Parent(s): 612ae19

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +33 -0
app.py ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoModelForSequenceClassification, AutoTokenizer
3
+ import torch
4
+ import joblib # Assuming label_encoder is saved as a .pkl file
5
+
6
+ # Load model and tokenizer
7
+ model_name = "mmuzamilai/distilbert-review-bug-classifier"
8
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
9
+ model = AutoModelForSequenceClassification.from_pretrained(model_name)
10
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
11
+ model.to(device)
12
+
13
+ # Load label encoder
14
+ label_encoder = joblib.load("label_encoder.pkl") # Adjust if you have another format
15
+
16
+ # Classification function
17
+ def classify_review(text):
18
+ inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=128).to(device)
19
+ with torch.no_grad():
20
+ outputs = model(**inputs)
21
+ predicted_label = torch.argmax(outputs.logits).item()
22
+ decoded_label = label_encoder.inverse_transform([predicted_label])[0]
23
+ return decoded_label
24
+
25
+ # Gradio interface
26
+ iface = gr.Interface(
27
+ fn=classify_review,
28
+ inputs=gr.Textbox(lines=2, placeholder="Enter your review..."),
29
+ outputs=gr.Label(label="Predicted Category"),
30
+ title="Review Bug Classifier"
31
+ )
32
+
33
+ iface.launch()