one more tiny step

CALM_pytorch/CALM.py

@@ -143,10 +143,13 @@ def __call__(self, *args):
         hidden = self.get_output_fn(*args)
         self.output = hidden.detach()
 
-    def pop_saved(self):
+    def release_saved_(self):
+        self.output = None
+
+    def pop_saved_(self):
         output = self.output
         assert exists(output)
-        self.output = None
+        self.release_saved_()
         return output
 
 # cross attention wrapper class
@@ -204,7 +207,7 @@ def forward(self, *hook_args):
         x = get_block_output_from_hook_outputs(self.forward_hook_get_hidden, *hook_args)
 
         if self.release_recorder_output:
-            context = self.recorder.pop_saved()
+            context = self.recorder.pop_saved_()
         else:
             context = self.recorder.output
 
@@ -441,6 +444,16 @@ def load_state_dict(self, pkg, strict = False):
     def parameters(self):
         return self.cross_attns.parameters()
 
+    def set_release_recorder_output(self, release_recorder_output: bool):
+        for module in self.modules():
+            if isinstance(module, CrossAttentionBlock):
+                module.release_recorder_output = release_recorder_output
+
+    def clear_recorded_augment_hiddens(self):
+        for module in self.modules():
+            if isinstance(module, CrossAttentionBlock):
+                module.recorder.release_saved_()
+
     @beartype
     def forward(
         self,

README.md

@@ -248,7 +248,8 @@ loss.backward()
 
 - [ ] take care of caching the augment hiddens when sampling. forget about anchor kv cache for now
     - [x] logic for not releasing the saved output from recorder, for inference
-    - [ ] use a contextmanager for managing cross attention block state for popping the saved output from the recorder
+    - [x] managing cross attention block state for popping the saved output from the recorder
+    - [ ] move the augmentation forwards into one shared method, and craft out sampling method for anchor
 
 - [ ] show an example with giving the LLM ability to hear as well, using <a href="https://github.com/lucidrains/audiolm-pytorch">hubert or wav2vec</a> wrappers
 - [ ] handle a wrapper or function that takes in the sequence and prompt length, and auto derives the inputs to CALM

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'CALM-Pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.1.6',
+  version = '0.1.7',
   license='MIT',
   description = 'CALM - Pytorch',
   author = 'Phil Wang',