import threading

import gradio as gr
import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

BASE_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
ADAPTER = "ali4566544/qwen2.5-0.5b-instruct-python-lora"
# The adapter was trained with this exact system prompt.
SYSTEM_PROMPT = (
    "You are a helpful Python coding assistant. Write clear, correct, idiomatic Python, "
    "explain the code concisely, and point out bugs and how to fix them when asked."
)

tokenizer = AutoTokenizer.from_pretrained(ADAPTER)
base = AutoModelForCausalLM.from_pretrained(BASE_MODEL, torch_dtype=torch.float32)
# Merge LoRA weights into the base model for faster CPU inference.
model = PeftModel.from_pretrained(base, ADAPTER).merge_and_unload()
model.eval()


def respond(message, history, max_new_tokens, temperature):
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    messages += [{"role": m["role"], "content": m["content"]} for m in history]
    messages.append({"role": "user", "content": message})

    inputs = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt", return_dict=True
    )
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    gen_kwargs = dict(**inputs, streamer=streamer, max_new_tokens=int(max_new_tokens))
    if temperature > 0:
        gen_kwargs.update(do_sample=True, temperature=temperature, top_p=0.9)
    else:
        gen_kwargs.update(do_sample=False)

    threading.Thread(target=model.generate, kwargs=gen_kwargs).start()
    output = ""
    for chunk in streamer:
        output += chunk
        yield output


demo = gr.ChatInterface(
    respond,
    type="messages",
    title="🐍 PyCoder",
    description=(
        "A small Python coding assistant: Qwen2.5-0.5B-Instruct fine-tuned with a "
        f"[LoRA adapter](https://huggingface.co/{ADAPTER}). It's a 0.5B learning project, "
        "so always run and test the code it writes."
    ),
    additional_inputs=[
        gr.Slider(64, 1024, value=512, step=32, label="Max new tokens"),
        gr.Slider(0.0, 1.5, value=0.0, step=0.1, label="Temperature (0 = greedy)"),
    ],
    examples=[
        ["Write a function that removes duplicates from a list while keeping order."],
        ["How do I read a CSV file with pandas and drop rows with missing values?"],
        ["Write a minimal FastAPI app with a /health endpoint."],
        ["Why does this fail? `for i in range(len(xs)): xs.remove(xs[i])`"],
    ],
    cache_examples=False,
)

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