File size: 4,398 Bytes
3467e0f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | #!/usr/bin/env python3
"""Отфильтровать готовый датасет (`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 # noqa: E402 (правила отбраковки капшенов)
import train as T # noqa: E402 (только find_dataset_chunks — как train.py ищет чанки)
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()
|