from transformers import AutoTokenizer
import transformers
import os
import sys
import fire
import torch
import gradio as gr


MAGICODER_PROMPT = """You are an exceptionally intelligent coding assistant that consistently delivers accurate and reliable responses to user instructions.

@@ Instruction
{instruction}

@@ Response
"""

title = "<h1 style='text-align: center; margin-bottom: 1rem'>🎩 Magicoder</h1>"

description = """This is a playground for Magicoder-S-DS-6.7B! Follow us on Github: https://github.com/ise-uiuc/magicoder and Huggingface: https://huggingface.co/ise-uiuc.

💡 Tips: You can include more details in your instruction for Magicoder to write better code for you! Details can be, but not limited to, as follows:
1. Specify your programming language (e.g., "... in Python" and "... in Java")
2. Specify the external libraries/packages that are necessary in your task (e.g., "... using the turtle library" and "Write a gradio application to ...")
3. Demonstrate your requirements as clear and comprehensive as possible (e.g., "use CNN as the model structure" and "draw a chart of the training loss")"""


def main(
    base_model="ise-uiuc/Magicoder-S-DS-6.7B"
):
    pipeline = transformers.pipeline(
        "text-generation", model=base_model, torch_dtype=torch.bfloat16, device_map="auto"
    )

    def evaluate_magicoder(
        instruction,
        temperature=1,
        max_length=2048,
    ):
        prompt = MAGICODER_PROMPT.format(instruction=instruction)

        if temperature > 0:
            sequences = pipeline(
                prompt,
                do_sample=True,
                temperature=temperature,
                max_length=max_length,
            )
        else:
            sequences = pipeline(
                prompt,
                max_length=max_length,
            )
        for seq in sequences:
            generated_text = seq["generated_text"].replace(prompt, "")
            return generated_text

    with gr.Blocks() as demo:
        gr.Markdown(title)
        gr.Markdown(description)
        with gr.Row(equal_height=False):
            with gr.Column(variant="panel"):
                interface_input = gr.Textbox(
                    lines=3,
                    label="Instruction",
                    placeholder="Anything you want to ask Magicoder ?",
                    value="Write a snake game in Python using the turtle library (the game is created by Magicoder).",
                )
                temperature = gr.Slider(minimum=0, maximum=1, value=0, label="Temperature")
                max_new_tokens = gr.Slider(
                    minimum=1, maximum=2048, step=1, value=2048, label="Max tokens"
                )
                submit = gr.Button(value="Generate")
                gr.Examples(
                    examples=[
                        ["Write a snake game in Python using the turtle library (the game is created by Magicoder)."],
                        ["Build a console-based Othello game in Java with row and column numbers shown on the board. The game should end when there are no more valid moves for either player."],
                        ["Write a gradio (3.48.0) application for the following use case: Take an input image and return a 45 degree clockwise rotated image. You should also add text description under the output showing the rotation degree."],
                        ["Build a simple neural network in Python using Pytorch to classify handwritten digits from the MNIST dataset. You should use CNN as the model structure, train the model for 5 epochs, draw a chart of the training loss, and show the final result."],
                    ],
                    inputs = [interface_input]
                )
            output = gr.Markdown(
                label="Output",
            )
            submit.click(
                evaluate_magicoder,
                inputs=[interface_input, temperature, max_new_tokens],
                outputs=[output],
            )
    demo.queue().launch()

if __name__ == "__main__":
    fire.Fire(main)