from transformers import pipeline
from transformers import AutoModelForSeq2SeqLM
from transformers import AutoTokenizer


# Load trained model
model = AutoModelForSeq2SeqLM.from_pretrained("/output/reframer")
tokenizer = AutoTokenizer.from_pretrained("/output/reframer")
reframer = pipeline('summarization', model=model, tokenizer=tokenizer)


def reframe(text, strategy):
    text_with_strategy = text +  "Strategy: ['" + strategy + "']"
    return reframer(text_with_strategy)[0]['summary_text']


import gradio as gr
with gr.Blocks() as demo:
    text = gr.Textbox(label="Original Text")
    radio = gr.Radio(
        ["thankfulness", "neutralizing", "optimism", "growth", "impermanence", "self_affirmation"], label="Strategy to use?"
    )
    output = gr.Textbox(label="Reframed Output")
    radio.change(fn=reframe, inputs=[text, radio], outputs=output)
demo.launch()