from transformers import pipeline
import gradio as gr

model_id = 'okite97/distilbert-base-uncased-najianews'
pipe = pipeline('text-classification',model=model_id)

sample_title1 = 'Uefa Opens Proceedings against Barcelona, Juventus and Real Madrid Over European Super League Plan'
sample_desc1 = 'Uefa has opened disciplinary proceedings against Barcelona, Juventus and Real Madrid over their involvement in the proposed European Super League.'
sample_title2 = 'Nigeria’s Parliament Passes Amended Electoral Bill'
sample_desc2 = 'Nigeria\'s Senate on Tuesday passed the harmonised Clause 84 of the 2010 Electoral Act (Amendment) Bill 2022, which allows political'

def predict_news(title, description):
  news = title + '. ' + description
  outputs = pipe(news, top_k=6)
  labels = {}
  for val in outputs:
    labels[val['label']] = val['score']
  return labels
  
title = 'News Categorization'
description = '''Given a news title and an excerpt from the news article, can we place the news in its proper category?'''
article = "To learn more, check out [Github repo](https://github.com/chimaobi-okite/NLP-Projects-Competitions/tree/main/NewsCategorization) that this demo is based off of."
examples = [[sample_title1, sample_desc1], [sample_title2, sample_desc2]]

    
textbox1  = gr.Textbox(label="News Title" , lines=1)
textbox2  = gr.Textbox(label="Excerpt", lines=2)
label = gr.Label()
demo = gr.Interface(fn=predict_news, inputs = [textbox1, textbox2], 
                     outputs = label,
                    title=title,
                    description=description,
                    article = article, examples = examples)
demo.launch()