обновлен формат генерации датасета

assembly_dataset/__main__.py

@@ -0,0 +1,51 @@
+import argparse
+import asyncio
+import json
+from pathlib import Path
+import sys
+
+from .config import load_config
+from .pipeline import generate, export, validate
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Датасет высокоуровневых планов сборки: full / step / branch')
+    parser.add_argument('--config', type=Path, default=Path(__file__).resolve().parents[1] / 'config.yaml')
+    parser.add_argument('command', choices=['generate', 'export', 'validate'], nargs='?', default='generate')
+    parser.add_argument('--count', type=int, help='Точное суммарное число SFT-записей; 0 допустим')
+    parser.add_argument('--output-dir', type=Path, help='Папка корпуса (относительно текущего каталога)')
+    parser.add_argument('--dry-run', action='store_true', help='Показать настройки без записи и запросов')
+    parser.add_argument('--no-progress', action='store_true', help='Отключить индикаторы прогресса')
+    args = parser.parse_args()
+    try:
+        config = load_config(args.config)
+        data = config.model_dump()
+        if args.count is not None:
+            data['target_examples'] = args.count
+        if args.output_dir is not None:
+            data['output_dir'] = args.output_dir.resolve()
+        config = type(config).model_validate(data)
+        if args.dry_run:
+            print(config.model_dump_json(indent=2))
+            return 0
+        if args.command == 'generate':
+            if config.teacher.enabled:
+                from dotenv import load_dotenv
+                load_dotenv(Path(__file__).resolve().parents[1] / '.env')
+            asyncio.run(generate(config, show_progress=not args.no_progress))
+            result = export(config, show_progress=not args.no_progress)
+            result = {k: result[k] for k in ('examples', 'scenarios', 'counts', 'teacher_failures')}
+        elif args.command == 'export':
+            result = export(config, show_progress=not args.no_progress)
+            result = {k: result[k] for k in ('examples', 'scenarios', 'counts', 'teacher_failures')}
+        else:
+            result = validate(config)
+        print(json.dumps(result, ensure_ascii=False, indent=2))
+        return 0
+    except (ValueError, OSError) as exc:
+        print(f'Ошибка: {exc}', file=sys.stderr)
+        return 1
+
+
+if __name__ == '__main__':
+    raise SystemExit(main())