From fde167655f4e10b31efd796e3b946ffa761a825a Mon Sep 17 00:00:00 2001
From: Martin Schrimpf <m4rtinsch@gmail.com>
Date: Sun, 21 Mar 2021 14:51:05 -0400
Subject: [PATCH] fit to fitting_stimuli in LabelBehavior instead of assuming
 pretrained imagenet

---
 model_tools/brain_transformation/__init__.py |  4 +--
 model_tools/brain_transformation/behavior.py | 36 ++++++++------------
 2 files changed, 17 insertions(+), 23 deletions(-)

diff --git a/model_tools/brain_transformation/__init__.py b/model_tools/brain_transformation/__init__.py
index 19e2750..683b8ab 100644
--- a/model_tools/brain_transformation/__init__.py
+++ b/model_tools/brain_transformation/__init__.py
@@ -3,7 +3,7 @@
 from brainscore.model_interface import BrainModel
 from brainscore.utils import LazyLoad
 from model_tools.brain_transformation.temporal import TemporalIgnore
-from .behavior import BehaviorArbiter, LogitsBehavior, ProbabilitiesMapping
+from .behavior import BehaviorArbiter, LabelBehavior, ProbabilitiesMapping
 from .neural import LayerMappedModel, LayerSelection, LayerScores
 
 
@@ -28,10 +28,10 @@ def __init__(self, identifier,
         self.region_benchmarks = {**self.standard_region_benchmarks, **(region_benchmarks or {})}
         layer_model = LayerMappedModel(identifier=identifier, activations_model=activations_model)
         self.layer_model = TemporalIgnore(layer_model)
-        logits_behavior = LogitsBehavior(identifier=identifier, activations_model=activations_model)
         behavioral_readout_layer = behavioral_readout_layer or layers[-1]
         probabilities_behavior = ProbabilitiesMapping(identifier=identifier, activations_model=activations_model,
                                                       layer=behavioral_readout_layer)
+        logits_behavior = LabelBehavior(probabilities_mapping=probabilities_behavior)
         self.behavior_model = BehaviorArbiter({BrainModel.Task.label: logits_behavior,
                                                BrainModel.Task.probabilities: probabilities_behavior})
         self.do_behavior = False
diff --git a/model_tools/brain_transformation/behavior.py b/model_tools/brain_transformation/behavior.py
index b70d05e..7027cf9 100644
--- a/model_tools/brain_transformation/behavior.py
+++ b/model_tools/brain_transformation/behavior.py
@@ -1,4 +1,3 @@
-import os
 from collections import OrderedDict
 
 import sklearn.linear_model
@@ -22,32 +21,27 @@ def look_at(self, stimuli, *args, **kwargs):
         return self.current_executor.look_at(stimuli, *args, **kwargs)
 
 
-class LogitsBehavior(BrainModel):
-    def __init__(self, identifier, activations_model):
-        self.identifier = identifier
-        self.activations_model = activations_model
-        self.current_task = None
+class LabelBehavior(BrainModel):
+    def __init__(self, probabilities_mapping):
+        self.probabilities_mapping = probabilities_mapping
 
     def start_task(self, task: BrainModel.Task, fitting_stimuli):
         assert task in [BrainModel.Task.passive, BrainModel.Task.label]
         if task == BrainModel.Task.label:
-            assert fitting_stimuli == 'imagenet'
-        self.current_task = task
+            self.probabilities_mapping.start_task(BrainModel.Task.probabilities, fitting_stimuli=fitting_stimuli)
 
     def look_at(self, stimuli, number_of_trials=1):
-        if self.current_task is BrainModel.Task.passive:
-            return
-        logits = self.activations_model(stimuli, layers=['logits'])
-        assert len(logits['neuroid']) == 1000
-        logits = logits.transpose('presentation', 'neuroid')
-        prediction_indices = logits.values.argmax(axis=1)
-        with open(os.path.join(os.path.dirname(__file__), 'imagenet_classes.txt')) as f:
-            synsets = f.read().splitlines()
-        prediction_synsets = [synsets[index] for index in prediction_indices]
-        return BehavioralAssembly([prediction_synsets], coords={
-            **{coord: (dims, values) for coord, dims, values in walk_coords(logits['presentation'])},
-            **{'synset': ('presentation', prediction_synsets), 'logit': ('presentation', prediction_indices)}},
-                                  dims=['choice', 'presentation'])
+        probabilities = self.probabilities_mapping.look_at(stimuli=stimuli, number_of_trials=number_of_trials)
+        assert len(probabilities['choice']) == 1_000
+        prediction_indices = probabilities.argmax('choice').values
+        predictions = probabilities['choice'].values[prediction_indices]
+        return BehavioralAssembly(predictions, coords={
+            **{coord: (dims, values) for coord, dims, values in walk_coords(probabilities['presentation'])},
+            **{'synset': ('presentation', predictions),
+               'logit': ('presentation', predictions),
+               'label': ('presentation', predictions),
+               'choice': ('presentation', predictions)}},
+                                  dims=['presentation'])
 
 
 class ProbabilitiesMapping(BrainModel):