56 lines
2.6 KiB
Python
56 lines
2.6 KiB
Python
|
"""样本加载：ml/data/samples/*.{X,Y,V,SG,TID}.npy + *.tasks.jsonl（ml/prep/build_dataset.cjs 产出）
|
|||
|
|
python -m ml.data --index ml/data/samples # 汇总样本/任务数，写 index.json
|
|||
|
|
全部载入内存（200 万样本约 1.5 GB），按任务 split 字段划分 train / val。"""
|
|||
|
|
import argparse, glob, json, os
|
|||
|
|
import numpy as np
|
|||
|
|
import torch
|
|||
|
|
|
|||
|
|
|
|||
|
|
def load_tasks(dirs):
|
|||
|
|
tasks = []
|
|||
|
|
for d in dirs:
|
|||
|
|
for f in sorted(glob.glob(os.path.join(d, '*.tasks.jsonl'))):
|
|||
|
|
for line in open(f, encoding='utf-8'):
|
|||
|
|
if line.strip():
|
|||
|
|
tasks.append(json.loads(line))
|
|||
|
|
return tasks
|
|||
|
|
|
|||
|
|
|
|||
|
|
class Samples:
|
|||
|
|
def __init__(self, dirs, device='cpu', split=None):
|
|||
|
|
if isinstance(dirs, str): dirs = [dirs]
|
|||
|
|
xs, ys, vs, ss, ts = [], [], [], [], []
|
|||
|
|
for d in dirs:
|
|||
|
|
for fx in sorted(glob.glob(os.path.join(d, '*.X.npy'))):
|
|||
|
|
b = fx[:-6]
|
|||
|
|
xs.append(np.load(fx)); ys.append(np.load(b + '.Y.npy')); vs.append(np.load(b + '.V.npy'))
|
|||
|
|
ss.append(np.load(b + '.SG.npy')); ts.append(np.load(b + '.TID.npy'))
|
|||
|
|
if not xs: raise SystemExit(f'没有样本：{dirs}（先跑 ml/prep/run_prep.sh）')
|
|||
|
|
self.tasks = load_tasks(dirs)
|
|||
|
|
val_tid = {t['tid'] for t in self.tasks if t['split'] == 'val'}
|
|||
|
|
X, Y, V, SG, TID = map(np.concatenate, (xs, ys, vs, ss, ts))
|
|||
|
|
is_val = np.isin(TID, np.fromiter(val_tid, dtype=np.int64)) if val_tid else np.zeros(len(TID), bool)
|
|||
|
|
sel = {'train': ~is_val, 'val': is_val, None: np.ones(len(TID), bool)}[split]
|
|||
|
|
t = lambda a: torch.from_numpy(np.ascontiguousarray(a[sel])).to(device)
|
|||
|
|
self.X, self.Y, self.V, self.SG, self.TID = t(X), t(Y), t(V), t(SG), t(TID)
|
|||
|
|
self.n = self.X.shape[0]
|
|||
|
|
|
|||
|
|
def batch(self, bs, gen=None):
|
|||
|
|
i = torch.randint(0, self.n, (bs,), device=self.X.device, generator=gen)
|
|||
|
|
return self.X[i], self.Y[i], self.V[i], self.SG[i]
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
ap = argparse.ArgumentParser(); ap.add_argument('--index', default='ml/data/samples'); a = ap.parse_args()
|
|||
|
|
tr, va = Samples(a.index, split='train'), Samples(a.index, split='val')
|
|||
|
|
tasks = load_tasks([a.index])
|
|||
|
|
cats = {}
|
|||
|
|
for t in tasks: cats[t['cat'] + '/' + t['mode']] = cats.get(t['cat'] + '/' + t['mode'], 0) + 1
|
|||
|
|
info = {'train_samples': tr.n, 'val_samples': va.n, 'tasks': len(tasks), 'val_tasks': sum(t['split'] == 'val' for t in tasks), 'by_cat_mode': cats}
|
|||
|
|
json.dump(info, open(os.path.join(a.index, 'index.json'), 'w'), indent=1)
|
|||
|
|
print(json.dumps(info, indent=1))
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
main()
|