Spaces:
Running
Running
File size: 1,966 Bytes
d407fa8 754963a d407fa8 754963a d407fa8 974d582 754963a d407fa8 754963a adb5ea8 754963a d407fa8 754963a d407fa8 754963a d407fa8 754963a 63f5189 d407fa8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
import gradio as gr
from transformers import AutoTokenizer
import torch
import matplotlib.pyplot as plt
from model import EnergySmellsDetector
from config import SMELLS, BEST_THRESHOLD
TOKENIZER = "microsoft/graphcodebert-base"
tokenizer = AutoTokenizer.from_pretrained(TOKENIZER)
model = EnergySmellsDetector.load_model_from_hf()
def get_predictions(code_snippet):
inputs = tokenizer(code_snippet, return_tensors="pt", truncation=True)
with torch.no_grad():
probs = model(**inputs)[0].cpu().numpy().flatten() # Model output is already sigmoid applied
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(f"Energy Smells Probabilities - Threshold: {BEST_THRESHOLD}")
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="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=get_predictions,
outputs=[
gr.Json(label="Detection Results"),
gr.Image(label="Confidence Bar Chart")
],
).launch()
|