обновлен формат генерации датасета

assembly_dataset/structured.py

@@ -0,0 +1,37 @@
+"""Public skill plans derived from the same validated actions as Russian text."""
+from .language import component_for, object_name, target_name
+
+PLACEMENT_SKILLS = {
+    'position', 'place', 'align', 'insert', 'insert_fastener', 'tighten',
+    'press_fit', 'route_wire', 'connect', 'seat_component', 'bond', 'clamp',
+    'weld', 'move_to_storage',
+}
+OPERATOR_SKILLS = {'move_to_operator', 'hold', 'handover'}
+
+
+def public_action(action, scene, prompts):
+    component = component_for(action, scene)
+    skill = action['skill']
+    result = {'skill': skill, 'object': object_name(component, scene, prompts)}
+    if skill in PLACEMENT_SKILLS:
+        result['target'] = target_name(component, scene)
+    elif skill in OPERATOR_SKILLS:
+        result['target'] = 'оператор'
+    elif skill == 'return_home':
+        result.update(object='робот', target='исходное положение')
+    return result
+
+
+def structured_answer(episode, mode, index, prompts):
+    actions = episode['plan']['actions']
+    status = episode['plan']['status']
+    if mode == 'step' and actions:
+        actions = actions[index:index + 1]
+        status = 'ready'
+    plan = [public_action(action, episode['scene'], prompts) for action in actions]
+    if mode == 'branch':
+        for i, action in enumerate(plan):
+            action.update(id=f'step_{i + 1}',
+                          on_success=f'step_{i + 2}' if i + 1 < len(plan) else status,
+                          on_failure='replan_required')
+    return {'status': status, 'plan': plan}