76 lines
4.3 KiB
Python
76 lines
4.3 KiB
Python
import argparse
|
|
import asyncio
|
|
import json
|
|
import signal
|
|
from pathlib import Path
|
|
import sys
|
|
|
|
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')
|
|
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('--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:
|
|
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')
|
|
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, allow_partial=args.allow_partial)
|
|
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())
|