feat(tasks): describe the answer's form with TaskContext.response

.github/workflows/umi-v2.yml

@@ -11,9 +11,9 @@ jobs:
   coordinated:
     name: UMI v2
     # Update the workflow and runner revision together.
-    uses: KartikP/brainscore-unified/.github/workflows/offline-tests.yml@af5309500c48b9cc27a6a468ebd2db93f7ec883f
+    uses: KartikP/brainscore-unified/.github/workflows/offline-tests.yml@5bf978697cd312116ac9aadbf297f8b0da61a364
     with:
       candidate: core
       candidate_repository: ${{ github.event.pull_request.head.repo.full_name || github.repository }}
       candidate_revision: ${{ github.event.pull_request.head.sha || github.sha }}
-      unified_revision: af5309500c48b9cc27a6a468ebd2db93f7ec883f
+      unified_revision: 5bf978697cd312116ac9aadbf297f8b0da61a364

brainscore_core/_model/behavior.py

@@ -24,6 +24,28 @@
 )
 
 
+class RatingRegressor:
+    """Ridge regression readout for numeric answers (ratings).
+
+    Scales features, fits ridge regression on numeric labels, predicts one
+    value per presentation. The rating counterpart of ProbabilitiesClassifier.
+    """
+
+    def __init__(self, alpha: float = 1.0):
+        import sklearn.linear_model
+        import sklearn.preprocessing
+        self._regressor = sklearn.linear_model.Ridge(alpha=alpha)
+        self._scaler = sklearn.preprocessing.StandardScaler()
+
+    def fit(self, features, labels: Iterable) -> 'RatingRegressor':
+        X = self._scaler.fit_transform(np.asarray(features))
+        self._regressor.fit(X, np.asarray(list(labels), dtype=float))
+        return self
+
+    def predict(self, features) -> np.ndarray:
+        return self._regressor.predict(self._scaler.transform(np.asarray(features)))
+
+
 class ProbabilitiesClassifier:
     """Logistic regression readout that produces label probabilities.
 

brainscore_core/_sessions/drivers.py

@@ -402,19 +402,24 @@ def drive_behavior_and_neural_session(model, session, driver: str = "interact")
     if time_bins is not None and len(time_bins) > 1:
         raise NotImplementedError(
             "Recording during a task returns one time bin; request at most one.")
-    regions = []
-    for channel in _requested_output_channels(session):
-        family, region = parse_channel(channel)
-        if family == "neural":
-            regions.append(region)
+    requested = _requested_output_channels(session)
+    regions = [parse_channel(c)[1] for c in requested if parse_channel(c)[0] == "neural"]
+    answers = [c for c in requested if parse_channel(c)[0] == "behavior"]
     _refuse_variance(model, session)
     model.start_task(task_context)
     stimuli = _behavior_stimuli_from_session(session, task_context)
     if stimuli is None:
         raise ValueError("Recording during a task needs stimuli to present.")
-    behavior, neural = model._behavioral.respond_and_record(stimuli, regions)
-    if "behavior" in _requested_output_channels(session):
-        session.emit(StreamEvent(channel="behavior", payload=behavior, t_ms=0.0,
+    if task_context.response is not None:
+        behavior, neural = model._behavioral.respond(stimuli, answers, regions)
+    else:
+        if set(answers) - {"behavior"}:
+            raise ValueError(f"{answers} need TaskContext(response=...); task_type "
+                             f"{task_context.task_type!r} answers on 'behavior'.")
+        plain, neural = model._behavioral.respond_and_record(stimuli, regions)
+        behavior = {"behavior": plain} if answers else {}
+    for channel in answers:
+        session.emit(StreamEvent(channel=channel, payload=behavior[channel], t_ms=0.0,
                                  meta={"driver": driver, "task_type": task_context.task_type}))
     for region in regions:
         session.emit(StreamEvent(channel=f"neural:{region}", payload=neural[region], t_ms=0.0,

brainscore_core/brainscore_model.py

@@ -94,6 +94,8 @@ def __init__(
         answer: Optional[Callable] = None,
         policy: Optional[Callable] = None,
         intervention: Optional[Callable] = None,
+        reply: Optional[Callable] = None,
+        likelihood: Optional[Callable] = None,
         behavioral_readout_layer: Optional[str] = None,
         generation_fn: Optional[Callable] = None,
         action_fn: Optional[Callable] = None,
@@ -130,6 +132,13 @@ def __init__(
         :param answer: ``(stimulus_row, instruction, label_set) -> label``;
             generates and parses a behavioral answer for instruction-following
             tasks.
+        :param reply: ``(stimulus_row, prompt) -> str``; the model's raw reply to
+            a prompt the task builds (instruction plus the response's directive).
+            Parsing is the task's. Gives choices, ratings, reported confidence
+            and written answers.
+        :param likelihood: ``(stimulus_row, prompt, candidates) -> {candidate:
+            log probability}``; scores candidate replies without generating.
+            Gives choices and their implied probabilities.
         :param policy: ``EnvironmentStep -> EnvironmentResponse``; chooses an
             action for closed-loop / embodied evaluation.
         :param intervention: ``StateChange -> (PerturbationApplied, cleanup)``;
@@ -223,6 +232,7 @@ def __init__(
         self._behavioral_readout_layer = capability_values.get(
             'behavioral_readout_layer')
         self._generation_fn = capability_values.get('generation_fn')
+        self._reply_fn, self._likelihood_fn = reply, likelihood
         self._action_fn = capability_values.get('action_fn')
         self._state_change_fn = capability_values.get('state_change_fn')
 
@@ -613,7 +623,8 @@ def reset(self) -> None:
                 failures.append(error)
         # Callable capabilities and extractors may own episode/history state.
         # Reset each provider once, including bound-method registrations.
-        providers = [self._action_fn, self._generation_fn, self._activations_model,
+        providers = [self._action_fn, self._generation_fn, self._reply_fn, self._likelihood_fn,
+                     self._activations_model,
                      *self._preprocessors.values(), *self._extractors.values()]
         seen = set()
         from .capabilities import capability_registry

brainscore_core/capabilities/sessions.py

@@ -60,22 +60,38 @@ class BehaviorSessions(_SessionCapability):
 
     def enabled_for(self, model) -> bool:
         return (model._behavioral_readout_layer is not None
-                or model._generation_fn is not None)
+                or model._generation_fn is not None
+                or model._reply_fn is not None or model._likelihood_fn is not None)
 
     def declared_input_channels(self, model):
-        return {'instruction'} if model._generation_fn is not None else set()
+        replies = (model._generation_fn, model._reply_fn, model._likelihood_fn)
+        return {'instruction'} if any(part is not None for part in replies) else set()
 
     def declared_output_channels(self, model):
-        return {'behavior'}
+        # Each part declares the kinds of answer it can give (rule 2.4).
+        channels = {'behavior'}
+        if model._behavioral_readout_layer is not None:
+            channels |= {'behavior:choice', 'behavior:probabilities', 'behavior:rating'}
+        if model._generation_fn is not None:
+            channels |= {'behavior:choice'}
+        if model._reply_fn is not None:
+            channels |= {'behavior:choice', 'behavior:rating', 'behavior:confidence',
+                         'behavior:text'}
+        if model._likelihood_fn is not None:
+            channels |= {'behavior:choice', 'behavior:probabilities'}
+        return channels
 
     def supports_session(self, model, channels) -> bool:
         # Behavior alone, or behavior plus regions recorded during the same answers.
-        recorded = set(channels) - {'behavior'}
-        return 'behavior' in channels and recorded <= {
+        behavior = {c for c in channels if c == 'behavior' or c.startswith('behavior:')}
+        recorded = set(channels) - behavior
+        return bool(behavior) and behavior <= self.declared_output_channels(model) and recorded <= {
             f'neural:{region}' for region in model.region_layer_map}
 
     def participate(self, model, session) -> None:
-        if set(_drivers._requested_output_channels(session)) != {'behavior'}:
+        requested = set(_drivers._requested_output_channels(session))
+        task = getattr(session, 'task_context', None)
+        if requested != {'behavior'} or (task is not None and task.response is not None):
             _drivers.drive_behavior_and_neural_session(model, session)
         elif getattr(session, 'streaming', False):
             _drivers._drive_behavior_session_streaming(model, session)

brainscore_core/contract.py

@@ -39,19 +39,42 @@ class TaskContext:
     answer (for example a held-out test set) in ``metadata['stimulus_set']``
     when using ``behavioral_response``.
 
+    ``response`` is the form of the answer (``brainscore_core.responses``):
+    ``Choice``, ``Scale`` or ``Text``; ``task_type='passive'`` means no answer.
+    With a response, outputs are requested on ``behavior:<kind>`` channels.
+    The older ``task_type`` values ``'label'``, ``'classification'`` and
+    ``'probabilities'`` keep their meaning on the plain ``behavior`` channel.
+
     ``readout_fit`` says which network a readout is fitted on while a change
     (a lesion, a stimulation) is active: ``'current'`` (default) fits on the
     changed network; ``'intact'`` fits on the network with every change undone,
     then tests the changed network with that readout, as an animal trained
     before a lesion is tested after it.
     """
-    task_type: str
+    task_type: Optional[str] = None
     label_set: Optional[List[str]] = None
     fitting_stimuli: Optional[Any] = None
     instruction: Optional[str] = None
     prefer_path: Optional[str] = None  # 'auto' | 'generation' | 'readout'
     metadata: Dict[str, Any] = field(default_factory=dict)
     readout_fit: str = 'current'  # 'current' | 'intact'
+    response: Any = None  # a brainscore_core.responses Choice, Scale or Text
+
+    def __post_init__(self):
+        # The response, when given, decides the task type; label lists follow it.
+        kinds = {'Choice': 'choice', 'Scale': 'rating', 'Text': 'text'}
+        if self.response is not None:
+            kind = kinds.get(type(self.response).__name__)
+            if kind is None:
+                raise TypeError('response must be a Choice, Scale or Text from brainscore_core.responses')
+            if self.task_type is None:
+                self.task_type = kind
+            elif self.task_type != kind:
+                raise ValueError(f'task_type {self.task_type!r} contradicts a {type(self.response).__name__} response')
+            if kind == 'choice' and self.label_set is None:
+                self.label_set = list(self.response.labels)
+        elif self.task_type is None:
+            raise ValueError("Give a response (Choice, Scale, Text) or task_type='passive'.")
 
 
 

brainscore_core/io_catalog.py

@@ -359,6 +359,8 @@ def validate(
                  "a label, a probability vector, or generated text",
                  "a fitted readout or the model's generation path",
                  meta_keys=("label_set",),
+                 addressing=("address is the kind of answer a TaskContext response asks for: "
+                             "choice, probabilities, confidence, rating or text"),
                  materializers=("start_task", "BehavioralAssembly")),
     # --- outputs: motor (carried by EnvironmentResponse, harness-defined) ---
     CatalogEntry("motor", OUTPUT, "EnvironmentResponse.action",

brainscore_core/responses.py

@@ -0,0 +1,120 @@
+"""The form a behavioral answer takes.
+
+A ``TaskContext`` says what the subject is asked to do; its ``response`` says
+what form the answer takes. A response knows how to phrase itself after the
+instruction (``directive``), how to read a reply (``parse``), and which output
+channels it produces. A benchmark may subclass a response to change either;
+prompt wording and parsing are the benchmark's science, not the model's.
+
+Reported answers come from the model's reply. Implied measures come from the
+model's internals by a stated link: a readout's or a language model's
+probabilities over the labels (``behavior:probabilities``). Every result says
+which it is.
+"""
+import re
+from dataclasses import dataclass
+from typing import Optional, Sequence, Tuple
+
+
+@dataclass(frozen=True)
+class Scale:
+    """A number on a scale, such as a 1 to 7 rating. ``step=None`` allows any value."""
+    low: float
+    high: float
+    step: Optional[float] = None
+
+    channels = ('behavior:rating',)
+
+    def directive(self) -> str:
+        whole = self.step is not None and float(self.step).is_integer()
+        bounds = (f'{int(self.low)} to {int(self.high)}' if whole
+                  else f'{self.low} to {self.high}')
+        return f'Answer with a number from {bounds}. Reply with only that number.'
+
+    def parse(self, reply: str) -> Tuple[Optional[float], Optional[str]]:
+        found = re.search(r'-?\d+(?:\.\d+)?', reply or '')
+        if found is None:
+            return None, 'no number in the reply'
+        value = float(found.group())
+        if not self.low <= value <= self.high:
+            return None, f'{value} is outside {self.low} to {self.high}'
+        if self.step is not None and abs((value - self.low) / self.step
+                                         - round((value - self.low) / self.step)) > 1e-9:
+            return None, f'{value} is not on the scale step {self.step}'
+        return value, None
+
+
+@dataclass(frozen=True)
+class Choice:
+    """One label from ``labels``.
+
+    ``probabilities=True`` also returns the implied probability of each label
+    (from a readout or a language model's likelihoods). ``confidence=Scale(...)``
+    also asks the subject to report how sure it is, in the same reply.
+    """
+    labels: Tuple[str, ...]
+    probabilities: bool = False
+    confidence: Optional[Scale] = None
+
+    def __post_init__(self):
+        if isinstance(self.labels, str) or len(self.labels) < 2:
+            raise ValueError('Choice needs at least two labels')
+        object.__setattr__(self, 'labels', tuple(self.labels))
+
+    @property
+    def channels(self) -> Tuple[str, ...]:
+        return (('behavior:choice',)
+                + (('behavior:probabilities',) if self.probabilities else ())
+                + (('behavior:confidence',) if self.confidence is not None else ()))
+
+    def directive(self) -> str:
+        text = f'Answer with exactly one of: {", ".join(self.labels)}.'
+        if self.confidence is not None:
+            text += (f' Then, on a new line, rate your confidence from '
+                     f'{self.confidence.low:g} to {self.confidence.high:g}.')
+        return text
+
+    def parse(self, reply: str) -> Tuple[Optional[str], Optional[str]]:
+        first = (reply or '').strip().splitlines()[0].strip().lower() if (reply or '').strip() else ''
+        for label in self.labels:
+            if first == label.lower():
+                return label, None
+        found = [label for label in self.labels
+                 if re.search(rf'\b{re.escape(label.lower())}\b', first)]
+        if len(found) == 1:
+            return found[0], None
+        return None, ('no label in the reply' if not found
+                      else f'several labels in the reply: {found}')
+
+    def parse_confidence(self, reply: str) -> Tuple[Optional[float], Optional[str]]:
+        lines = (reply or '').strip().splitlines()
+        if len(lines) < 2:
+            return None, 'no confidence line in the reply'
+        return self.confidence.parse(' '.join(lines[1:]))
+
+
+@dataclass(frozen=True)
+class Text:
+    """A written answer, kept as written. Brain-Score never scores text itself:
+    the benchmark brings its own parser, judge or comparison."""
+    channels = ('behavior:text',)
+
+    def directive(self) -> str:
+        return ''
+
+    def parse(self, reply: str) -> Tuple[Optional[str], Optional[str]]:
+        return (reply, None) if isinstance(reply, str) else (None, 'the reply is not text')
+
+
+def response_channels(response) -> Tuple[str, ...]:
+    return () if response is None else tuple(response.channels)
+
+
+def prompt_for(instruction: Optional[str], response) -> str:
+    """The text a replying model receives: the instruction, then the directive."""
+    parts = [part for part in ((instruction or '').strip(), response.directive()) if part]
+    return '\n\n'.join(parts)
+
+
+def labels_of(response) -> Optional[Sequence[str]]:
+    return response.labels if isinstance(response, Choice) else None

tests/test_compatibility.py

@@ -394,7 +394,9 @@ def _model(self):
 
     def test_every_part_declares_registered_channels(self):
         model = self._model()
-        assert model.out_channels == {'neural:IT', 'behavior', 'motor', 'perturbation'}
+        assert model.out_channels == {'neural:IT', 'behavior', 'behavior:choice',
+                                      'behavior:probabilities', 'behavior:rating',
+                                      'motor', 'perturbation'}
         assert {'vision', 'observation', 'lesion', 'stimulation'} <= model.in_channels
         for output in sorted(model.out_channels):
             bench = FakeBenchmark(identifier=output)

tests/test_subject_channels.py

@@ -99,7 +99,8 @@ def test_brainscore_model_gets_channel_identity_from_preprocessors():
     )
 
     assert model.in_channels == {"vision"}
-    assert model.out_channels == {"neural:IT", "behavior"}
+    assert model.out_channels == {"neural:IT", "behavior", "behavior:choice",
+                                  "behavior:probabilities", "behavior:rating"}
     assert model.required_channels == set()