import gradio as gr
from transformers import T5Tokenizer, T5ForConditionalGeneration

# xl size run out of memory on 16GB vm
tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large")

title = ""

def get_examples ():
    return [
    ["Write a three sentence outline in markdown code on Being a Happier and Healthier Person"],
    ["Write a three sentence outline in markdown code on Learn to Use Mindfulness to Affect Well Being"],
    ["Write a three sentence outline in markdown code on Find Healthy Nutrition Habits"],
    ["Write a three sentence outline in markdown code on Find Reasons and Cut Back or Quit Entirely"],
    ["Write a three sentence outline in markdown code on Stress is relieved by quieting your mind, getting exercise and time with nature"],
    ["Write a three sentence outline in markdown code on Reprogram Pain Stress Reactions"],
    ["Write a three sentence outline in markdown code on Brain gamification"],
    ["Write a three sentence outline in markdown code on Mental Body Scan"],
    ["Write a three sentence outline in markdown code on Stretch, Calm, Breath"],
    ["Write a three sentence outline in markdown code on Relaxed Seat Breath"],
    ["Write a three sentence outline in markdown code on Walk Feel"],
    ["Write a three sentence outline in markdown code on alleviating stress"],
    ["Write a three sentence outline in markdown code on helping breathing, satisfaction"],
    ["Write a three sentence outline in markdown code on Relieve Stress, Build Support"],
    ["Write a three sentence outline in markdown code on Relaxation Response"],
    ["Write a three sentence outline in markdown code on Deep Breaths"],
    ["Write a three sentence outline in markdown code on Delete Not Helpful Thoughts"],
    ["Write a three sentence outline in markdown code on Strengthen Helpful"],
    ["Write a three sentence outline in markdown code on Sleep Better and Find Joy"],
    ["Write a three sentence outline in markdown code on Yoga Sleep"],
    ["Write a three sentence outline in markdown code on Relieve Pain"],
    ["Write a three sentence outline in markdown code on Build and Boost Mental Strength"],
    ["Write a three sentence outline in markdown code on Spending Time Outdoors"],
    ["Write a three sentence outline in markdown code on Daily Routine Tasks"],
    ["Write a three sentence outline in markdown code on Feel better each day when you awake by"],
    ["Write a three sentence outline in markdown code on Feel better physically by"],
    ["Write a three sentence outline in markdown code on Practicing mindfulness each day"],
    ["Write a three sentence outline in markdown code on Be happier by"],
    ["Write a three sentence outline in markdown code on Meditation can improve health"],
    ["Write a three sentence outline in markdown code on Spending time outdoors"],
    ["Write a three sentence outline in markdown code on Break the cycle of stress and anxiety"],
    ["Write a three sentence outline in markdown code on Feel calm in stressful situations"],
    ["Write a three sentence outline in markdown code on Deal with work pressure"],
    ["Write a three sentence outline in markdown code on Learn to reduce feelings of overwhelmed"]    
    ]
    
    
def text2text(input_text):
    input_ids = tokenizer(input_text, return_tensors="pt").input_ids

    outputs = model.generate(input_ids, max_length=200)
    return tokenizer.decode(outputs[0])


with gr.Blocks() as demo:
    gr.Markdown(
    """
    # Flan T5 Large Demo
    780M parameter Large language model fine tuned on diverse tasks.
    Prompt the model in the Input box.
    """)
    txt_in = gr.Textbox(label="Input", lines=3)
    correct_label = gr.Label(label="Correct")
    txt_out = gr.Textbox(value="", label="Output", lines=4)
    

    btn = gr.Button(value="Submit")
    btn.click(text2text, inputs=[txt_in], outputs=[txt_out])


    gr.Examples(
        examples=get_examples(),
        inputs=[txt_in,correct_label]
    )
    

if __name__ == "__main__":
    demo.launch()