TRUFELLINI/python-bugfix

1import gradio as gr2import torch3from transformers import AutoModelForCausalLM, AutoTokenizer4from peft import PeftModel5 6# ── Config ──────────────────────────────────────────────────────────────7BASE_MODEL    = "Qwen/Qwen2.5-Coder-0.5B-Instruct"8ADAPTER_REPO = "TRUFELLINI/qwen2.5-coder-0.5b-bugfix"  9 10SYSTEM_PROMPT = """You are a Python bug fixing assistant.11Rules:12- Make the minimal change possible13- Do not rewrite unrelated code14- Return only fixed code15- Preserve original behavior unless required"""16 17# ── Load model (runs once when Space starts) ─────────────────────────────18print("Loading base model...")19tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)20model = AutoModelForCausalLM.from_pretrained(21    BASE_MODEL,22    torch_dtype=torch.float32,  # CPU — float3223    device_map="cpu",24)25 26print("Loading LoRA adapter...")27model = PeftModel.from_pretrained(model, ADAPTER_REPO)28model.eval()29print("Model ready.")30 31# ── Inference ────────────────────────────────────────────────────────────32def fix_bug(buggy_code, max_new_tokens=400):33    if not buggy_code.strip():34        return "Paste some buggy Python code above."35 36    messages = [37        {"role": "system", "content": SYSTEM_PROMPT},38        {"role": "user",   "content": f"Fix the bug:\n{buggy_code}"}39    ]40    text   = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)41    inputs = tokenizer(text, return_tensors="pt")42 43    with torch.no_grad():44        output = model.generate(45            **inputs,46            max_new_tokens = max_new_tokens,47            temperature    = 0.1,48            do_sample      = True,49            pad_token_id   = tokenizer.eos_token_id,50        )51    new_tokens = output[0][inputs["input_ids"].shape[1]:]52    return tokenizer.decode(new_tokens, skip_special_tokens=True)53 54# ── UI ───────────────────────────────────────────────────────────────────55EXAMPLES = [56    ["def calculate_average(numbers):\n    total = 0\n    for n in numbers:\n        total += n\n    return total / len(numbers)"],57    ["def add_item(item, items=[]):\n    items.append(item)\n    return items"],58    ["def sum_list(nums):\n    total = 0\n    for i in range(len(nums)+1):\n        total += nums[i]\n    return total"],59    ["def first_element(lst):\n    return lst[1]"],60]61 62demo = gr.Interface(63    fn          = fix_bug,64    inputs      = gr.Code(65                    label       = "Paste buggy Python code here",66                    language    = "python",67                    lines       = 14,68                  ),69    outputs     = gr.Code(70                    label       = "Fixed code",71                    language    = "python",72                  ),73    title       = "🐛 Python Bug Fixer",74    description = """**Qwen2.5-Coder 0.5B fine-tuned on real buggy→fixed Python pairs.**  75Makes minimal, correct fixes — not rewrites. Trained using QLoRA on [alexjercan/bugnet](https://huggingface.co/datasets/alexjercan/bugnet).  76[[Model]](https://huggingface.co/your-username/qwen2.5-coder-0.5b-bugfix) · [[GitHub]](https://github.com/your-username/python-bugfix-llm)""",77    examples    = EXAMPLES,78    theme       = gr.themes.Soft(),79    cache_examples = False,80)81 82demo.launch()83