Spaces:
Running
Running
File size: 3,971 Bytes
9d0d4e9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | #!/usr/bin/env python3
"""
Post-process dataset: clean teacher explanations by stripping <think> blocks.
Qwen3 instruct outputs reasoning in <think>...</think> tags.
We want only the text after </think>.
Also generates an improved dataset with augmented prompt templates.
"""
import json, re, random
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parents[2]
ARTIFACTS_DIR = REPO_ROOT / "artifacts" / "tiny_nla"
SIDECAR_PATH = Path(__file__).resolve().parent / "nla_meta.yaml"
JSONL_PATH = ARTIFACTS_DIR / "dataset.jsonl"
ACT_PATH = ARTIFACTS_DIR / "activations.pt"
def clean_explanation(text: str) -> str:
"""Strip <think> blocks, extract only the answer text."""
if not text:
return ""
# Remove <think>...</think> blocks (possibly with newlines)
cleaned = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
cleaned = cleaned.strip()
# Also strip leading/trailing whitespace per line
cleaned = '\n'.join(line.strip() for line in cleaned.split('\n') if line.strip())
# If after stripping think there's still nothing, keep minimal
if not cleaned:
# Try to salvage something from within think
inner = re.search(r'<think>(.*?)</think>', text, re.DOTALL)
if inner:
cleaned = inner.group(1).strip()[:100]
return cleaned
def is_empty_or_useless(text: str) -> bool:
"""Check if explanation is empty or just boilerplate."""
if not text or len(text) < 5:
return True
useless = ["[空输出]", "[生成失败", "<think>"]
if any(u in text for u in useless):
return True
return False
def main():
print("=" * 60)
print("🧹 Post-Processing Dataset")
print("=" * 60)
# Load dataset
records = []
with open(JSONL_PATH, "r", encoding="utf-8") as f:
for line in f:
records.append(json.loads(line))
print(f" Loaded {len(records)} records from {JSONL_PATH}")
# Check stats before
empty_before = sum(1 for r in records if is_empty_or_useless(r.get("teacher_explanation", "")))
has_think = sum(1 for r in records if "<think>" in r.get("teacher_explanation", ""))
print(f" Before cleaning:")
print(f" Empty/useless: {empty_before}")
print(f" Contains <think>: {has_think}")
# Clean explanations
for r in records:
raw = r.get("teacher_explanation", "")
cleaned = clean_explanation(raw)
r["teacher_explanation_raw"] = raw # keep original
r["teacher_explanation"] = cleaned
empty_after = sum(1 for r in records if is_empty_or_useless(r.get("teacher_explanation", "")))
print(f" After cleaning:")
print(f" Empty/useless: {empty_after}")
# Check quality
print("\n Sample cleaned explanations:")
sample = random.Random(42).sample(records, 8)
for r in sample:
print(f" [{r['token_text']!r:10}] {r['teacher_explanation'][:80]}")
# Save cleaned dataset
with open(JSONL_PATH, "w", encoding="utf-8") as f:
for r in records:
f.write(json.dumps(r, ensure_ascii=False) + "\n")
print(f"\n ✅ Saved cleaned dataset to {JSONL_PATH}")
# Also save a separate "av ready" format (for AV training directly)
av_data = []
for r in records:
if not is_empty_or_useless(r.get("teacher_explanation", "")):
av_data.append({
"text": r["text"],
"token_text": r["token_text"],
"pos": r["pos"],
"teacher_explanation": r["teacher_explanation"],
"top_tokens": r["top_tokens"],
})
av_path = ARTIFACTS_DIR / "av_training_data.json"
with open(av_path, "w", encoding="utf-8") as f:
json.dump(av_data, f, ensure_ascii=False, indent=2)
print(f" ✅ Saved AV-ready data ({len(av_data)} samples) to {av_path}")
print("\n✅ Post-processing complete!")
if __name__ == "__main__":
main()
|