#!/usr/bin/env python3 import json, sys, statistics if len(sys.argv) < 5: print(json.dumps({"ok": False, "stage": "usage", "error": "usage: validate-dataset.py FILE MODEL MAX_EXAMPLES MAX_LENGTH"})) raise SystemExit(0) file_path, model_name, max_examples_s, max_length_s = sys.argv[1:5] max_examples = int(max_examples_s) max_length = int(max_length_s) try: from tinker_cookbook import model_info from tinker_cookbook.renderers import TrainOnWhat, get_renderer from tinker_cookbook.tokenizer_utils import get_tokenizer except Exception as e: print(json.dumps({"ok": False, "stage": "import", "error": f"{type(e).__name__}: {e}"})) raise SystemExit(0) try: renderer_name = model_info.get_recommended_renderer_name(model_name) tokenizer = get_tokenizer(model_name) renderer = get_renderer(renderer_name, tokenizer) except Exception as e: print(json.dumps({"ok": False, "stage": "renderer", "error": f"{type(e).__name__}: {e}"})) raise SystemExit(0) BUCKETS = [512, 2048, 8192, 16384, 32768, 65536] def flatten_values(x): if hasattr(x, "detach"): x = x.detach().cpu() if hasattr(x, "tolist"): x = x.tolist() if isinstance(x, (list, tuple)): out = [] for item in x: out.extend(flatten_values(item)) return out try: return [float(x)] except Exception: return [] def content_text(content): if isinstance(content, str): return content if isinstance(content, list): parts = [] for part in content: if isinstance(part, dict): if isinstance(part.get("text"), str): parts.append(part["text"]) elif isinstance(part.get("content"), str): parts.append(part["content"]) elif isinstance(part, str): parts.append(part) return "\n".join(parts) return "" def stats(xs): if not xs: return {"count": 0} xs_sorted = sorted(xs) return { "count": len(xs), "min": xs_sorted[0], "p25": xs_sorted[max(0, int((len(xs_sorted)-1)*0.25))], "median": statistics.median(xs_sorted), "p90": xs_sorted[max(0, int((len(xs_sorted)-1)*0.90))], "p95": xs_sorted[max(0, int((len(xs_sorted)-1)*0.95))], "max": xs_sorted[-1], "mean": sum(xs_sorted) / len(xs_sorted), } def histogram(xs): counts = {f"<= {b}": 0 for b in BUCKETS} counts[f"> {BUCKETS[-1]}"] = 0 for x in xs: placed = False for b in BUCKETS: if x <= b: counts[f"<= {b}"] += 1 placed = True break if not placed: counts[f"> {BUCKETS[-1]}"] += 1 return counts def contiguous_ranges(indices): if not indices: return [] ranges = [] start = prev = indices[0] for idx in indices[1:]: if idx == prev + 1: prev = idx else: ranges.append((start, prev + 1)) start = prev = idx ranges.append((start, prev + 1)) return ranges rows = 0 valid_rows = 0 roles = {} bad = [] no_user = [] no_assistant = [] empty_assistant = [] assistant_char_lengths = [] checked = 0 token_lengths = [] trainable_tokens = [] trainable_ratios = [] zero_trainable = [] over_max_length = [] top_longest = [] previews = [] trainable_previews = [] length_mismatch = [] try: with open(file_path) as f: for line_number, line in enumerate(f, start=1): line = line.strip() if not line: continue rows += 1 try: obj = json.loads(line) except Exception as e: bad.append({"line": line_number, "issue": f"invalid JSON: {type(e).__name__}: {e}"}) continue messages = obj.get("messages") if not isinstance(messages, list): bad.append({"line": line_number, "issue": "missing messages[]"}) continue valid_rows += 1 seen_user = False seen_assistant = False assistant_chars_this_row = 0 for j, msg in enumerate(messages): if not isinstance(msg, dict): bad.append({"line": line_number, "issue": f"message {j} is not an object"}) continue role = msg.get("role", "") roles[role] = roles.get(role, 0) + 1 if role not in {"system", "user", "assistant", "tool"}: bad.append({"line": line_number, "issue": f"message {j} has unexpected role {role!r}"}) if role == "user": seen_user = True if role == "assistant": seen_assistant = True txt = content_text(msg.get("content")) assistant_chars_this_row += len(txt.strip()) if not txt.strip() and not msg.get("tool_calls"): empty_assistant.append(line_number) content = msg.get("content") if not isinstance(content, (str, list)): bad.append({"line": line_number, "issue": f"message {j} content should be string or content-part array"}) if not seen_user: no_user.append(line_number) if not seen_assistant: no_assistant.append(line_number) assistant_char_lengths.append(assistant_chars_this_row) if checked >= max_examples: continue try: model_input, weights = renderer.build_supervised_example( messages, train_on_what=TrainOnWhat.ALL_ASSISTANT_MESSAGES ) toks = model_input.to_ints() if hasattr(model_input, "to_ints") else [] w = flatten_values(weights) positive = [i for i, weight in enumerate(w) if weight > 0] checked += 1 token_lengths.append(len(toks)) trainable_tokens.append(len(positive)) trainable_ratios.append((len(positive) / len(toks)) if toks else 0.0) if len(w) not in {len(toks), max(0, len(toks) - 1)}: length_mismatch.append({"line": line_number, "tokens": len(toks), "weights": len(w)}) if len(positive) == 0: zero_trainable.append(line_number) if len(toks) > max_length: over_max_length.append({"line": line_number, "tokens": len(toks)}) top_longest.append({"line": line_number, "tokens": len(toks), "trainable": len(positive)}) top_longest = sorted(top_longest, key=lambda x: x["tokens"], reverse=True)[:8] if len(previews) < 3: previews.append({"line": line_number, "text": tokenizer.decode(toks[: min(len(toks), 500)])}) if len(trainable_previews) < 5 and positive: snippets = [] for start, end in contiguous_ranges(positive)[:4]: snippet_tokens = toks[start:min(end, start + 160)] snippets.append({"range": [start, end], "text": tokenizer.decode(snippet_tokens)}) trainable_previews.append({"line": line_number, "snippets": snippets}) except Exception as e: bad.append({"line": line_number, "issue": f"renderer/tokenization failed: {type(e).__name__}: {e}"}) except Exception as e: print(json.dumps({"ok": False, "stage": "file", "error": f"{type(e).__name__}: {e}"})) raise SystemExit(0) serious = [] if bad: serious.append(f"{len(bad)} data/renderer issue(s)") if no_assistant: serious.append(f"{len(no_assistant)} row(s) with no assistant message") if empty_assistant: serious.append(f"{len(empty_assistant)} empty assistant message(s)") if zero_trainable: serious.append(f"{len(zero_trainable)} checked row(s) with zero trainable tokens") if length_mismatch: serious.append(f"{len(length_mismatch)} checked row(s) with token/weight length mismatch") warnings = [] if rows < 20: warnings.append("very small dataset; good for smoke tests only") if over_max_length: warnings.append(f"{len(over_max_length)} checked row(s) exceed max_length={max_length}") if checked < min(max_examples, valid_rows): warnings.append("not all rows were tokenized; increase --examples for a fuller audit") if token_lengths and max(token_lengths) > max_length * 0.9: warnings.append("some examples are close to the max length; inspect truncation behavior") if assistant_char_lengths and statistics.median(assistant_char_lengths) < 20: warnings.append("median assistant completion is very short; confirm this is intentional") if serious: readiness = "FIX DATA FIRST" elif warnings: readiness = "SMOKE ONLY" else: readiness = "READY" mean_tokens = (sum(token_lengths) / len(token_lengths)) if token_lengths else 0 mean_trainable = (sum(trainable_tokens) / len(trainable_tokens)) if trainable_tokens else 0 estimate = { "estimated_total_input_tokens_per_epoch": int(mean_tokens * valid_rows) if checked else 0, "estimated_trainable_tokens_per_epoch": int(mean_trainable * valid_rows) if checked else 0, "basis": f"estimated from {checked} checked row(s)", } print(json.dumps({ "ok": not serious, "stage": "done", "readiness": readiness, "serious": serious, "warnings": warnings, "model": model_name, "renderer": renderer_name, "rows": rows, "valid_rows": valid_rows, "checked_examples": checked, "roles": roles, "token_lengths": stats(token_lengths), "trainable_tokens": stats(trainable_tokens), "trainable_ratios": stats(trainable_ratios), "assistant_chars": stats(assistant_char_lengths), "token_length_histogram": histogram(token_lengths), "trainable_token_histogram": histogram(trainable_tokens), "estimate": estimate, "bad": bad[:80], "bad_count": len(bad), "no_user": no_user[:20], "no_user_count": len(no_user), "no_assistant": no_assistant[:20], "no_assistant_count": len(no_assistant), "empty_assistant": empty_assistant[:20], "empty_assistant_count": len(empty_assistant), "zero_trainable": zero_trainable[:20], "zero_trainable_count": len(zero_trainable), "over_max_length": over_max_length[:20], "over_max_length_count": len(over_max_length), "length_mismatch": length_mismatch[:20], "length_mismatch_count": len(length_mismatch), "top_longest": top_longest, "previews": previews, "trainable_previews": trainable_previews, }, ensure_ascii=False))