import gradio as gr
from transformers import AutoTokenizer
import torch
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 round_logit(logits, threshold):
    logits = (logits > threshold).to(int)
    return logits.cpu().numpy()


def greet(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)))}"



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."
title = "Energy Smells Detector"

gr.Interface(
    title=title,
    description=description,
    inputs=textbox,
    fn=greet,
    outputs="text"
).launch()