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()).detach().numpy()
    values, indices = torch.topk(probs, k=10)
    # 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)
    print("values: ", values)
    print("indices: ", indices)
    print({i: v.item() for i, v in zip(indices, values)})
    return {cpv[i]: v.item() for i, v in zip(indices, values)}

gr.Interface(fn=askcpv, inputs="textbox", outputs="label").launch()