| |
| """Отфильтровать готовый датасет (`make_dataset.py`) в новый, не пересчитывая латенты. |
| |
| Латенты VAE — самая дорогая часть нарезки (часы GPU), а фильтр работает по тексту капшена. |
| Поэтому здесь ничего не кодируется заново: читаем arrow-файлы, отбрасываем строки |
| по правилам из `train.py` (FILTER_RULES) и пишем урезанный датасет в `--save`. |
| |
| Структуру сохраняем один-в-один: каждый исходный `chunk_*/sub_*` даёт один новый |
| `chunk_*/sub_*` (только без отброшенных строк). Поэтому группы уже однородны по размеру |
| картинки — `make_dataset.py` писал их именно так, — и батчи в `train.py` остаются |
| одинаковыми по HxW. Никакой перегруппировки и, значит, никаких копий латентов в RAM. |
| |
| Запуск: |
| python3 src/filter_dataset.py --src src/work/ds_testg_full --save src/work/ds_testg_filtered |
| python3 src/filter_dataset.py --src ... --save ... --filter views,text,multi,long |
| """ |
| import argparse |
| import os |
| import shutil |
| import sys |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| import filters as F |
| import train as T |
|
|
|
|
| def parse_args(): |
| p = argparse.ArgumentParser(description="Датасет -> отфильтрованный датасет (без пересчёта латентов)") |
| p.add_argument("--src", required=True, help="папка с готовым датасетом (chunk_*/sub_*)") |
| p.add_argument("--save", required=True, help="куда сохранить отфильтрованный") |
| p.add_argument("--filter", default="all", |
| help=f"группы через запятую ({', '.join(F.FILTER_GROUPS)}) или 'all'") |
| return p.parse_args() |
|
|
|
|
| def main(): |
| args = parse_args() |
| groups = tuple(g.strip() for g in args.filter.split(",") if g.strip()) |
| if "all" in groups: |
| groups = F.FILTER_GROUPS |
| bad = [g for g in groups if g not in F.FILTER_GROUPS] |
| if bad: |
| raise SystemExit(f"--filter: неизвестные группы {bad}; есть {list(F.FILTER_GROUPS)}") |
|
|
| from datasets import load_from_disk |
| srcs = sorted(p for p in T.find_dataset_chunks(args.src) if p.startswith( |
| os.path.join(args.src, "chunk_"))) |
| if not srcs: |
| raise SystemExit(f"в {args.src} нет чанков chunk_*/sub_*") |
|
|
| print(f"группы: {', '.join(groups)}", flush=True) |
| os.makedirs(args.save, exist_ok=True) |
|
|
| dropped = {g: 0 for g in F.FILTER_GROUPS} |
| total = kept = 0 |
| for src in srcs: |
| ds = load_from_disk(src) |
| keep = [i for i, text in enumerate(ds["text"]) |
| if not _is_musor(text, groups, dropped)] |
| total += len(ds) |
| if not keep: |
| print(f" {os.path.relpath(src, args.src)}: 0 из {len(ds)} — группа пустая", flush=True) |
| continue |
| rel = os.path.relpath(src, args.src) |
| dst = os.path.join(args.save, rel) |
| if os.path.isdir(dst): |
| shutil.rmtree(dst) |
| os.makedirs(os.path.dirname(dst), exist_ok=True) |
| ds.select(keep).save_to_disk(dst) |
| open(os.path.join(os.path.dirname(dst), ".done"), "w").close() |
| kept += len(keep) |
| print(f" {rel}: {len(keep)} из {len(ds)}", flush=True) |
|
|
| total_dropped = total - kept |
| print(f"\nвсего {total}, отброшено {total_dropped} " |
| f"({total_dropped / max(total, 1):.1%}), осталось {kept}", flush=True) |
| for g, c in dropped.items(): |
| if c: |
| print(f" {g:8} {c:6d} {c / max(total, 1):5.1%}", flush=True) |
| print(f"готово -> {args.save}", flush=True) |
|
|
|
|
| def _is_musor(text, groups, dropped): |
| reason = F.filter_reason(text, groups) |
| if reason is not None: |
| dropped[reason] += 1 |
| return reason is not None |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
