src/fpgm/prompting/task_prompt.py

"""Turn a DROID task instruction into a SAM 3.1 text prompt.

DROID episode metadata carries a free-form ``current_task`` sentence, e.g.
``"Put brick in drawer shelf and close drawer"``. SAM 3.1's text prompt wants a
short noun phrase naming the object to segment, not a full imperative sentence.
This module derives that phrase with a small dependency-free heuristic --
no spaCy/NLTK, to keep the environment lean -- rather than a parser: strip the
leading imperative verb (and its particle, e.g. "pick *up*"), then split the
remainder on prepositions/"and" so the first noun phrase (the direct object) is
the primary candidate and later phrases (destinations, containers, second
clauses) become fallback candidates for retrying a SAM 3.1 query that found
nothing.
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field

from fpgm.config import SegmentationConfig

# Leading imperative verbs seen across DROID instructions. Only the *first*
# token is ever checked against this set (see `_split_leading_verb`), so it is
# deliberately broad -- a false positive here just means one word is dropped
# from an already-imperative sentence, which is harmless.
_VERBS = frozenset(
    {
        "pick",
        "place",
        "put",
        "move",
        "push",
        "pull",
        "open",
        "close",
        "grasp",
        "lift",
        "take",
        "set",
        "insert",
        "pour",
        "wipe",
        "press",
        "turn",
        "rotate",
        "slide",
        "stack",
        "drop",
        "remove",
        "grab",
        "hold",
        "release",
        "flip",
        "sweep",
        "cover",
        "uncover",
        "screw",
        "unscrew",
        "twist",
        "squeeze",
        "touch",
        "tap",
        "nudge",
        "align",
        "arrange",
        "sort",
        "clean",
        "fill",
        "empty",
        "throw",
        "toss",
        "carry",
        "transfer",
        "swap",
        "position",
    }
)

# Particles that can immediately follow a leading verb as part of a phrasal
# verb ("pick *up*", "take *out*") rather than starting the object phrase.
_PARTICLES = frozenset({"up", "down", "out", "off", "away", "back", "aside"})

# Determiners/pronouns dropped from every phrase -- they never belong in a SAM
# 3.1 text prompt and stripping them everywhere (not just leading) is simpler
# than tracking phrase-internal position.
_STOPWORDS = frozenset(
    {"the", "a", "an", "it", "its", "this", "that", "these", "those", "them", "there"}
)

# Two-word connector checked before the single-word table below.
_TWO_WORD_CONNECTOR = ("next", "to")

# Prepositions/"and" that separate the direct object from a destination,
# container, or a second clause. Exactly the list called out in the design:
# in, into, on, onto, to, from, with, near, next to, inside, under, over, at, and.
_CONNECTORS = frozenset(
    {
        "in",
        "into",
        "on",
        "onto",
        "to",
        "from",
        "with",
        "near",
        "inside",
        "under",
        "over",
        "at",
        "and",
    }
)

_TOKEN_RE = re.compile(r"[A-Za-z']+")


@dataclass(frozen=True)
class PromptCandidates:
    """Ranked SAM 3.1 text-prompt candidates derived from a task instruction.

    Attributes:
        primary: Best-guess direct-object noun phrase, e.g. ``"brick"``.
        alternatives: Fallback phrases (destination, container, second clause
            object, ...) to retry if SAM 3.1 finds nothing for ``primary``.
        source: ``"override"`` if ``primary`` came from an explicit override
            rather than the heuristic.
    """

    primary: str
    alternatives: list[str] = field(default_factory=list)
    source: str = "heuristic"

    @property
    def all(self) -> list[str]:
        """``primary`` followed by ``alternatives``, in retry order."""
        return [self.primary, *self.alternatives]


class TaskPromptDeriver:
    """Derives SAM 3.1 text prompts from DROID ``current_task`` instructions."""

    def __init__(self, config: SegmentationConfig | None = None) -> None:
        """Args:
        config: If given, ``config.prompt_overrides`` is consulted by
            :meth:`derive` when an ``episode_uuid`` is supplied.
        """
        self.config = config

    def derive(
        self,
        instruction: str,
        override: str | None = None,
        episode_uuid: str | None = None,
    ) -> PromptCandidates:
        """Derive prompt candidates from a task instruction.

        Args:
            instruction: Raw DROID ``current_task`` sentence.
            override: Explicit prompt that wins outright over both the
                heuristic and any config-level override.
            episode_uuid: If given (and ``override`` is not), looked up in
                ``config.prompt_overrides``; a hit wins outright over the
                heuristic, exactly like ``override`` does.

        Returns:
            :class:`PromptCandidates` with ``primary`` plus ranked
            ``alternatives`` to retry if SAM 3.1 finds nothing.

        Raises:
            ValueError: If ``instruction`` has no usable words, or the
                heuristic cannot extract any noun phrase from it.
        """
        resolved_override = override
        if resolved_override is None and self.config is not None and episode_uuid is not None:
            resolved_override = self.config.prompt_overrides.get(episode_uuid)
        if resolved_override is not None:
            return PromptCandidates(primary=resolved_override, alternatives=[], source="override")

        tokens = _TOKEN_RE.findall(instruction.lower())
        if not tokens:
            raise ValueError(f"task_prompt: instruction has no words: {instruction!r}")

        remaining = _strip_leading_verb(tokens)
        if not remaining:
            raise ValueError(f"task_prompt: nothing left after the leading verb: {instruction!r}")

        chunks = _split_on_connectors(remaining)
        phrases = [p for p in (_clean_chunk(c) for c in chunks) if p]
        if not phrases:
            raise ValueError(f"task_prompt: could not extract a noun phrase from: {instruction!r}")

        primary, *rest = phrases
        alternatives: list[str] = []
        seen = {primary}
        for phrase in rest:
            if phrase not in seen:
                alternatives.append(phrase)
                seen.add(phrase)

        return PromptCandidates(primary=primary, alternatives=alternatives, source="heuristic")


def _strip_leading_verb(tokens: list[str]) -> list[str]:
    """Drop a recognised leading imperative verb and its particle, if present."""
    if tokens and tokens[0] in _VERBS:
        tokens = tokens[1:]
        if tokens and tokens[0] in _PARTICLES:
            tokens = tokens[1:]
    return tokens


def _split_on_connectors(tokens: list[str]) -> list[list[str]]:
    """Split a token list into phrases on prepositions/"and" (and "next to")."""
    chunks: list[list[str]] = [[]]
    i = 0
    n = len(tokens)
    while i < n:
        is_two_word = (
            tokens[i] == _TWO_WORD_CONNECTOR[0]
            and i + 1 < n
            and tokens[i + 1] == _TWO_WORD_CONNECTOR[1]
        )
        if is_two_word:
            chunks.append([])
            i += 2
            continue
        if tokens[i] in _CONNECTORS:
            chunks.append([])
            i += 1
            continue
        chunks[-1].append(tokens[i])
        i += 1
    return chunks


def _clean_chunk(words: list[str]) -> str:
    """Strip a secondary-clause leading verb (e.g. "and *close* drawer") and stopwords."""
    words = _strip_leading_verb(list(words))
    words = [w for w in words if w not in _STOPWORDS]
    return " ".join(words)