обновлен формат генерации датасета
This commit is contained in:
Даниил Грабарь
@@ -0,0 +1,247 @@
|
||||
"""Bounded generation, transactional checkpoints, streaming sharded export."""
|
||||
import asyncio
|
||||
import csv
|
||||
from collections import Counter
|
||||
from contextlib import contextmanager
|
||||
import fcntl
|
||||
import hashlib
|
||||
import json
|
||||
from itertools import zip_longest
|
||||
from pathlib import Path
|
||||
import shutil
|
||||
import sqlite3
|
||||
import tempfile
|
||||
|
||||
from .config import Config
|
||||
from .language import validate_texts
|
||||
from .csv_export import COLUMNS, csv_record
|
||||
from .progress import ProgressDisplay
|
||||
from .world import canonical, digest, make_episode, training_rows, validate_plan, validate_branch
|
||||
|
||||
VERSION = 3
|
||||
STORAGE_VERSION = 2
|
||||
|
||||
|
||||
@contextmanager
|
||||
def storage(config):
|
||||
root = config.output_dir
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
with (root / '.lock').open('a') as lock:
|
||||
try:
|
||||
fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
except BlockingIOError:
|
||||
raise ValueError('Другой процесс уже работает с output_dir') from None
|
||||
db = sqlite3.connect(root / 'episodes.sqlite3')
|
||||
try:
|
||||
db.execute('PRAGMA journal_mode=WAL')
|
||||
db.execute('CREATE TABLE IF NOT EXISTS meta (key TEXT PRIMARY KEY, value TEXT NOT NULL)')
|
||||
db.execute('CREATE TABLE IF NOT EXISTS episodes (seq INTEGER PRIMARY KEY, signature TEXT UNIQUE NOT NULL, payload TEXT NOT NULL, rows INTEGER NOT NULL)')
|
||||
settings = config.model_dump(mode='json', exclude={'target_examples', 'output_dir', 'shard_size', 'prompts_file'})
|
||||
# Structured prompts only affect export, so existing v2 episodes remain reusable.
|
||||
settings['prompts'].pop('structured_system', None)
|
||||
fingerprint = digest({'version': STORAGE_VERSION, 'settings': settings})
|
||||
saved = db.execute("SELECT value FROM meta WHERE key='fingerprint'").fetchone()
|
||||
if saved and saved[0] != fingerprint:
|
||||
raise ValueError('Изменён config: для других настроек используйте новый output_dir')
|
||||
if not saved:
|
||||
with db:
|
||||
db.execute('INSERT INTO meta VALUES (?, ?)', ('fingerprint', fingerprint))
|
||||
db.execute('INSERT INTO meta VALUES (?, ?)', ('settings', canonical(settings)))
|
||||
db.execute('INSERT INTO meta VALUES (?, ?)', ('next_index', '0'))
|
||||
yield db
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
async def generate(config: Config, agent=None, *, show_progress=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])
|
||||
with ProgressDisplay('Генерация', config.target_examples, initial=count, enabled=show_progress) as progress:
|
||||
if count >= config.target_examples:
|
||||
return count
|
||||
if config.teacher.enabled and config.teacher.fraction > 0 and agent is None:
|
||||
from .teacher import create_agent
|
||||
progress.update(status='подключение LLM')
|
||||
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:
|
||||
# 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)]
|
||||
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))]
|
||||
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')
|
||||
progress.update(count, status='')
|
||||
return count
|
||||
|
||||
|
||||
class Shards:
|
||||
def __init__(self, root, size):
|
||||
self.root, self.size = root, size
|
||||
self.counts = Counter()
|
||||
self.handles = {}
|
||||
self.csv_handles = {}
|
||||
self.csv_writers = {}
|
||||
self.files = []
|
||||
|
||||
def write(self, row, structured_row):
|
||||
key = row['split'], row['mode']
|
||||
count = self.counts[key]
|
||||
for variant, record in (('text', row), ('structured', structured_row)):
|
||||
handle_key = (*key, variant)
|
||||
if count % self.size == 0:
|
||||
if handle_key in self.handles:
|
||||
self.handles[handle_key].close()
|
||||
self.csv_handles[handle_key].close()
|
||||
suffix = '.jsonl' if variant == 'text' else '.structured.jsonl'
|
||||
path = self.root / key[0] / key[1] / f'{count // self.size:06d}{suffix}'
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self.handles[handle_key] = path.open('w', encoding='utf-8')
|
||||
self.files.append(path)
|
||||
csv_path = path.with_suffix('.csv')
|
||||
self.csv_handles[handle_key] = csv_path.open('w', encoding='utf-8-sig', newline='')
|
||||
self.csv_writers[handle_key] = csv.DictWriter(self.csv_handles[handle_key], fieldnames=COLUMNS, delimiter=';')
|
||||
self.csv_writers[handle_key].writeheader()
|
||||
self.files.append(csv_path)
|
||||
self.handles[handle_key].write(canonical(record) + '\n')
|
||||
self.csv_writers[handle_key].writerow(csv_record(record))
|
||||
self.counts[key] += 1
|
||||
|
||||
def close(self):
|
||||
for file in (*self.handles.values(), *self.csv_handles.values()):
|
||||
file.close()
|
||||
|
||||
|
||||
def file_digest(path):
|
||||
hasher = hashlib.sha256()
|
||||
with path.open('rb') as file:
|
||||
for block in iter(lambda: file.read(1024 * 1024), b''):
|
||||
hasher.update(block)
|
||||
return hasher.hexdigest()
|
||||
|
||||
|
||||
def export(config, *, show_progress=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:
|
||||
raise ValueError('Недостаточно сценариев: сначала выполните generate')
|
||||
destination = config.output_dir / 'sft'
|
||||
backup = config.output_dir / '.sft-previous'
|
||||
# Recover an interrupted directory swap before making another export.
|
||||
if backup.exists() and not destination.exists():
|
||||
backup.rename(destination)
|
||||
if destination.exists() and not (destination / 'manifest.json').exists():
|
||||
raise ValueError('sft не является управляемым экспортом; выберите другой output_dir')
|
||||
stage = Path(tempfile.mkdtemp(prefix='.sft-', dir=config.output_dir))
|
||||
writer = Shards(stage, config.shard_size)
|
||||
total, scenarios, failures = 0, 0, 0
|
||||
domains, sources = Counter(), Counter()
|
||||
try:
|
||||
for (payload,) in db.execute('SELECT payload FROM episodes ORDER BY seq'):
|
||||
if total >= config.target_examples:
|
||||
break
|
||||
episode = json.loads(payload)
|
||||
scenarios += 1
|
||||
domains[episode['scene']['domain']] += 1
|
||||
sources[episode['source']['provider']] += 1
|
||||
failures += int('teacher_error' in episode)
|
||||
for record, structured_record in zip(training_rows(episode, config),
|
||||
training_rows(episode, config, output_format='structured'), strict=True):
|
||||
if total >= config.target_examples:
|
||||
break
|
||||
writer.write(record, structured_record)
|
||||
total += 1
|
||||
progress.update(total)
|
||||
progress.update(status="контрольные суммы")
|
||||
writer.close()
|
||||
manifest = {'schema_version': VERSION, 'examples': total, 'scenarios': scenarios,
|
||||
'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,
|
||||
'validation_level': 'logical_only',
|
||||
'config': config.model_dump(mode='json'),
|
||||
'files': {str(p.relative_to(stage)): file_digest(p) for p in writer.files}}
|
||||
(stage / 'manifest.json').write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding='utf-8')
|
||||
if backup.exists():
|
||||
shutil.rmtree(backup)
|
||||
if destination.exists():
|
||||
destination.rename(backup)
|
||||
stage.rename(destination)
|
||||
if backup.exists():
|
||||
shutil.rmtree(backup)
|
||||
progress.update(total, status='')
|
||||
return manifest
|
||||
finally:
|
||||
writer.close()
|
||||
if stage.exists():
|
||||
shutil.rmtree(stage)
|
||||
|
||||
|
||||
def validate(config):
|
||||
with storage(config) as db:
|
||||
episodes, paths = 0, 0
|
||||
for signature, payload, stored_rows in db.execute('SELECT signature,payload,rows FROM episodes'):
|
||||
episode = json.loads(payload)
|
||||
if signature != digest(episode['scene']) or episode['signature'] != signature or episode['scenario_id'] != signature:
|
||||
raise ValueError('Неверная сигнатура сценария')
|
||||
if 'step_texts' in episode:
|
||||
validate_texts(episode, episode['step_texts'], config.prompts)
|
||||
trace = validate_plan(episode, episode['plan'])
|
||||
if trace != episode['trace']:
|
||||
raise ValueError('Trace не соответствует плану')
|
||||
paths += validate_branch(episode, episode['branch'])
|
||||
if sum(1 for _ in training_rows(episode, config)) != stored_rows:
|
||||
raise ValueError('Неверное число записей эпизода')
|
||||
episodes += 1
|
||||
root = config.output_dir / 'sft'
|
||||
manifest_path = root / 'manifest.json'
|
||||
if not manifest_path.exists():
|
||||
raise ValueError('Нет экспорта SFT: выполните generate или export')
|
||||
manifest = json.loads(manifest_path.read_text())
|
||||
total = 0
|
||||
expected_files = set(manifest['files'])
|
||||
actual_files = {str(p.relative_to(root)) for pattern in ('*/*/*.jsonl', '*/*/*.csv') for p in root.glob(pattern)}
|
||||
if expected_files != actual_files:
|
||||
raise ValueError('Набор шардов не соответствует manifest')
|
||||
for relative, checksum in manifest['files'].items():
|
||||
path = root / relative
|
||||
if path.resolve().is_relative_to(root.resolve()) is False or file_digest(path) != checksum:
|
||||
raise ValueError(f'Повреждён шард {relative}')
|
||||
for relative in manifest['files']:
|
||||
path = root / relative
|
||||
if path.suffix == '.jsonl' and '.structured.' not in path.name:
|
||||
counterpart = path.with_suffix('.structured.jsonl')
|
||||
if str(counterpart.relative_to(root)) not in expected_files:
|
||||
raise ValueError('Отсутствует структурированная пара JSONL')
|
||||
with path.open(encoding='utf-8') as file, counterpart.open(encoding='utf-8') as other:
|
||||
for text_line, structured_line in zip_longest(file, other):
|
||||
if text_line is None or structured_line is None:
|
||||
raise ValueError('Разное число записей в паре JSONL')
|
||||
text_row, structured_row = json.loads(text_line), json.loads(structured_line)
|
||||
if (text_row['record_id'] != structured_row['record_id'] or
|
||||
text_row['messages'][1] != structured_row['messages'][1]):
|
||||
raise ValueError('Записи в паре JSONL не совпадают')
|
||||
total += 1
|
||||
if total != manifest['examples']:
|
||||
raise ValueError('Неверное число экспортированных записей')
|
||||
return {'episodes_checked': episodes, 'branch_paths_checked': paths, 'examples_checked': total}
|
||||
Reference in New Issue
Block a user