Spaces:
Running
Running
| import gradio as gr | |
| from transformers import AutoModelForSeq2SeqLM, AutoTokenizer | |
| # Load the model and tokenizer | |
| model_name = "google/flan-t5-large" | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| model = AutoModelForSeq2SeqLM.from_pretrained(model_name) | |
| def concatenate_and_generate(source_text, target_example_texts, reranking, temperature, top_p): | |
| concatenated_text = source_text + " " + target_example_texts | |
| inputs = tokenizer(concatenated_text, return_tensors="pt") | |
| # Generate the output with specified temperature and top_p | |
| output = model.generate( | |
| inputs["input_ids"], | |
| do_sample=True, | |
| temperature=temperature, | |
| top_p=top_p, | |
| max_length=100 | |
| ) | |
| generated_text = tokenizer.decode(output[0], skip_special_tokens=True) | |
| return generated_text | |
| # Preset examples with cached generations | |
| preset_examples = [ | |
| { | |
| "source_text": "Once upon a time in a small village", | |
| "target_example_texts": "In a land far away, there was a kingdom ruled by a wise king. Every day, the people of the kingdom would gather to listen to the king's stories, which were full of wisdom and kindness.", | |
| "reranking": 5, | |
| "temperature": 1.0, | |
| "top_p": 1.0, | |
| "output": "Once upon a time in a small village in a land far away, there was a kingdom ruled by a wise king. Every day, the people of the kingdom would gather to listen to the king's stories, which were full of wisdom and kindness." | |
| }, | |
| { | |
| "source_text": "The quick brown fox", | |
| "target_example_texts": "A nimble, chocolate-colored fox swiftly darted through the emerald forest, weaving between trees with grace and agility.", | |
| "reranking": 5, | |
| "temperature": 0.9, | |
| "top_p": 0.9, | |
| "output": "The quick brown fox, a nimble, chocolate-colored fox, swiftly darted through the emerald forest, weaving between trees with grace and agility." | |
| } | |
| ] | |
| # Define Gradio interface | |
| with gr.Blocks(theme="ParityError/[email protected]") as demo: | |
| gr.Markdown("# TinyStyler Demo") | |
| gr.Markdown("Style transfer the source text into the target style, given some example texts of the target style. You can adjust re-ranking and top_p to your desire to control the quality of style transfer. A higher re-ranking value will generally result in better results, at slower speed.") | |
| source_text = gr.Textbox(lines=3, placeholder="Enter the source text to transform into the target style...", label="Source Text") | |
| target_example_texts = gr.Textbox(lines=5, placeholder="Enter example texts of the target style (one per line)...", label="Example Texts of the Target Style") | |
| reranking = gr.Slider(1, 10, value=5, step=1, label="Re-ranking") | |
| temperature = gr.Slider(0.0, 2.0, value=1.0, step=0.1, label="Temperature") | |
| top_p = gr.Slider(0.0, 1.0, value=1.0, step=0.1, label="Top-P") | |
| output = gr.Markdown(label="Output") | |
| def set_example(example): | |
| return example["source_text"], example["target_example_texts"], example["reranking"], example["temperature"], example["top_p"], example["output"] | |
| example_dropdown = gr.Dropdown(label="Preset Examples", choices=[f"Example {i+1}" for i in range(len(preset_examples))]) | |
| example_button = gr.Button("Load Example") | |
| example_button.click( | |
| lambda example_index: set_example(preset_examples[int(example_index.split()[-1])-1]), | |
| inputs=[example_dropdown], | |
| outputs=[source_text, target_example_texts, reranking, temperature, top_p, output] | |
| ) | |
| btn = gr.Button("Generate") | |
| btn.click(concatenate_and_generate, [source_text, target_example_texts, reranking, temperature, top_p], output) | |
| demo.launch() |