import gradio as gr
import pandas as pd
import torch
import numpy as np
from transformers import AutoTokenizer, AutoModelForSequenceClassification

tokenizer = AutoTokenizer.from_pretrained("model")
model = AutoModelForSequenceClassification.from_pretrained("model")

cpv = pd.read_csv("cpv.csv", sep=";")
df = pd.read_csv("code-desc.csv", sep=";")

labels = df.columns[1:]
cpv = cpv.code

id2label = {idx:label for idx, label in enumerate(labels)}
label2id = {label:idx for idx, label in enumerate(labels)}
codeDesc = dict(zip(df.code, df.desc))
   
def askcpv(description):
    
    encoding = tokenizer(description, return_tensors="pt")
    encoding = {k: v.to(model.device) for k,v in encoding.items()}
    outputs = model(**encoding)
    sigmoid = torch.nn.Sigmoid()
    probs = sigmoid(outputs.logits.squeeze().cpu())
    values, indices = torch.topk(probs, k=5)
    # turn predicted id's into actual label names
    # predicted_labels = [id2label[idx] for idx, label in enumerate(predictions) if label == 1.0]
    # return predicted_labels
    print("probs: ", probs.detach().numpy())
    print("values: ", values.detach().numpy())
    print("indices: ", indices.detach().numpy())
    # print({i: v.item() for i, v in zip(indices, values)})
    return {cpv[i]: v.item() for i, v in zip(indices.detach().numpy(), values.detach().numpy())}

gr.Interface(fn=askcpv, inputs="textbox", outputs="label", description = "Devuelve los dos primeros dígitos más probables dada la descripción de una licitación en español.",
title = "Clasificador de códigos CPV",  examples=[["Nuevo edificio de acceso en piscinas de Binaced."], ["Servicio de ambulancias y personal sanitario para el servicio de salvamento y socorrismo de las playas de Cádiz"]], css=".text-2xl{font-size:1.2rem !important; line-height: 0.5rem;} .text-sm{font-size:0.7rem !important} .mb-2{margin-bottom: 0.1rem};" ).launch()