from transformers import pipeline
import gradio as gr

# loading the model
model_checkpoint = 'zinoubm/e-comerce-category-classification'
model = pipeline(
            "text-classification", model=model_checkpoint,
            )

def predict(input):
  predictions = model(input)
  predictions = [prediction['label'] for prediction in predictions]
  return ' '.join(predictions)

# defining demo content
title = 'E-Commerce Category Prediction.'

description = '''
    This is a classification model that predicts the category of an input.
    We have 4 Categories, Electronics, Household, Books and Clothing & Accessories.
'''

article = '''
# How to use this interface
Using the interface is straight forward, just type some text that falls in one of these 4 categories: **Electronics**, **Household**, **Books** or **Clothing & Accessories**.
and then hit **Submit**. the results will be in the output cell. You can also try one of the provided examples.
Here is the [notebook](https://colab.research.google.com/drive/1MGlDaJXcjECrSRmXgHzUgFBTDuKEwPV0?usp=share_link) used to train the model.
'''

examples = [
    ['I want to sell a laptop'],
    ['This is a beatiful T-shirt'],
    ['Save 50% on detergent powder']
]

# launching the interface
gr.Interface(fn=predict,
             inputs="text",
             title=title,
             description=description,
             article=article,
             outputs="text",
             examples = examples,
             theme='default',
             ).launch()