import gradio as gr
from transformers import pipeline

pipe_flan_t5 = pipeline(
    "text2text-generation", model="pszemraj/flan-t5-small-instructiongen"
)
pipe_bart_base = pipeline(
    "text2text-generation", model="pszemraj/bart-base-instructiongen"
)

examples = [
    "You'll need to start by choosing the right venue. Consider the type of atmosphere and the size of the area that will be suitable for the number of guests you plan to invite. Choose the right decorations based on your brother's interests, such as balloons in his favorite colors, banners, and streamers. Next, decide on the food and drinks, making sure they are tasty and appropriate for the occasion. Then decide on the other games, music, and entertainment that will make the party memorable. Finally, involve your brother's friends and family to help create the perfect surprise.",
    "Start by selecting clothing that is futuristic and edgy, such as leather jackets, neon-colored accessories, and tech-inspired patterns. Add accessories like goggles, cybernetic implants, and LED lights to enhance the cyberpunk vibe. Use makeup and body paint to create a futuristic look, such as metallic skin or neon makeup. Consider adding functional elements to your costume, such as a built-in backpack or hidden pockets for your tech gadgets. Finally, practice your confident walk and embrace your inner cyberpunk for a memorable and immersive costume experience.",
    "Start by creating a base terrain with mountains, valleys, and other natural features. Use fractal noise and displacement mapping to add texture and detail to the terrain, and experiment with different materials like rock, grass, and water. Add surreal elements like floating islands, giant mushrooms, or impossible geometry to create a dreamlike atmosphere. Use lighting and color grading to enhance the mood and tone of the scene, and render the final image at a high resolution for maximum impact. Share your surreal landscape with the world and inspire others to explore the possibilities of 3D art.",
    "Start by setting a realistic goal and creating a training plan. Build up your mileage gradually over time, and incorporate cross-training and strength exercises to prevent injury and improve endurance. Be sure to stay hydrated and properly fuel your body with nutritious foods. Listen to your body and adjust your training as needed to avoid overexertion or burnout. Finally, taper your training in the weeks leading up to the race to give your body time to rest and recover before the big day.",
]

title = "InstructionGen: Instructions from Text"
description = "This demo compares the [flan-t5-small-instructiongen](https://huggingface.co/pszemraj/flan-t5-small-instructiongen) and [bart-base-instructiongen](https://huggingface.co/pszemraj/bart-base-instructiongen) models on 'creating' an instruction for arbitrary text. Note that [the dataset](https://huggingface.co/datasets/pszemraj/fleece2instructions) & models are trained to **only** generate an `instruction`relevant to some text, and **do not** expect/recover the (potential) `inputs`"
article = """---
These models generate instructions **only** for Large Language Models (LLMs) from arbitrary text. They are fine-tuned on the [fleece2instructions](https://huggingface.co/datasets/pszemraj/fleece2instructions) dataset, which is a filtered/formatted version of the [alpaca](https://huggingface.co/datasets/tatsu-lab/alpaca) dataset.

Example of the difference:
- LLM instruction: "What is the capital of France?"
- Instruction+specific inputs:
  Instruction: "Provide information on the following:"
  Specific Inputs: {"category": "geography", "question": "capital of France"}
"""


def inference(text):
    params = {
        "num_beams": 2,
        "early_stopping": True,
        "max_length": 100,
        "truncation": True,
    }

    output_flan_t5 = pipe_flan_t5(text, **params)[0]["generated_text"]
    output_bart_base = pipe_bart_base(text, **params)[0]["generated_text"]

    return [output_flan_t5, output_bart_base]


io = gr.Interface(
    inference,
    gr.Textbox(
        lines=3,
        max_lines=10,
        placeholder="Add text here & find out what LLM intruction/prompt might output that result!",
        interactive=True,
        show_label=False,
    ),
    outputs=[
        gr.Textbox(lines=1, label="Flan T5 Small", interactive=False),
        gr.Textbox(lines=1, label="Bart Base", interactive=False),
    ],
    title=title,
    description=description,
    article=article,
    examples=examples,
    cache_examples=False,
)
io.launch()