import gradio as gr
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
import torch
import torch.distributed.run as distributed_run
from git import Repo
from huggingface_hub import HfApi
import multiprocessing as mp


# Clone the medusa repo locally
print("Cloning the medusa repo locally...")
Repo.clone_from("https://github.com/FasterDecoding/Medusa.git", "medusa")
print("Cloning the Vicuna data locally...")
Repo.clone_from("https://huggingface.co/datasets/Aeala/ShareGPT_Vicuna_unfiltered", "data")
print("Done")


def create_medusa_heads(model_id: str):
    parser = distributed_run.get_args_parser()
    args = parser.parse_args([
        "--nproc_per_node", "2",
        "medusa/medusa/train/train.py",
        "--model_name_or_path", model_id,
        "--data_path", "data/ShareGPT_V4.3_unfiltered_cleaned_split.json",
        "--bf16", "True",
        "--output_dir", "medusa_heads",
        "--num_train_epochs", "1",
        "--per_device_train_batch_size", "4",
        "--per_device_eval_batch_size", "4",
        "--gradient_accumulation_steps", "4",
        "--evaluation_strategy", "no",
        "--save_strategy", "no",
        "--learning_rate", "1e-3",
        "--weight_decay", "0.0",
        "--warmup_ratio", "0.1",
        "--lr_scheduler_type", "cosine",
        "--logging_steps", "1",
        "--tf32", "True",
        "--model_max_length", "2048",
        "--lazy_preprocess", "True",
        "--medusa_num_heads", "3",
        "--medusa_num_layers", "1",
    ])
    distributed_run.run(args)

    # Upload the medusa heads to the Hub
    repo_id = f"{model_id}-medusa"
    api = HfApi()
    api.create_repo(
        repo_id=repo_id,
        exist_ok=True,
    )
    api.upload_folder(
        folder_path="medusa_heads",
        repo_id=repo_id,
    )
    return repo_id

def run(model_id: str) -> str:
    print(f"\n\n\nNEW RUN: {model_id}")

    # Input validation
    if model_id == "":
        return """
        ### Invalid input 🐞

        Please fill a model_id.
        """
    print(f"Valid inputs ✅\nValidating model_id: {model_id}")

    # Attempt to load the base model
    try:
        config = AutoConfig.from_pretrained(model_id)
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16)
        del config, tokenizer, model
    except Exception as e:
        return f"""
        ### {model_id} can't be loaded with AutoClasses 🐞

        {e}
        """
    print(f"{model_id} can be loaded ✅\nCreating medusa heads (will take a few hours)")

    # Run the medusa heads creation
    try:
        proc = mp.Process(target=create_medusa_heads, args=(model_id, ))
        proc.start()
        proc.join()
        repo_id = f"{model_id}/medusa"
        print("Success ✅\nMedusa heads uploaded to: ", repo_id)
        return f"""
        ### Success 🔥

        Yay! Medusa heads were successfully created and uploaded to, {repo_id}
        """
    except Exception as e:
        print("Error ❌\n", e)
        return f"""
        ### Error 😢😢😢

        {e}
        """


DESCRIPTION = """
The steps to create [medusa](https://sites.google.com/view/medusa-llm) heads are the following:

1. Input a public model id from the Hub
2. Click "Submit"
3. That's it! You'll get feedback if it works or not, and if it worked, you'll get the URL of the new repo 🔥
"""

title="Create LLM medusa heads in a new repo 🐍"

with gr.Blocks(title=title) as demo:
    description = gr.Markdown(f"""# {title}""")
    description = gr.Markdown(DESCRIPTION)

    with gr.Row() as r:
        with gr.Column() as c:
            model_id = gr.Text(max_lines=1, label="model_id")
            with gr.Row() as c:
                clean = gr.ClearButton()
                submit = gr.Button("Submit", variant="primary")

        with gr.Column() as d:
            status_box = gr.Markdown()

    submit.click(run, inputs=[model_id], outputs=status_box, concurrency_limit=1)

demo.queue(max_size=10).launch(show_api=True)