allow for easy transformation of text and superclass embeds into text encodings and superclass text encodings, if the clip transformer (post embed -> post transformer embed after the last layernorm) is passed in

Author: lucidrains

perfusion_pytorch/embedding.py

@@ -118,7 +118,8 @@ def forward(
         self,
         x: Union[Tensor, List[str]],
         concept_id: Optional[Union[int, Tuple[int, ...]]] = None,
-        return_embed_with_superclass = True
+        return_embed_with_superclass = True,
+        clip_transformer_fn: Optional[Callable[[Tensor], Tensor]] = None
     ) -> EmbeddingReturn:
 
         assert not (self.training and self.num_concepts > 1), 'cannot train with multiple concepts'
@@ -195,15 +196,30 @@ def forward(
         # if training, and superclass embed id given
         # also return embeddings with superclass, for deriving superclass_text_enc
 
+        superclass_embeds = None
+
         if self.training and exists(self.superclass_embed_id) and return_embed_with_superclass:
             x = x.masked_fill(concept_masks[0], self.superclass_embed_id)
 
             with torch.no_grad():
                 superclass_embeds = self.embed(x)
 
-            return EmbeddingReturn(embeds, superclass_embeds, embed_mask, concept_indices)
+        # if the clip transformer function is passed in, transform the embeds and superclass_embeds into the text_enc and superclass_text_enc, to be forwarded by cross attentions into the Rank1EditModules
+
+        if exists(clip_transformer_fn):
+            with torch.no_grad():
+                embeds = clip_transformer_fn(embeds)
+
+                if exists(superclass_embeds):
+                    superclass_embeds = clip_transformer_fn(superclass_embeds)
+
+        # return tuple, with
+        # 1. text embeds | encodings
+        # 2. superclass text embeds | encoding
+        # 3. text mask
+        # 4. concept indices
 
-        return EmbeddingReturn(embeds, None, embed_mask, concept_indices)
+        return EmbeddingReturn(embeds, superclass_embeds, embed_mask, concept_indices)
 
 @beartype
 def merge_embedding_wrappers(

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'perfusion-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.1.14',
+  version = '0.1.15',
   license='MIT',
   description = 'Perfusion - Pytorch',
   author = 'Phil Wang',