Update app.py

app.py

@@ -1,49 +1,42 @@
-import requests
-import torch
-from transformers import XLNetTokenizer, XLNetForSequenceClassification
+from transformers import pipeline
 import gradio as gr
 
-# URL of the saved model on GitHub
-model_url = 'https://github.com/AliArshadswl/severity_prediction/raw/main/XLNet_model_project_Core.pt'
+# Load the model using the pipeline
+pipe = pipeline("text-classification", model="AliArshad/Severity_Predictor")
+from transformers import pipeline
+import gradio as gr
 
-# Function to download the model from URL and load it
-def download_model(url):
-    response = requests.get(url)
-    with open('XLNet_model_project_Core.pt', 'wb') as f:
-        f.write(response.content)
+# Load the model using the pipeline
+pipe = pipeline("text-classification", model="AliArshad/Severity_Predictor")
 
-# Download the model
-download_model(model_url)
+# Function to predict severity and return confidence score
+def predict_severity(text):
+    # Get prediction from the pipeline
+    prediction = pipe(text)
 
-# Load the saved model
-tokenizer = XLNetTokenizer.from_pretrained('xlnet-base-cased')
-model = XLNetForSequenceClassification.from_pretrained('xlnet-base-cased', num_labels=2)
-model.load_state_dict(torch.load('XLNet_model_project_Core.pt', map_location=torch.device('cpu')))
-model.eval()
+    # Interpret the label and get the confidence score
+    label = prediction[0]['label']
+    confidence = prediction[0]['score']
+    severity = "Severe" if label == "LABEL_1" else "Non-Severe"
 
-# Function for prediction
-def xl_net_predict(text):
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=100)
-    with torch.no_grad():
-        outputs = model(**inputs)
-    logits = outputs.logits
-    probabilities = torch.softmax(logits, dim=1)
-    predicted_class = torch.argmax(probabilities).item()
-    return "Severe" if predicted_class == 1 else "Non-severe"
+    # Return severity and confidence as separate outputs
+    return severity, confidence
 
-# Customizing the interface
+# Define the Gradio interface with a title, specific placeholder message, and a progress bar for confidence
 iface = gr.Interface(
-    fn=xl_net_predict,
-    inputs=gr.Textbox(lines=2, label="Summary", placeholder="Enter text here..."),
-    outputs=gr.Textbox(label="Predicted Severity"),
-    title="SevPrecit - A GPT-2 Based Bug Report Severity Prediction",
+    fn=predict_severity,
+    inputs=gr.Textbox(lines=2, placeholder="Please Enter Bug Report Summary"),
+    outputs=[
+        gr.Textbox(label="Prediction"),
+        gr.Number(label="Confidence", precision=2)
+    ],
+    title="SevPredict: GPT-2 Based Severity Prediction",
     description="Enter text and predict its severity (Severe or Non-severe).",
-    theme="huggingface",
     examples=[
         ["Can't open multiple bookmarks at once from the bookmarks sidebar using the context menu"],
         ["Minor enhancements to make-source-package.sh"]
-    ],
-    allow_flagging=False
+    ]
 )
 
+# Launch the interface
 iface.launch()