return the concept indices from the embedding wrapper as well

Author: lucidrains

README.md

@@ -92,11 +92,11 @@ wrapped_embed = EmbeddingWrapper(
 
 # now just pass in your prompts with the superclass id
 
-embeds_with_new_concept, embeds_with_superclass, embed_mask = wrapped_embed([
+embeds_with_new_concept, embeds_with_superclass, embed_mask, concept_indices = wrapped_embed([
     'a portrait of dog',
     'dog running through a green field',
     'a man walking his dog'
-]) # (3, 77, 512), (3, 77, 512), (3, 77)
+]) # (3, 77, 512), (3, 77, 512), (3, 77), (3,)
 
 # now pass both embeds through clip text transformer
 # the embed_mask needs to be passed to the cross attention as key padding mask

perfusion_pytorch/embedding.py

@@ -17,7 +17,8 @@
 EmbeddingReturn = namedtuple('EmbeddingReturn', [
     'embed_with_concept',
     'embed_with_superclass',
-    'embed_mask'
+    'embed_mask',
+    'concept_indices'
 ])
 
 # helper functions
@@ -120,6 +121,11 @@ def forward(
         return_embed_with_superclass = True
     ) -> EmbeddingReturn:
 
+        assert not (self.training and self.num_concepts > 1), 'cannot train with multiple concepts'
+
+        if self.training:
+            concept_id = default(concept_id, 0)
+
         if exists(concept_id):
             if not isinstance(concept_id, tuple):
                 concept_id = (concept_id,)
@@ -136,6 +142,7 @@ def forward(
             assert superclass_mask.any(dim = -1).all(), 'superclass embed id must be present for all prompts'
 
             # automatically replace the superclass id with the concept id
+
             x = torch.where(superclass_mask, inferred_concept_id, x)
 
         # get the embedding mask, defined as not padding id
@@ -177,6 +184,14 @@ def forward(
                 embeds
             )
 
+        # whether to return concept indices for the rank-1-edit modules
+
+        concept_indices = None
+
+        if self.training and exists(concept_id) and len(concept_id) == 1:
+            concept_mask, = concept_masks
+            concept_indices = (concept_mask.cumsum(dim = -1) == 0).sum(dim = -1).long()
+
         # if training, and superclass embed id given
         # also return embeddings with superclass, for deriving superclass_text_enc
 
@@ -186,9 +201,9 @@ def forward(
             with torch.no_grad():
                 superclass_embeds = self.embed(x)
 
-            return EmbeddingReturn(embeds, superclass_embeds, embed_mask)
+            return EmbeddingReturn(embeds, superclass_embeds, embed_mask, concept_indices)
 
-        return EmbeddingReturn(embeds, None, embed_mask)
+        return EmbeddingReturn(embeds, None, 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.10',
+  version = '0.1.11',
   license='MIT',
   description = 'Perfusion - Pytorch',
   author = 'Phil Wang',