for embedding wrapper, allow for passing in concept ids, which would assert that the necessary concept embedding ids are present for each prompt across the batch

Author: lucidrains

perfusion_pytorch/embedding.py

@@ -1,13 +1,25 @@
 import torch
 from torch import nn
 from torch.nn import Module
+
 from beartype import beartype
+from beartype.typing import Optional, Tuple, Union
 
 from einops import rearrange
 
+# helper functions
+
 def exists(val):
     return val is not None
 
+def is_all_unique(arr):
+    return len(set(arr)) == len(arr)
+
+def filter_tuple_indices(tup, indices):
+    return tuple(tup[i] for i in indices)
+
+# embedding wrapper class
+
 class EmbeddingWrapper(Module):
     @beartype
     def __init__(
@@ -19,17 +31,34 @@ def __init__(
         self.embed = embed
         num_embeds, dim = embed.weight.shape
 
+        self.num_embeds = num_embeds
         self.num_concepts = num_concepts
         self.concepts = nn.Parameter(torch.zeros(num_concepts, dim))
         nn.init.normal_(self.concepts, std = 0.02)
 
-        self.concept_ids = tuple(range(num_embeds, num_embeds + num_concepts))
+        self.concept_embed_ids = tuple(range(num_embeds, num_embeds + num_concepts))
 
     def parameters(self):
         return [self.concepts]
 
-    def forward(self, x):
-        concept_masks = tuple(concept_id == x for concept_id in self.concept_ids)
+    def forward(
+        self,
+        x,
+        concept_id: Optional[Union[int, Tuple[int, ...]]] = None
+    ):
+        concept_masks = tuple(concept_id == x for concept_id in self.concept_embed_ids)
+
+        if exists(concept_id):
+            if not isinstance(concept_id, tuple):
+                concept_id = (concept_id,)
+
+            assert is_all_unique(concept_id), 'concept ids must be all unique'
+            assert all([cid < self.num_concepts for cid in concept_id])
+
+            has_concept = tuple(concept_mask.any(dim = -1).all() for concept_mask in concept_masks)
+
+            assert all(filter_tuple_indices(has_concept, concept_id)), f'concept ids {filter_tuple_indices(self.concept_embed_ids, concept_id)} not found in ids passed in'
+            concept_masks = filter_tuple_indices(concept_masks, concept_id)
 
         for concept_mask in concept_masks:
             x = x.masked_fill(concept_mask, 0)
@@ -38,7 +67,7 @@ def forward(self, x):
             embeds = self.embed(x)
             embeds.detach_()
 
-        for concept_id, concept, concept_mask in zip(self.concept_ids, self.concepts, concept_masks):
+        for concept, concept_mask in zip(self.concepts, concept_masks):
             embeds = torch.where(
                 rearrange(concept_mask, '... -> ... 1'),
                 concept,

perfusion_pytorch/perfusion.py

@@ -20,7 +20,7 @@
 def exists(val):
     return val is not None
 
-def all_unique(arr):
+def is_all_unique(arr):
     return len(set(arr)) == len(arr)
 
 IndicesTensor = Union[LongTensor, IntTensor]
@@ -254,7 +254,7 @@ def forward(
 
         if is_multi_concepts:
             assert not self.training, 'multi concepts can only be done at inference'
-            assert all_unique(concept_id)
+            assert is_all_unique(concept_id)
             assert all([cid < self.num_concepts for cid in concept_id])
 
             concept_id_tensor = torch.tensor(concept_id, dtype = torch.long, device = self.device)

setup.py

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