import gradio as gr

import os

# Use a pipeline as a high-level helper
from transformers import pipeline
from huggingface_hub import login

HF_TOKEN = os.getenv('HF_TOKEN')
login(token=HF_TOKEN)

generator = pipeline("text-generation", model="artificial-nerds/gpt2-finetune-miio-v0.1")

def generate_prompt(row):
  prompt = f'Ciclo: {row["Ciclo"]}\nCiclo_Max: {row["Ciclo_Max"]}\nnumero_creditos: {row["numero_creditos"]}\nNroDiasAtraso: {row["NroDiasAtraso"]}\nNroCuotaDias: {row["NroCuotaDias"]}\nDiasCancel: {row["DiasCancel"]}\nDiaRestPago: {row["DiaRestPago"]}\nPctPagoDias_Prom: {row["PctPagoDias_Prom"]}\nPctPagoDias: {row["PctPagoDias"]}\nTipo_Cancelacion: {row["Tipo_Cancelacion"]}\nRango_PagoPlazo: {row["Rango_PagoPlazo"]}\nNroCuotasC: {row["NroCuotasC"]}\nSemanaCancelacion: {row["SemanaCancelacion"]}\ndeposit: {row["deposit"]}\nESTADO_CIVIL: {row["ESTADO_CIVIL"]}\n\nA partir de la información anterior, indica si el registro es Bueno o Malo para un crédito:'
  return prompt

def test(prompt):
    outputs = generator("The White man worked as a", max_length=768, num_return_sequences=1, do_sample=True, top_p=50, top_k=0.95)
    return outputs[0]["generated_text"]

description = """
Test GPT2 Model
"""

default_prompt = """Ciclo: 1
Ciclo_Max: 1
numero_creditos: 1
NroDiasAtraso: 2
NroCuotaDias: 28
DiasCancel: 2
DiaRestPago: -26
PctPagoDias_Prom: 7,143
PctPagoDias: 7,143
Tipo_Cancelacion: No_Liquidado_Vig
Rango_%PagoPlazo: e. No_Liq_vig
NroCuotasC: 4
SemanaCancelacion: 2
deposit: 500
ESTADO_CIVIL: Casado

A partir de la información anterior, indica si el registro es Bueno o Malo para un crédito:"""
iface = gr.Interface(fn=test,
                      title= "GPT2 Model",
                      description = description,
                      inputs=[
                              gr.inputs.Textbox(lines=4, placeholder="Prompt", label='Prompt')
                              ],
                      outputs=[
                               gr.outputs.Textbox(label="Tu poema"),
                              ],
                      examples= [default_prompt]
                     )
iface.launch(enable_queue=True)