61 lines
3.3 KiB
Python
61 lines
3.3 KiB
Python
import asyncio
|
|
import json
|
|
from pathlib import Path
|
|
import tempfile
|
|
import unittest
|
|
|
|
from assembly_dataset.config import Config
|
|
from assembly_dataset.pipeline import generate, export, validate
|
|
from assembly_dataset.world import make_episode, training_rows
|
|
|
|
|
|
class PairedExportTests(unittest.TestCase):
|
|
def test_paired_shards_and_exact_count(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
config = Config(output_dir=Path(tmp), target_examples=10, shard_size=2)
|
|
asyncio.run(generate(config))
|
|
report = export(config)
|
|
self.assertEqual(report['examples'], 10)
|
|
pairs = 0
|
|
for path in Path(tmp).glob('sft/*/*/[0-9]*.jsonl'):
|
|
if '.structured.' in path.name:
|
|
continue
|
|
structured_path = path.with_suffix('.structured.jsonl')
|
|
self.assertTrue(structured_path.exists())
|
|
text_rows = list(map(json.loads, path.read_text().splitlines()))
|
|
json_rows = list(map(json.loads, structured_path.read_text().splitlines()))
|
|
self.assertEqual(len(text_rows), len(json_rows))
|
|
for text, structured in zip(text_rows, json_rows):
|
|
self.assertEqual(text['record_id'], structured['record_id'])
|
|
self.assertEqual(text['messages'][1], structured['messages'][1])
|
|
self.assertNotEqual(text['messages'][0], structured['messages'][0])
|
|
answer = json.loads(structured['messages'][2]['content'])
|
|
self.assertEqual(set(answer), {'status', 'plan'})
|
|
self.assertNotRegex(str(answer), r'part_\d|slot_\d')
|
|
pairs += 1
|
|
self.assertEqual(pairs, 10)
|
|
self.assertEqual(validate(config)['examples_checked'], 10)
|
|
|
|
def test_structured_skills_and_branches(self):
|
|
config = Config(domains=['operator_assistance'], min_components=1, max_components=1,
|
|
blocked_fraction=0, step_sampling='all')
|
|
episode = make_episode(0, config)
|
|
rows = list(training_rows(episode, config, output_format='structured'))
|
|
full = json.loads(rows[0]['messages'][2]['content'])
|
|
actions = full['plan']
|
|
self.assertEqual(actions[2], {'skill': 'pick', 'object': 'распределительная коробка №1'})
|
|
self.assertEqual(actions[3]['target'], 'оператор')
|
|
self.assertEqual(actions[6]['target'], 'зона хранения №1')
|
|
self.assertNotIn('target', actions[7])
|
|
for row in rows[1:-1]:
|
|
self.assertEqual(len(json.loads(row['messages'][2]['content'])['plan']), 1)
|
|
branch = json.loads(rows[-1]['messages'][2]['content'])
|
|
for i, action in enumerate(branch['plan']):
|
|
self.assertEqual(action['on_failure'], 'replan_required')
|
|
self.assertEqual(action['on_success'], f'step_{i+2}' if i+1 < len(actions) else 'ready')
|
|
|
|
def test_blocked_empty_plan(self):
|
|
config = Config(domains=['operator_assistance'], min_components=1, max_components=1, blocked_fraction=1)
|
|
for row in training_rows(make_episode(0, config), config, output_format='structured'):
|
|
self.assertEqual(json.loads(row['messages'][2]['content']), {'status': 'blocked', 'plan': []})
|