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 [
    ["Being a Happier and Healthier Person"],
    ["Learn to Use Mindfulness to Affect Well Being"],
    ["Eating and Drinking - Find Healthy Nutrition Habits"],
    ["Drinking - Find Reasons and Cut Back or Quit Entirely"],
    ["Stress is relieved by quieting your mind, getting exercise and time with nature"],
    ["Reprogram Pain Stress Reactions"],
    ["Brain gamification"],
    ["Mental Body Scan"],
    ["Stretch, Calm, Breath"],
    ["Relaxed Seat Breath"],
    ["Walk Feel"],
    ["alleviating stress"],
    ["helping breathing, satisfaction"],
    ["Relieve Stress, Build Support"],
    ["Relaxation Response"],
    ["Deep Breaths"],
    ["Delete Not Helpful Thoughts"],
    ["Strengthen Helpful"],
    ["Sleep Better and Find Joy"],
    ["Yoga Sleep"],
    ["Relieve Pain"],
    ["Build and Boost Mental Strength"],
    ["Spending Time Outdoors"],
    ["Daily Routine Tasks"],
    ["Feel better each day when you awake by"],
    ["Feel better physically by"],
    ["Practicing mindfulness each day"],
    ["Be happier by"],
    ["Meditation can improve health"],
    ["Spending time outdoors"],
    ["Break the cycle of stress and anxiety"],
    ["Feel calm in stressful situations"],
    ["Deal with work pressure"],
    ["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()