scripts/build_sft_dataset.py

#!/usr/bin/env python3
"""Build claim-drafter dataset v3: all six industry domains, multi-slice validation.

Run with the tokenizer venv so token counts are exact:
    tokvenv/bin/python build_dataset_v3.py <hupd_dir> <tokenizer.json> [out_dir]

Changes over v2:
  * All six target industry domains (see domains.py), not just software/mechanical.
  * Drops claims that reference material absent from the text -- drawn chemical
    structures ("a compound of formula (I)"), sequence listings ("SEQ ID NO: 1"),
    Markush R-group tables, and figures. These are unlearnable as text: the model
    would emit a placeholder pointing at an image that does not exist.
  * Exact token counts from the real Qwen tokenizer, with a hard cap. Tinker's
    cookbook truncates over-length examples silently FROM THE RIGHT, which would
    eat the tail of the claim set without raising -- so length is enforced here.
  * Caps any single domain's share so electronics does not swamp the mix.
  * Emits a held-out eval pool sliced several ways for multi-dataset validation.

Retained from v2: granted (not as-filed) claims, prompts built from Background +
Abstract with the Summary dropped, abstract-leakage filter.
"""

import argparse
import glob
import hashlib
import json
import os
import random
import re
import sys
from collections import Counter, defaultdict

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from claim_drafter.domains import DOMAINS, classify
from claim_drafter.filters import (claims_numbered_correctly, clean_section, find_unlearnable,
                                   leak_ratio, truncate_at_sentence)
from claim_drafter.patents import normalize_claim

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
CACHE = os.path.join(ROOT, "data", "cache", "granted_claims.jsonl")

SYSTEM_PROMPT = (
    "You are an expert US patent attorney. Draft a set of independent and "
    "dependent claims based on the provided invention description. Use "
    "correct USPTO formatting."
)

SEED = 42
EVAL_POOL_FRACTION = 0.12
MAX_DOMAIN_SHARE = 0.30

MIN_BACKGROUND_CHARS, MAX_BACKGROUND_CHARS = 400, 4500
MIN_ABSTRACT_CHARS, MAX_ABSTRACT_CHARS = 150, 2500
MIN_CLAIMS, MAX_CLAIMS = 3, 30
MIN_CLAIMS_CHARS = 500
MAX_TOKENS = 8192          # aligns with the recommended Tinker max_length
MAX_ABSTRACT_LEAK = 0.50








def load_granted():
    granted = {}
    with open(CACHE) as f:
        for line in f:
            try:
                r = json.loads(line)
            except ValueError:
                continue
            claims = [normalize_claim(c) for c in r.get("claims", [])]
            claims = [c for c in claims if c]
            if claims:
                granted[r["patent_number"]] = claims
    return granted


def build_example(d, granted, tok, stats):
    if d.get("decision") != "ACCEPTED":
        stats["not_granted"] += 1
        return None
    domain = classify(d.get("main_ipcr_label"))
    if not domain:
        stats["unmapped_ipc"] += 1
        return None

    pn = d.get("patent_number") or ""
    claims = granted.get(pn)
    if not claims:
        stats["no_granted_claims"] += 1
        return None
    if not (MIN_CLAIMS <= len(claims) <= MAX_CLAIMS):
        stats["claim_count"] += 1
        return None
    if not claims_numbered_correctly(claims):
        stats["claim_numbering"] += 1
        return None

    claims_text = "\n\n".join(claims)
    if len(claims_text) < MIN_CLAIMS_CHARS:
        stats["claim_length"] += 1
        return None
    bad = find_unlearnable(claims_text)
    if bad:
        stats["unlearnable_" + bad] += 1
        return None

    background = clean_section(d.get("background"))
    abstract = clean_section(d.get("abstract"))
    if len(background) < MIN_BACKGROUND_CHARS:
        stats["thin_background"] += 1
        return None
    if len(abstract) < MIN_ABSTRACT_CHARS:
        stats["thin_abstract"] += 1
        return None

    leak = leak_ratio(abstract, claims[0])
    if leak > MAX_ABSTRACT_LEAK:
        stats["abstract_leak"] += 1
        return None

    background = truncate_at_sentence(background, MAX_BACKGROUND_CHARS)
    abstract = truncate_at_sentence(abstract, MAX_ABSTRACT_CHARS)
    title = " ".join((d.get("title") or "").split()).title()
    user_msg = ("Title: %s\n\nTechnical Field and Background:\n%s\n\nInvention Disclosure:\n%s"
                % (title, background, abstract))

    n_prompt = len(tok.encode(SYSTEM_PROMPT + user_msg).ids)
    n_completion = len(tok.encode(claims_text).ids)
    if n_prompt + n_completion > MAX_TOKENS:
        stats["too_long"] += 1
        return None

    return {
        "domain": domain,
        "ipc": d.get("main_ipcr_label") or "",
        "application_number": d.get("application_number", ""),
        "patent_number": pn,
        "issue_year": (d.get("patent_issue_date") or "")[:4],
        "filing_year": (d.get("filing_date") or "")[:4],
        "n_claims": len(claims),
        "abstract_leak": round(leak, 4),
        "n_prompt_tokens": n_prompt,
        "n_completion_tokens": n_completion,
        "dedup_key": hashlib.sha1(abstract.encode("utf8")).hexdigest(),
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_msg},
            {"role": "assistant", "content": claims_text},
        ],
    }


def write_split(path, examples):
    with open(path, "w") as f:
        for ex in examples:
            f.write(json.dumps({"messages": ex["messages"]}) + "\n")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("tokenizer_json")
    ap.add_argument("--sources", nargs="+", required=True,
                    help="one or more directories of HUPD application JSON files")
    ap.add_argument("--out-dir", default=os.path.join(ROOT, "data", "sft"))
    args = ap.parse_args()
    args.out_dir = args.out_dir

    from tokenizers import Tokenizer
    tok = Tokenizer.from_file(args.tokenizer_json)

    granted = load_granted()
    print("Loaded granted claims for %d patents" % len(granted))

    files = []
    for src in args.sources:
        found = sorted(glob.glob(os.path.join(src, "*.json")))
        print("  source %s: %d files" % (src, len(found)))
        files.extend(found)

    examples, seen, stats = [], set(), Counter()
    for fp in files:
        with open(fp) as f:
            try:
                d = json.load(f)
            except ValueError:
                stats["bad_json"] += 1
                continue
        ex = build_example(d, granted, tok, stats)
        if ex is None:
            continue
        if ex["dedup_key"] in seen:
            stats["duplicate"] += 1
            continue
        seen.add(ex["dedup_key"])
        examples.append(ex)

    rng = random.Random(SEED)
    rng.shuffle(examples)

    # Cap any one domain's share so electronics does not dominate.
    cap = int(len(examples) * MAX_DOMAIN_SHARE)
    by_domain = defaultdict(list)
    for ex in examples:
        by_domain[ex["domain"]].append(ex)
    kept, capped_out = [], Counter()
    for dom, exs in by_domain.items():
        kept.extend(exs[:cap])
        if len(exs) > cap:
            capped_out[dom] = len(exs) - cap
    rng.shuffle(kept)

    # Hold out a stratified eval pool: proportional within each domain, so even
    # the thinnest domain gets a usable validation slice.
    kept_by_domain = defaultdict(list)
    for ex in kept:
        kept_by_domain[ex["domain"]].append(ex)
    pool, train = [], []
    for dom, exs in kept_by_domain.items():
        n_hold = min(len(exs) // 2, max(8, int(len(exs) * EVAL_POOL_FRACTION)))
        pool.extend(exs[:n_hold])
        train.extend(exs[n_hold:])
    rng.shuffle(pool)
    rng.shuffle(train)

    os.makedirs(args.out_dir, exist_ok=True)
    val_dir = os.path.join(args.out_dir, "validation")
    os.makedirs(val_dir, exist_ok=True)

    write_split(os.path.join(args.out_dir, "train.jsonl"), train)
    write_split(os.path.join(args.out_dir, "val.jsonl"), pool)

    # --- multi-dataset validation: slice the SAME held-out pool several ways ---
    slices = {}
    for dom in DOMAINS:
        s = [e for e in pool if e["domain"] == dom]
        if s:
            slices["domain_" + dom] = s
    leaks = sorted(e["abstract_leak"] for e in pool)
    hard_cut = leaks[max(0, int(len(leaks) * 0.25))]
    slices["hard_low_leakage"] = [e for e in pool if e["abstract_leak"] <= hard_cut]
    slices["long_prosecution_2018"] = [e for e in pool if e["issue_year"] == "2018"]
    slices["many_claims_15plus"] = [e for e in pool if e["n_claims"] >= 15]
    for fy in sorted(set(e["filing_year"] for e in pool if e["filing_year"])):
        s = [e for e in pool if e["filing_year"] == fy]
        if len(s) >= 20:
            slices["cohort_filed_" + fy] = s
    for name, exs in slices.items():
        write_split(os.path.join(val_dir, name + ".jsonl"), exs)

    with open(os.path.join(args.out_dir, "manifest.jsonl"), "w") as mf:
        for split_name, exs in (("train", train), ("val", pool)):
            for ex in exs:
                rec = {k: ex[k] for k in ("domain", "ipc", "application_number",
                                          "patent_number", "issue_year", "filing_year",
                                          "n_claims", "abstract_leak", "n_prompt_tokens",
                                          "n_completion_tokens")}
                rec["split"] = split_name
                mf.write(json.dumps(rec) + "\n")

    tot_tok = sum(e["n_prompt_tokens"] + e["n_completion_tokens"] for e in train)
    comp_tok = sum(e["n_completion_tokens"] for e in train)
    print("\nScanned corpus -> kept %d examples (train %d / eval pool %d)"
          % (len(kept), len(train), len(pool)))
    print("\nBy domain (train):")
    dc = Counter(e["domain"] for e in train)
    for dom, n in dc.most_common():
        print("  %-26s %5d  (%5.1f%%)" % (dom, n, 100.0 * n / len(train)))
    if capped_out:
        print("  capped out by MAX_DOMAIN_SHARE: %s" % dict(capped_out))
    print("\nValidation slices:")
    for name, exs in sorted(slices.items()):
        print("  %-32s %4d" % (name, len(exs)))
    print("\nTraining tokens: %s total, %s completion (weight=1)"
          % ("{:,}".format(tot_tok), "{:,}".format(comp_tok)))
    print("Dropped: %s" % "  ".join("%s=%d" % kv for kv in sorted(stats.items())))
    print("\nWrote %s" % args.out_dir)


if __name__ == "__main__":
    main()