# ### 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
import gradio as gr

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)

def title_gen(keywords, diversity, temp):

 if keywords!= "":
    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)
    
    if diversity:
            num_beams = 20,
            num_beam_groups = 20,
            diversity_penalty=0.8,
            early_stopping = True,
    else:
             penalty_alpha = 0.8,
     
    beam_outputs = model.generate(
        input_ids = input_ids,
        attention_mask = attention_masks,
        max_new_tokens = 30,
        do_sample = True,
        num_return_sequences = 5,
        temperature = temp,
        top_k = 15,
        no_repeat_ngram_size = 3,
        #top_p = 0.60,
    )

    titles = "<p align=center>Title Suggestions:</p>"
    
    for i in range(len(beam_outputs)):
        result = tokenizer.decode(beam_outputs[i], skip_special_tokens=True)
        titles += f"<p align=center><b>{result}</b></p></p>" #Create string with titles and <br> tag for html reading in gradio html
    
    return titles



iface = gr.Interface(fn=title_gen, 
                     inputs=[gr.Textbox(label="Paste one or more keywords searated by a comma and hit 'Submit'.", lines=1), "checkbox", gr.Slider(0.1, 1.9, 1.2)],
                     outputs=[gr.HTML(label="Title suggestions:")],
                     title="AI Keywords to Title Generator", 
                     #description="Turn keywords into creative suggestions",
                     article="<div align=left><h1>AI Creative Title Generator</h1><li>With just keywords, generate a list of creative titles.</li><li>Click on Submit to generate more title options.</li><li>Tweak slider for less or more creative titles</li><li>Check 'diversity' to turn on diversity beam search</li><p>AI Model:<br><li>T5 Model trained on a dataset of titles and related keywords</li><li>Original model id: EnglishVoice/t5-base-keywords-to-headline by English Voice AI Labs</li></p><p>Default parameter details:<br><li><code>temperature = 1.2</code>, <code>no_repeat_ngram_size=3</code>, <code>top_k = 15</code>, <code>penalty_alpha = 0.8</code>, <code>max_new_tokens = 30</code></li><p>Diversity beam search params:<br><li><code>num_beams=20</code>, <code>diversity_penalty=0.8</code>, <code>num_beam_groups=20</code></li></div>",
                     flagging_mode='never',
                     examples=[
                                ["new, weight loss, lifestyle"],
                                ["launch, free, dating, app"],
                                ["AI, text to video, app"],
                                ["new movie, watch, free streaming"],
                              ],
                     cache_examples=True,
                    )

iface.launch()