# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "torch",
#     "transformers[torch]",
#     "peft",
#     "huggingface_hub",
#     "bitsandbytes",
#     "accelerate",
# ]
# ///
"""Bisect: base+adapter (no merge) vs merged checkpoint.

If base+adapter emits real newlines, the MERGE is what breaks decoding.
If it also emits literal backslash-n, the ADAPTER (i.e. the training) is at fault.
"""
import os
import re
import sys

import torch
from peft import PeftModel
from transformers import (AutoModelForCausalLM, AutoTokenizer,
                          BitsAndBytesConfig)

BASE = "unsloth/Qwen3-Coder-30B-A3B-Instruct"
ADAPTER = "Taimwe/securecoder-30b-pro-v2"

PROMPTS = [
    "Write a Python function `def is_palindrome(s: str) -> bool:` that checks whether a string reads the same forwards and backwards, ignoring case and non-alphanumeric characters.",
    "Write a Python function `def two_sum(nums: list[int], target: int) -> list[int]:` returning the indices of the two numbers adding to target, or an empty list.",
    "Write a Python function `def flatten(nested: list) -> list:` that flattens an arbitrarily nested list.",
]

token = os.environ.get("HF_TOKEN")
tok = AutoTokenizer.from_pretrained(BASE, token=token)
bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                         bnb_4bit_compute_dtype=torch.bfloat16,
                         bnb_4bit_use_double_quant=True)
model = AutoModelForCausalLM.from_pretrained(
    BASE, token=token, quantization_config=bnb,
    dtype=torch.bfloat16, device_map="auto")
model = PeftModel.from_pretrained(model, ADAPTER, token=token)
model.eval()

ok = 0
for i, p in enumerate(PROMPTS, 1):
    text = tok.apply_chat_template([{"role": "user", "content": p}],
                                   tokenize=False, add_generation_prompt=True)
    ids = tok(text, return_tensors="pt", add_special_tokens=False).input_ids.to(model.device)
    with torch.no_grad():
        gen = model.generate(ids, max_new_tokens=1024, do_sample=False,
                             eos_token_id=tok.eos_token_id, pad_token_id=tok.eos_token_id)
    new_ids = gen[0, ids.shape[-1]:]
    reply = tok.decode(new_ids, skip_special_tokens=True)

    has_real_nl = "\n" in reply
    has_lit = "\\n" in reply
    blocks = re.findall(r"```(?:python)?\s*\n(.*?)```", reply, re.S)
    block = max(blocks, key=len) if blocks else None
    if block is None:
        s = reply.find("def ")
        block = reply[s:] if s >= 0 else None
    err = None
    if block:
        try:
            compile(block, "<t>", "exec")
        except Exception as exc:
            err = "%s: %s" % (type(exc).__name__, exc)
    if not err:
        ok += 1
    print("\n" + "=" * 72)
    print("PROMPT %d  new_tokens=%d  fenced=%d  real_nl=%s  literal_backslash_n=%s"
          % (i, len(new_ids), len(blocks), has_real_nl, has_lit))
    print("COMPILE:", err or "OK")
    print("-" * 72)
    print(reply[:1200])

print("\n\nSUMMARY base+adapter: %d/%d compiled cleanly" % (ok, len(PROMPTS)))