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

assembly_dataset/pipeline.py

@@ -53,12 +53,13 @@ def storage(config):
             db.close()
 
 
-async def generate(config: Config, agent=None, *, show_progress=False):
+async def generate(config: Config, agent=None, *, show_progress=False, teacher_only=False):
     with storage(config) as db:
         count = db.execute('SELECT COALESCE(SUM(rows), 0) FROM episodes').fetchone()[0]
         index = int(db.execute("SELECT value FROM meta WHERE key='next_index'").fetchone()[0])
+        resume_until = db.execute('SELECT COALESCE(MAX(seq), 0) FROM episodes').fetchone()[0]
         with ProgressDisplay('Генерация', config.target_examples, initial=count, enabled=show_progress) as progress:
-            if count >= config.target_examples:
+            if count >= config.target_examples and index >= resume_until:
                 return count
             if config.teacher.enabled and config.teacher.fraction > 0 and agent is None:
                 from .teacher import create_agent
@@ -66,28 +67,49 @@ async def generate(config: Config, agent=None, *, show_progress=False):
                 agent = await create_agent(config.teacher, config.prompts)
             batch_size = config.teacher.concurrency if agent is not None else 32
             duplicates = 0
-            while count < config.target_examples:
+            while count < config.target_examples or index < resume_until:
                 # Never construct 500k coroutines or keep the entire dataset in RAM.
                 size = min(batch_size, max(1, config.target_examples - count))
                 episodes = [make_episode(i, config) for i in range(index, index + size)]
+                async def save_episode(episode):
+                    nonlocal count, duplicates
+                    # A restarted batch can contain responses already committed before cancellation.
+                    if db.execute('SELECT 1 FROM episodes WHERE signature=?',
+                                  (episode['signature'],)).fetchone():
+                        duplicates += 1
+                        return
+                    if agent is not None:
+                        from .teacher import enrich
+                        episode, _ = await enrich(episode, config.teacher, agent, config.prompts)
+                    if teacher_only and (config.teacher.llm is None or
+                                         episode['source']['provider'] != config.teacher.llm.provider or
+                                         'teacher_error' in episode):
+                        raise ValueError('Режим --teacher-only: ответ модели не прошёл проверку; символическая замена не сохранена')
+                    rows = sum(1 for _ in training_rows(episode, config))
+                    # No await between receiving a validated response and committing it.
+                    with db:
+                        cursor = db.execute('INSERT OR IGNORE INTO episodes(seq,signature,payload,rows) VALUES (?,?,?,?)',
+                                            (episode['candidate_index'] + 1, episode['signature'], canonical(episode), rows))
+                    if cursor.rowcount:
+                        count += rows
+                        duplicates = 0
+                    else:
+                        duplicates += 1
+                    progress.update(count, status='ожидание LLM' if agent is not None else '')
+
                 if agent is not None:
-                    from .teacher import enrich
                     progress.update(status='ожидание LLM')
-                    episodes = [item[0] for item in await asyncio.gather(*(
-                        enrich(e, config.teacher, agent, config.prompts) for e in episodes))]
+                tasks = [asyncio.create_task(save_episode(episode)) for episode in episodes]
+                try:
+                    await asyncio.gather(*tasks)
+                finally:
+                    for task in tasks:
+                        if not task.done():
+                            task.cancel()
+                    await asyncio.gather(*tasks, return_exceptions=True)
+                # Advance only after the whole batch succeeds; holes are retried on resume.
+                index += size
                 with db:
-                    for episode in episodes:
-                        rows = sum(1 for _ in training_rows(episode, config))
-                        cursor = db.execute('INSERT OR IGNORE INTO episodes(signature,payload,rows) VALUES (?,?,?)',
-                                            (episode['signature'], canonical(episode), rows))
-                        if cursor.rowcount:
-                            count += rows
-                            duplicates = 0
-                        else:
-                            duplicates += 1
-                        index = episode['candidate_index'] + 1
-                        if count >= config.target_examples:
-                            break
                     db.execute("UPDATE meta SET value=? WHERE key='next_index'", (str(index),))
                 if duplicates > 10000:
                     raise ValueError('Пространство сценариев исчерпано: расширьте диапазон компонентов или domains')
@@ -140,10 +162,10 @@ def file_digest(path):
     return hasher.hexdigest()
 
 
-def export(config, *, show_progress=False):
+def export(config, *, show_progress=False, allow_partial=False):
     with storage(config) as db, ProgressDisplay('Экспорт', config.target_examples, enabled=show_progress) as progress:
         capacity = db.execute('SELECT COALESCE(SUM(rows), 0) FROM episodes').fetchone()[0]
-        if capacity < config.target_examples:
+        if capacity < config.target_examples and not allow_partial:
             raise ValueError('Недостаточно сценариев: сначала выполните generate')
         destination = config.output_dir / 'sft'
         backup = config.output_dir / '.sft-previous'
@@ -175,6 +197,7 @@ def export(config, *, show_progress=False):
             progress.update(status="контрольные суммы")
             writer.close()
             manifest = {'schema_version': VERSION, 'examples': total, 'scenarios': scenarios,
+                        'partial': total < config.target_examples,
                         'examples_per_variant': {'text': total, 'structured': total},
                         'counts': {f'{split}/{mode}': n for (split, mode), n in sorted(writer.counts.items())},
                         'domains': dict(domains), 'sources': dict(sources), 'teacher_failures': failures,