fix: Add code block normalization script (#750)
This commit is contained in:
141
scripts/normalize-code-blocks.py
Normal file
141
scripts/normalize-code-blocks.py
Normal file
@@ -0,0 +1,141 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
normalize-code-blocks.py — Fix inconsistent indentation in training data code blocks.
|
||||
|
||||
When code blocks are embedded in JSONL as triple-quoted strings, indentation
|
||||
accumulates from the surrounding context. This script normalizes code block
|
||||
content using textwrap.dedent and consistent 4-space indentation.
|
||||
|
||||
Usage:
|
||||
python3 scripts/normalize-code-blocks.py training/data/preference_pairs.jsonl
|
||||
python3 scripts/normalize-code-blocks.py --dry-run training/data/*.jsonl
|
||||
python3 scripts/normalize-code-blocks.py --check training/data/*.jsonl # CI mode
|
||||
"""
|
||||
import argparse
|
||||
import json
|
||||
import re
|
||||
import sys
|
||||
import textwrap
|
||||
from pathlib import Path
|
||||
|
||||
# Matches ```python ... ``` or ``` ... ``` blocks inside string values
|
||||
CODE_BLOCK_RE = re.compile(
|
||||
r"(?P<open>```(?:python|py|bash|sh|javascript|js|typescript|ts|go|rust|ruby)?\s*\n)"
|
||||
r"(?P<code>.*?)"
|
||||
r"(?P<close>```)",
|
||||
re.DOTALL,
|
||||
)
|
||||
|
||||
|
||||
def normalize_code_block(match: re.Match) -> str:
|
||||
"""Normalize indentation in a single code block."""
|
||||
open_tag = match.group("open")
|
||||
code = match.group("code")
|
||||
close_tag = match.group("close")
|
||||
|
||||
# Skip empty blocks
|
||||
if not code.strip():
|
||||
return match.group(0)
|
||||
|
||||
# Dedent the code
|
||||
dedented = textwrap.dedent(code)
|
||||
|
||||
# Strip leading/trailing blank lines
|
||||
lines = dedented.split("\n")
|
||||
while lines and not lines[0].strip():
|
||||
lines.pop(0)
|
||||
while lines and not lines[-1].strip():
|
||||
lines.pop()
|
||||
|
||||
normalized = "\n".join(lines)
|
||||
return f"{open_tag}{normalized}\n{close_tag}"
|
||||
|
||||
|
||||
def process_line(line: str) -> tuple:
|
||||
"""Process a single JSONL line. Returns (new_line, num_fixes)."""
|
||||
try:
|
||||
obj = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
return line, 0
|
||||
|
||||
fixes = 0
|
||||
|
||||
def fix_strings(obj):
|
||||
nonlocal fixes
|
||||
if isinstance(obj, str):
|
||||
original = obj
|
||||
fixed = CODE_BLOCK_RE.sub(normalize_code_block, obj)
|
||||
if fixed != original:
|
||||
fixes += 1
|
||||
return fixed
|
||||
elif isinstance(obj, dict):
|
||||
return {k: fix_strings(v) for k, v in obj.items()}
|
||||
elif isinstance(obj, list):
|
||||
return [fix_strings(item) for item in obj]
|
||||
return obj
|
||||
|
||||
fixed_obj = fix_strings(obj)
|
||||
return json.dumps(fixed_obj, ensure_ascii=False) + "\n", fixes
|
||||
|
||||
|
||||
def process_file(filepath: str, dry_run: bool = False) -> dict:
|
||||
"""Process a single JSONL file. Returns stats dict."""
|
||||
path = Path(filepath)
|
||||
if not path.exists():
|
||||
return {"file": filepath, "error": "not found", "fixes": 0, "lines": 0}
|
||||
|
||||
lines = path.read_text(encoding="utf-8").splitlines()
|
||||
fixed_lines = []
|
||||
total_fixes = 0
|
||||
|
||||
for line in lines:
|
||||
if not line.strip():
|
||||
fixed_lines.append(line)
|
||||
continue
|
||||
new_line, fixes = process_line(line)
|
||||
fixed_lines.append(new_line.rstrip("\n"))
|
||||
total_fixes += fixes
|
||||
|
||||
if total_fixes > 0 and not dry_run:
|
||||
path.write_text("\n".join(fixed_lines) + "\n", encoding="utf-8")
|
||||
|
||||
return {
|
||||
"file": filepath,
|
||||
"lines": len(lines),
|
||||
"fixes": total_fixes,
|
||||
"changed": total_fixes > 0,
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Normalize code block indentation in JSONL training data"
|
||||
)
|
||||
parser.add_argument("files", nargs="+", help="JSONL files to process")
|
||||
parser.add_argument("--dry-run", action="store_true", help="Show changes without writing")
|
||||
parser.add_argument("--check", action="store_true", help="CI mode: exit 1 if fixes needed")
|
||||
args = parser.parse_args()
|
||||
|
||||
total_fixes = 0
|
||||
results = []
|
||||
|
||||
for filepath in args.files:
|
||||
result = process_file(filepath, dry_run=args.dry_run or args.check)
|
||||
results.append(result)
|
||||
total_fixes += result["fixes"]
|
||||
|
||||
if result["fixes"] > 0:
|
||||
status = "FIXED" if not args.dry_run and not args.check else "WOULD FIX"
|
||||
print(f" {status}: {result['file']} — {result['fixes']} code blocks normalized")
|
||||
else:
|
||||
print(f" OK: {result['file']}")
|
||||
|
||||
print(f"\nTotal: {total_fixes} code blocks normalized across {len(results)} files")
|
||||
|
||||
if args.check and total_fixes > 0:
|
||||
print("FAIL: Code block indentation issues found. Run without --check to fix.")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user