patient-edu-qa / scripts /merge_router_data.py
chenhaodev's picture
Initial upload: patient-edu-qa harness (code, data, LoRA, router GGUF)
784ea73 verified
Raw History Blame Contribute Delete
1.73 kB
#!/usr/bin/env python3
"""Merge curated + augmented router data into a balanced train/val SFT dataset."""
import json
import os
import random
ROOT = "/workspace/TASK18/data/router"
CUR = os.path.join(ROOT, "router_train.jsonl") # 12093 grounded single-intent
AUG = os.path.join(ROOT, "router_augmented.jsonl") # alert + compound
random.seed(42)
def load(p):
return [json.loads(l) for l in open(p, encoding="utf-8")]
cur = load(CUR)
aug = load(AUG)
# split augmented: hold out some alert & compound for val
random.shuffle(aug)
n_aug_val = max(30, int(len(aug) * 0.2))
aug_val, aug_train = aug[:n_aug_val], aug[n_aug_val:]
# from curated, keep a representative grounded sample for train, plus a small val set
random.shuffle(cur)
val_cur = cur[:300]
pool_cur = cur[300:]
# limit grounded pool in train to balance (alert ~200, keep grounded ~ few thousand)
GROUND_TRAIN_CAP = 3000
train_cur = pool_cur[:GROUND_TRAIN_CAP]
train = aug_train + train_cur
val = aug_val + val_cur
# dedup train by utterance
seen = set()
train_dedup = []
for r in train:
u = r["utterance"]
if u in seen:
continue
seen.add(u)
train_dedup.append(r)
from collections import Counter
def mode_of(r):
return json.loads(r["messages"][1]["content"]).get("mode", "grounded")
print("train:", len(train_dedup), Counter(mode_of(r) for r in train_dedup))
print("val :", len(val), Counter(mode_of(r) for r in val))
for name, recs in (("router_sft_train.jsonl", train_dedup), ("router_sft_val.jsonl", val)):
with open(os.path.join(ROOT, name), "w", encoding="utf-8") as f:
for r in recs:
f.write(json.dumps(r, ensure_ascii=False) + "\n")
print("saved ->", os.path.join(ROOT, name))