import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

MODEL_ID = "arinbalyan/qwen-coder-python"

theme = (
    gr.themes.Soft(primary_hue="indigo", neutral_hue="slate")
    .set(button_primary_background_fill_hover="#4f46e5")
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    dtype=torch.float16,
    device_map="auto",
)

def generate(instruction: str):
    if not instruction.strip():
        return "// Describe what you want to code, then hit Generate."
    prompt = f"### Instruction:\n{instruction.strip()}\n\n### Python Code:\n"
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=256,
            temperature=0.7,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
        )
    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return text.split("### Python Code:\n")[-1].strip()

with gr.Blocks(theme=theme, title="CodeLM") as demo:
    gr.Markdown(
        """
        # 🐍 CodeLM
        Fine-tuned code model that writes Python from plain English.
        """
    )
    with gr.Row():
        with gr.Column(scale=1):
            instruction = gr.Textbox(
                label="What should it code?",
                placeholder="e.g., merge two sorted lists",
                lines=3,
            )
            run = gr.Button("Generate", variant="primary", size="lg")
            gr.Examples(
                examples=[
                    "Write a function to reverse a string",
                    "Check if a number is prime",
                    "Flatten a nested list recursively",
                    "Merge two sorted lists",
                    "Compute factorial of a number",
                ],
                inputs=instruction,
                label="Try one of these",
            )
        with gr.Column(scale=1):
            code = gr.Code(
                label="Python",
                language="python",
                lines=20,
            )
    run.click(generate, inputs=instruction, outputs=code)
    gr.Markdown(
        """
        <div style="text-align:center; margin-top:1rem; color:gray; font-size:0.9rem;">
        Base: StarCoder2-1B &nbsp;•&nbsp; Fine-tune: LoRA (r=8) &nbsp;•&nbsp; Trained on Kaggle P100
        </div>
        """
    )

if __name__ == "__main__":
    demo.launch(theme=theme)