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
"""

css = """
#q-output {
    max-height: 60vh;
    overflow: auto;
}
"""

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

    # componenents 
    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"
    )
    output = gr.Markdown(label="Output")
    examples = [
        ["Write a snake game in Python using the turtle library (the game is created by Magicoder).",0,2048],
        ["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.",0,2048],
        ["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.",0,2048],
        ["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.",0,2048],
        ]
    # interface
    demo = gr.Interface(evaluate_magicoder,inputs = [interface_input,temperature,max_new_tokens],outputs=[output],examples = examples,cache_examples=True,title=title,description=description,css=css)
    demo.queue().launch()

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