#!/usr/bin/env python
# coding: utf-8

# ### Keywords to Title Generator
# - https://huggingface.co/EnglishVoice/t5-base-keywords-to-headline?text=diabetic+diet+plan
# - Apache 2.0


import torch
from transformers import T5ForConditionalGeneration,T5Tokenizer

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = T5ForConditionalGeneration.from_pretrained("EnglishVoice/t5-base-keywords-to-headline")
tokenizer = T5Tokenizer.from_pretrained("EnglishVoice/t5-base-keywords-to-headline", clean_up_tokenization_spaces=True, legacy=False)
model = model.to(device)

keywords = "music, sleep, night"

text =  "headline: " + keywords
encoding = tokenizer.encode_plus(text, return_tensors = "pt")
input_ids = encoding["input_ids"].to(device)
attention_masks = encoding["attention_mask"].to(device)
beam_outputs = model.generate(
    input_ids = input_ids,
    attention_mask = attention_masks,
    max_new_tokens = 25,
    do_sample = True,
    num_return_sequences = 5,
    temperature = 1.2,
    #num_beams = 20,
    #num_beam_groups = 20,
    #diversity_penalty=0.8,
    no_repeat_ngram_size = 3,
    penalty_alpha = 0.8,
    #early_stopping = True,
    top_k = 15,
    #top_p = 0.60,
)

for i in range(len(beam_outputs)):
    result = tokenizer.decode(beam_outputs[i], skip_special_tokens=True)
    print(result)


'''
#Create a four button panel for changing parameters with one click

def fn(text):
    return ("Hello gradio!")

with gr.Blocks () as demo:

    with gr.Row(variant='compact') as PanelRow1: #first row: top
        
        with gr.Column(scale=0, min_width=180) as PanelCol5:
            gr.HTML("")
        with gr.Column(scale=0) as PanelCol4:
            submit = gr.Button("Temp++", scale=0)
        with gr.Column(scale=1) as PanelCol5:
            gr.HTML("")
        
       
    with gr.Row(variant='compact') as PanelRow2: #2nd row: left, right, middle

        with gr.Column(min_width=100) as PanelCol1:
            submit = gr.Button("Contrastive")
        with gr.Column(min_width=100) as PanelCol2:
            submit = gr.Button("Re-generate")
        with gr.Column(min_width=100) as PanelCol3:
            submit = gr.Button("Diversity Beam")
        
        with gr.Column(min_width=100) as PanelCol5:
            gr.HTML("")
        with gr.Column(min_width=100) as PanelCol5:
            gr.HTML("")
        with gr.Column(scale=0) as PanelCol5:
            gr.HTML("")
            
    with gr.Row(variant='compact') as PanelRow3: #last row: down
        with gr.Column(scale=0, min_width=180) as PanelCol7:
            gr.HTML("")
        with gr.Column(scale=1) as PanelCol6:
            submit = gr.Button("Temp--", scale=0)
            
        with gr.Column(scale=0) as PanelCol5:
            gr.HTML("")

demo.launch()
'''