Spaces:
Sleeping
Sleeping
assistive-robot-study / src /personalization_module /preference_acquisition /training_data /builder.py
| from __future__ import annotations | |
| import json | |
| from dataclasses import asdict, dataclass, field | |
| from pathlib import Path | |
| from typing import Any, Literal | |
| DecisionLabel = Literal["do_now", "do_later", "remind", "no_action"] | |
| MissingFieldStrategy = Literal["unknown", "synthetic"] | |
| _USER_STATE_FLAG_TO_TOKEN = { | |
| "user_asleep": "asleep", | |
| "user_in_rush": "in_rush", | |
| "user_injured_or_disabled": "injured_or_disabled", | |
| "user_nearby": "nearby", | |
| } | |
| _ENV_FLAG_DIRECT = { | |
| "adverse_weather": "adverse_weather", | |
| "guests_present": "guests_present", | |
| } | |
| _VALID_DECISIONS = {"do_now", "do_later", "remind", "no_action"} | |
| class BuilderConfig: | |
| language: Literal["en", "es"] = "en" | |
| missing_field_strategy: MissingFieldStrategy = "unknown" | |
| include_compact_task_fields: bool = True | |
| class BuiltTrainingSample: | |
| sample_id: str | |
| user_id: int | |
| user_external_id: str | |
| label_action: DecisionLabel | |
| action_input: dict[str, Any] | |
| context_input: dict[str, Any] | |
| structured_task_features: dict[str, Any] | |
| preference_snapshot: list[dict[str, Any]] = field(default_factory=list) | |
| source_metadata: dict[str, Any] = field(default_factory=dict) | |
| data_provenance: dict[str, str] = field(default_factory=dict) | |
| def to_dict(self) -> dict[str, Any]: | |
| return asdict(self) | |
| class TrainingDataBuilder: | |
| """ | |
| Builds model-ready samples from mapped TAACO JSONL files. | |
| Designed for MVP bootstrap training where feedback events are still sparse. | |
| """ | |
| def __init__(self, config: BuilderConfig | None = None): | |
| self.config = config or BuilderConfig() | |
| def build_from_taaco_mapped( | |
| self, | |
| *, | |
| input_path: Path, | |
| ) -> list[BuiltTrainingSample]: | |
| records = self._load_jsonl(input_path) | |
| persona_map = self._build_persona_id_map(records) | |
| out: list[BuiltTrainingSample] = [] | |
| for row in records: | |
| sample = self._build_sample(row=row, persona_id_map=persona_map) | |
| if sample is not None: | |
| out.append(sample) | |
| return out | |
| def write_jsonl(self, samples: list[BuiltTrainingSample], output_path: Path) -> None: | |
| output_path.parent.mkdir(parents=True, exist_ok=True) | |
| with output_path.open("w", encoding="utf-8") as f: | |
| for sample in samples: | |
| f.write(json.dumps(sample.to_dict(), ensure_ascii=False) + "\n") | |
| def summarize(self, samples: list[BuiltTrainingSample]) -> dict[str, Any]: | |
| label_counts: dict[str, int] = {} | |
| user_counts: dict[str, int] = {} | |
| with_time = 0 | |
| with_weekday = 0 | |
| with_available_objects = 0 | |
| for sample in samples: | |
| label_counts[sample.label_action] = label_counts.get(sample.label_action, 0) + 1 | |
| user_counts[sample.user_external_id] = ( | |
| user_counts.get(sample.user_external_id, 0) + 1 | |
| ) | |
| if sample.context_input.get("time_of_day", "unknown") != "unknown": | |
| with_time += 1 | |
| if sample.context_input.get("weekday", "unknown") != "unknown": | |
| with_weekday += 1 | |
| if sample.context_input.get("available_objects"): | |
| with_available_objects += 1 | |
| total = len(samples) | |
| return { | |
| "num_samples": total, | |
| "labels": label_counts, | |
| "users": user_counts, | |
| "coverage": { | |
| "time_of_day_known": with_time, | |
| "weekday_known": with_weekday, | |
| "available_objects_nonempty": with_available_objects, | |
| }, | |
| "config": asdict(self.config), | |
| } | |
| def _build_sample( | |
| self, | |
| *, | |
| row: dict[str, Any], | |
| persona_id_map: dict[str, int], | |
| ) -> BuiltTrainingSample | None: | |
| persona = str(row.get("source_persona", "unknown_persona")) | |
| user_id = persona_id_map[persona] | |
| task_input = self._select_task_input(row) | |
| if task_input is None: | |
| return None | |
| label = self._select_label(row) | |
| if label is None: | |
| return None | |
| action_text = str(task_input.get("action", "")).strip() | |
| activity = str(task_input.get("activity", "")).strip() or None | |
| locations = self._as_list(task_input.get("locations")) | |
| objects = self._as_list(task_input.get("objects")) | |
| conditions = self._as_list(task_input.get("conditions")) | |
| context_flags = task_input.get("context_flags") or {} | |
| context_input, provenance = self._build_context_input( | |
| locations=locations, | |
| objects=objects, | |
| conditions=conditions, | |
| context_flags=context_flags, | |
| user_busy=bool(task_input.get("user_busy", False)), | |
| quiet_hours=bool(task_input.get("quiet_hours", False)), | |
| ) | |
| sample_id = f"{persona}:{task_input.get('task_id', '')}:{row.get('source_preference_index', '')}" | |
| structured_task = { | |
| "kind": task_input.get("kind"), | |
| "urgency": task_input.get("urgency"), | |
| "sensitivity": task_input.get("sensitivity"), | |
| "user_busy": bool(task_input.get("user_busy", False)), | |
| "quiet_hours": bool(task_input.get("quiet_hours", False)), | |
| } | |
| if self.config.include_compact_task_fields: | |
| structured_task["conditions"] = conditions | |
| structured_task["context_flags"] = context_flags | |
| return BuiltTrainingSample( | |
| sample_id=sample_id, | |
| user_id=user_id, | |
| user_external_id=persona, | |
| label_action=label, | |
| action_input={ | |
| "action_text": action_text, | |
| "activity": activity, | |
| }, | |
| context_input=context_input, | |
| structured_task_features=structured_task, | |
| preference_snapshot=[], | |
| source_metadata={ | |
| "source_example_index": row.get("source_example_index"), | |
| "source_preference_index": row.get("source_preference_index"), | |
| "expected_action_taaco": self._extract_expected_action_taaco(row), | |
| "mapping_trace": row.get("mapping_trace", []), | |
| }, | |
| data_provenance=provenance, | |
| ) | |
| def _build_context_input( | |
| self, | |
| *, | |
| locations: list[str], | |
| objects: list[str], | |
| conditions: list[str], | |
| context_flags: dict[str, Any], | |
| user_busy: bool, | |
| quiet_hours: bool, | |
| ) -> tuple[dict[str, Any], dict[str, str]]: | |
| true_flags = {k for k, v in context_flags.items() if bool(v)} | |
| user_state: list[str] = [] | |
| for flag, token in _USER_STATE_FLAG_TO_TOKEN.items(): | |
| if flag in true_flags: | |
| user_state.append(token) | |
| if user_busy: | |
| user_state.append("busy") | |
| environment_flags: list[str] = [] | |
| for flag, token in _ENV_FLAG_DIRECT.items(): | |
| if flag in true_flags: | |
| environment_flags.append(token) | |
| if quiet_hours: | |
| environment_flags.append("quiet_hours") | |
| if "weekend" in true_flags: | |
| environment_flags.append("weekend") | |
| location_current = locations[0] if locations else None | |
| time_of_day = "unknown" | |
| weekday = "unknown" | |
| available_objects: list[str] = [] | |
| provenance = { | |
| "location_current": "derived_from_locations[0]", | |
| "objects_nearby": "direct_from_task.objects", | |
| "raw_conditions": "direct_from_task.conditions", | |
| "user_state": "derived_from_context_flags_and_user_busy", | |
| "environment_flags": "derived_from_context_flags_and_quiet_hours", | |
| } | |
| if self.config.missing_field_strategy == "synthetic": | |
| available_objects = sorted(set(objects)) | |
| provenance["available_objects"] = "synthetic_from_objects_nearby" | |
| if "early_morning" in true_flags: | |
| time_of_day = "morning" | |
| provenance["time_of_day"] = "synthetic_from_context_flags.early_morning" | |
| elif "user_asleep" in true_flags: | |
| time_of_day = "night" | |
| provenance["time_of_day"] = "synthetic_from_context_flags.user_asleep" | |
| else: | |
| provenance["time_of_day"] = "unknown_no_signal" | |
| if "weekend" in true_flags: | |
| weekday = "saturday" | |
| provenance["weekday"] = "synthetic_from_context_flags.weekend" | |
| else: | |
| provenance["weekday"] = "unknown_no_signal" | |
| else: | |
| provenance["available_objects"] = "unknown_default" | |
| provenance["time_of_day"] = "unknown_default" | |
| provenance["weekday"] = "unknown_default" | |
| return ( | |
| { | |
| "location_current": location_current, | |
| "objects_nearby": objects, | |
| "available_objects": available_objects, | |
| "raw_conditions": conditions, | |
| "time_of_day": time_of_day, | |
| "weekday": weekday, | |
| "user_state": sorted(set(user_state)), | |
| "environment_flags": sorted(set(environment_flags)), | |
| }, | |
| provenance, | |
| ) | |
| def _select_task_input(self, row: dict[str, Any]) -> dict[str, Any] | None: | |
| # Preferred format: all_personas_mapped_both.jsonl | |
| if self.config.language == "en" and isinstance(row.get("task_input_en"), dict): | |
| return row["task_input_en"] | |
| if self.config.language == "es" and isinstance(row.get("task_input_es"), dict): | |
| return row["task_input_es"] | |
| # Compact format: all_personas_mapped_en/es.jsonl | |
| compact = row.get("task_input") | |
| source = row.get("source") or {} | |
| if isinstance(compact, dict): | |
| if "action" not in compact and isinstance(source, dict): | |
| # Recover dropped fields from source payload. | |
| compact = dict(compact) | |
| compact["action"] = source.get("action", "") | |
| compact["activity"] = source.get("activity", "") | |
| compact["objects"] = self._as_list(source.get("object")) | |
| compact["locations"] = self._as_list(source.get("location")) | |
| compact["conditions"] = self._as_list(source.get("conditions")) | |
| compact["explanations"] = source.get("explanations", []) | |
| compact["explanations_opposing"] = source.get("explanations_opposing", []) | |
| compact["context_flags"] = row.get("context_flags", {}) | |
| return compact | |
| return None | |
| def _select_label(self, row: dict[str, Any]) -> DecisionLabel | None: | |
| expected = row.get("expected_action") | |
| if isinstance(expected, str): | |
| return self._normalize_label(expected) | |
| if isinstance(expected, dict): | |
| if self.config.language == "en": | |
| return self._normalize_label(expected.get("en")) | |
| if self.config.language == "es": | |
| # Training label remains canonical English token. | |
| return self._normalize_label(expected.get("taaco")) | |
| return None | |
| def _extract_expected_action_taaco(self, row: dict[str, Any]) -> str | None: | |
| expected = row.get("expected_action") | |
| if isinstance(expected, str): | |
| return expected | |
| if isinstance(expected, dict): | |
| taaco = expected.get("taaco") | |
| return str(taaco) if taaco is not None else None | |
| alt = row.get("expected_action_taaco") | |
| return str(alt) if alt is not None else None | |
| def _build_persona_id_map(self, rows: list[dict[str, Any]]) -> dict[str, int]: | |
| personas = sorted({str(r.get("source_persona", "unknown_persona")) for r in rows}) | |
| return {persona: i + 1 for i, persona in enumerate(personas)} | |
| def _normalize_label(self, label: Any) -> DecisionLabel | None: | |
| if label is None: | |
| return None | |
| text = str(label).strip().lower().replace(" ", "_") | |
| if text in _VALID_DECISIONS: | |
| return text # type: ignore[return-value] | |
| return None | |
| def _as_list(self, value: Any) -> list[str]: | |
| if value is None: | |
| return [] | |
| if isinstance(value, list): | |
| return [str(x) for x in value] | |
| return [str(value)] | |
| def _load_jsonl(self, path: Path) -> list[dict[str, Any]]: | |
| rows: list[dict[str, Any]] = [] | |
| with path.open("r", encoding="utf-8") as f: | |
| for raw in f: | |
| line = raw.strip() | |
| if not line: | |
| continue | |
| rows.append(json.loads(line)) | |
| return rows | |
| __all__ = [ | |
| "BuilderConfig", | |
| "BuiltTrainingSample", | |
| "DecisionLabel", | |
| "MissingFieldStrategy", | |
| "TrainingDataBuilder", | |
| ] | |