обновлен процесс обработки данных

assembly_dataset/__main__.py

@@ -1,6 +1,7 @@
 import argparse
 import asyncio
 import json
+import signal
 from pathlib import Path
 import sys
 
@@ -8,6 +9,18 @@ from .config import load_config
 from .pipeline import generate, export, validate
 
 
+async def run_generation(config, *, show_progress, teacher_only=False):
+    loop = asyncio.get_running_loop()
+    task = asyncio.current_task()
+    previous = signal.getsignal(signal.SIGTERM)
+    loop.add_signal_handler(signal.SIGTERM, task.cancel)
+    try:
+        return await generate(config, show_progress=show_progress, teacher_only=teacher_only)
+    finally:
+        loop.remove_signal_handler(signal.SIGTERM)
+        signal.signal(signal.SIGTERM, previous)
+
+
 def main():
     parser = argparse.ArgumentParser(description='Датасет высокоуровневых планов сборки: full / step / branch')
     parser.add_argument('--config', type=Path, default=Path(__file__).resolve().parents[1] / 'config.yaml')
@@ -15,6 +28,8 @@ def main():
     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('--teacher-only', action='store_true', help='Сохранять только проверенные ответы LLM; останавливать генерацию при символической замене')
+    parser.add_argument('--allow-partial', action='store_true', help='Экспортировать накопленные записи, даже если цель ещё не достигнута')
     parser.add_argument('--no-progress', action='store_true', help='Отключить индикаторы прогресса')
     args = parser.parse_args()
     try:
@@ -32,11 +47,20 @@ def main():
             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))
+            try:
+                asyncio.run(run_generation(config, show_progress=not args.no_progress, teacher_only=args.teacher_only))
+            except (Exception, KeyboardInterrupt, asyncio.CancelledError) as exc:
+                print(f'Генерация остановлена ({type(exc).__name__}). Экспорт накопленных данных…', file=sys.stderr)
+                try:
+                    saved = export(config, show_progress=not args.no_progress, allow_partial=True)
+                    print(f"Сохранено {saved['examples']} записей: {config.output_dir / 'sft'}", file=sys.stderr)
+                except Exception as export_error:
+                    print(f'Экспорт не выполнен ({type(export_error).__name__}); контрольные точки остаются в {config.output_dir / "episodes.sqlite3"}', file=sys.stderr)
+                return 130 if isinstance(exc, (KeyboardInterrupt, asyncio.CancelledError)) else 1
             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 = export(config, show_progress=not args.no_progress, allow_partial=args.allow_partial)
             result = {k: result[k] for k in ('examples', 'scenarios', 'counts', 'teacher_failures')}
         else:
             result = validate(config)