import gradio as gr
import numpy as np
from transformers import pipeline

# pipelines
pipeline_clf = pipeline("text-classification", model = "stinoco/beto-sentiment-analysis-finetuned", return_all_scores = True)
pipeline_pos = pipeline("token-classification", model = "sagorsarker/codeswitch-spaeng-pos-lince")


def predict(text: str):

  '''
  Función que recibe texto como input, devuelve la clasificación de texto para ser recibida por el demo.
  text: texto a clasificar (str)
  '''

  # Text Classification
  classes = pipeline_clf(text)[0]

  # POS
  classes = {element['label']: element['score'] for element in classes}

  labeled_text = {'text': text, 'entities': pipeline_pos(text)}

  return classes, labeled_text

demo = gr.Interface(fn = predict, 
                         inputs = [gr.Textbox(placeholder = "Ingresa el reclamo acá", label = 'Reclamo')],
                         outputs = [gr.outputs.Label(label = 'Categorías'),
                                    gr.Highlightedtext(label = 'Part of Speech')],
                         examples = [
                             ['al ser de region simpre esta con quiebre de stock'],
                             ['que tienen que tener vendedores que conozcan el rubro y que sepan lo que estan vendiendo'],
                             ['un solo vendedor no pude estar encargado de miles de articulos debe especificarse en cerveza'],
                             ['no hay mercaderia']
                         ],
                         title = 'Demo Clasificación NPS'
                         )

demo.launch()