better visualization

app.py

@@ -1,6 +1,8 @@
 import gradio as gr
 from transformers import AutoTokenizer
 import torch
+import pandas as pd
+import matplotlib.pyplot as plt
 from model import EnergySmellsDetector
 from config import SMELLS, BEST_THRESHOLD
 
@@ -8,29 +10,49 @@ TOKENIZER = "microsoft/graphcodebert-base"
 tokenizer = AutoTokenizer.from_pretrained(TOKENIZER)
 model = EnergySmellsDetector.load_model_from_hf()
 
-
-def round_logit(logits, threshold):
-    logits = (logits > threshold).to(int)
-    return logits.cpu().numpy()
-
-
-def greet(code_snippet):
+def get_predictions(code_snippet):
     inputs = tokenizer(code_snippet, return_tensors="pt", truncation=True)
     with torch.no_grad():
         logits = model(**inputs)[0]
-        rounded_logits = round_logit(logits, BEST_THRESHOLD)
-        return f"{dict(zip(SMELLS, map(int, rounded_logits)))}"
+        probs = torch.sigmoid(logits).cpu().numpy().flatten()
+        rounded_logits = (probs > BEST_THRESHOLD).astype(int)
+
+    # Prepare results in a dictionary
+    results = {label: {"Detected": bool(pred), "Confidence": round(prob * 100, 2)}
+               for label, pred, prob in zip(SMELLS, rounded_logits, probs)}
+
+    return results, plot_bar_chart(results)
+
+
+def plot_bar_chart(results):
+    labels = list(results.keys())
+    confidences = [results[label]["Confidence"] for label in labels]
 
+    plt.figure(figsize=(8, 4))
+    plt.barh(labels, confidences, color=['green' if results[label]["Detected"] else 'red' for label in labels])
+    plt.xlabel("Confidence (%)")
+    plt.xlim(0, 100)
+    plt.title("Energy Smells Detection Confidence")
+    plt.gca().invert_yaxis()  # Invert y-axis for better readability
+    plt.tight_layout()
+    img_path = "confidence_chart.png"
+    plt.savefig(img_path)
+    plt.close()
+    
+    return img_path
 
 
-textbox = gr.Textbox(label="Enter your code snippet", placeholder="Here goes your code")
-description = "An application to identify whether your code has energy smells or not. It predicts the presence of 9 different energy smells."
+textbox = gr.Textbox(label="Enter your code snippet", placeholder="Paste your code here...")
 title = "Energy Smells Detector"
+description = "Analyze your code for potential energy smells. The model detects 9 different energy inefficiencies in your code."
 
 gr.Interface(
     title=title,
     description=description,
     inputs=textbox,
-    fn=greet,
-    outputs="text"
+    fn=get_predictions,
+    outputs=[
+        gr.Json(label="Detection Results"),
+        gr.Image(label="Confidence Bar Chart")
+    ]
 ).launch()