from liqfit.pipeline import ZeroShotClassificationPipeline
from liqfit.models import T5ForZeroShotClassification
from transformers import T5Tokenizer
import gradio as gr

modelo = T5ForZeroShotClassification.from_pretrained('knowledgator/comprehend_it-multilingual-t5-base')
tokenizador = T5Tokenizer.from_pretrained('knowledgator/comprehend_it-multilingual-t5-base')
classificador = ZeroShotClassificationPipeline(model=modelo, tokenizer=tokenizador,hypothesis_template = '{}', encoder_decoder = True)

categorias = ['ferramentas', 'eletrônicos', 'beleza', 'brinquedos', 'cozinha']

def categorizar(descricao):
  resultado = classificador(descricao, categorias, multi_label=False)
  categoria_max = max(zip(resultado['labels'],resultado['scores']), key=lambda x: x[1])[0]
  return categoria_max

app = gr.Interface(
    fn=categorizar,
    inputs=gr.Textbox(label="Descrição do Produto", placeholder="Digite a descrição do produto"),
    outputs=gr.Textbox(label="Categoria Prevista"),
    title="Classificação de Produtos por Categoria",
    description="Digite a descrição de um produto para classificá-lo automaticamente em uma categoria."
)

if __name__ == "__main__":
    app.launch(share=True)