import gradio as gr
from transformers import pipeline

pipe = pipeline("summarization", model="RajSang/pegasus-sports-titles")

def sportyify(article,length_penalty=0.6,diversity_penalty=2.0):
  gen_kwargs = {"length_penalty":length_penalty,"num_beams":4, "num_return_sequences": 4,"num_beam_groups":4,"diversity_penalty":float(diversity_penalty)}
  titles=[i["summary_text"]+'\n\n' for i in pipe(article, **gen_kwargs,truncation=True)]
  returnable=''
  for i,title in enumerate(titles):
    returnable+='Title '+str(i+1)+' :\n\n'+title
  return returnable


article = "<p style='text-align: center'><a href='https://huggingface.co/RajSang/pegasus-sports-titles' target='_blank'>Model Page </a></p> "
examples = [["""After Pat Cummins side bowled England out for 188 on day two, Australia struggled once again early on in their second innings, with David Warner, Marnus Labuschagne and Usman Khawaja all falling in the final minutes of day two. But Steve Smith managed to steady the Australia ship and is currently leading Australia to an overall lead past 150, with the Hobart hosts looking to end their exceptional series in style.""",
             0.6,
             2.0]]

iface = gr.Interface(fn=sportyify, 
                     inputs=[gr.inputs.Textbox(lines=15, placeholder="Paste Article Text Here"),
                             gr.inputs.Slider(minimum=0.1,maximum=0.6,default=0.6,step=0.1),
                             gr.inputs.Slider(minimum=1.0,maximum=3.0,default=2.0,step=1.0)
                            ],        
                     outputs="text",
                     title="SPORT-Y Titles",
                     description="<h3 style='text-align: center'>Generate titles for sports articles and sports news.</h3><h6 style='text-align: center'>Created by <a href='https://rajsangani.me' target='_blank'>Raj</a></h6><p>Decreasing the length penalty ranks shorter titles higher.</p><p>Increasing the diversity penalty creates unique titles.</p><p>Read more <a href='https://arxiv.org/pdf/1610.02424.pdf' target='_blank'>here</a></p>",
                     article=article, 
                     examples=examples,
                     allow_flagging='never')

iface.launch(enable_queue=True)