|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| """Print FULL replies + extracted blocks for a few code prompts.
|
|
|
| The eval report only stores reply_first_160, which hides why v2 scored 0% ast_rate
|
| while the base model scored 93.3%. This dumps everything so we can tell whether the
|
| model emits broken code or the extractor mis-parses a well-formed reply.
|
| """
|
| import json
|
| import os
|
| import re
|
| import sys
|
|
|
| import torch
|
| from transformers import (AutoModelForCausalLM, AutoTokenizer,
|
| BitsAndBytesConfig)
|
|
|
| REPO = sys.argv[1] if len(sys.argv) > 1 else "Taimwe/securecoder-30b-pro-v2-merged"
|
| MAXNEW = 1024
|
|
|
| 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(REPO, token=token)
|
| model = AutoModelForCausalLM.from_pretrained(
|
| REPO, token=token,
|
| quantization_config=BitsAndBytesConfig(
|
| load_in_4bit=True, bnb_4bit_quant_type="nf4",
|
| bnb_4bit_compute_dtype=torch.bfloat16,
|
| bnb_4bit_use_double_quant=True),
|
| dtype=torch.bfloat16, device_map="auto")
|
|
|
| out = []
|
| 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=MAXNEW, 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)
|
|
|
|
|
| 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
|
| if block:
|
| block = re.split(r"\n```", block, maxsplit=1)[0]
|
| err = None
|
| if block:
|
| try:
|
| compile(block, "<t>", "exec")
|
| except Exception as exc:
|
| err = "%s: %s" % (type(exc).__name__, exc)
|
| out.append({"prompt": p[:80], "n_new_tokens": int(len(new_ids)),
|
| "n_fenced_blocks": len(blocks),
|
| "block_len": len(block) if block else 0,
|
| "compile_error": err, "full_reply": reply})
|
| print("\n" + "=" * 74)
|
| print("PROMPT %d new_tokens=%d fenced_blocks=%d block_len=%s"
|
| % (i, len(new_ids), len(blocks), len(block) if block else 0))
|
| print("COMPILE:", err or "OK")
|
| print("-" * 74)
|
| print(reply[:2500])
|
|
|
| dest = sys.argv[2] if len(sys.argv) > 2 else "/data/dbg.json"
|
| with open(dest, "w", encoding="utf-8") as fh:
|
| json.dump(out, fh, indent=2, default=str)
|
| print("\nWROTE", dest) |
