remove some code in favor of embedding and clip wrappers, also make clip wrapper more flexible

Author: lucidrains

perfusion_pytorch/embedding.py

@@ -226,21 +226,28 @@ def forward(
 # and on forward, passes the concept embeddings + superclass concept embeddings through the text transformer + final layernorm
 # as well as make the forward pass the ids and superclass_ids through the modified text encoder twice (will attempt to substitute the nn.Embedding with an nn.Identity)
 
-from open_clip import CLIP
-
 class OpenClipEmbedWrapper(Module):
     @beartype
     def __init__(
         self,
-        clip: CLIP,
+        clip: Module,
+        text_transformer_path = 'transformer',
+        ln_final_path = 'ln_final',  # in CLIP, they had the final layernorm separate from the transformer
         **embedding_wrapper_kwargs
     ):
         super().__init__()
         self.wrapped_embed = EmbeddingWrapper(clip.token_embedding, **embedding_wrapper_kwargs)
 
+        path_to_modules = dict([(path, mod) for path, mod in clip.named_modules()])
+
+        assert text_transformer_path in path_to_modules
+
+        text_transformer = path_to_modules[text_transformer_path]
+        ln_final = path_to_modules.get(ln_final_path, nn.Identity())
+
         self.text_transformer = nn.Sequential(
-            clip.transformer,
-            clip.ln_final
+            text_transformer,
+            ln_final
         )
 
     def forward(

perfusion_pytorch/perfusion.py

@@ -69,38 +69,6 @@ def calculate_input_covariance(
 
     return einsum('n d, n e -> d e', all_embeds, all_embeds) / length
 
-@beartype
-def find_first_index(
-    indices: IndicesTensor,
-    concept_or_superclass_id: int
-):
-    """
-    for deriving the concept_indices to be passed into the Rank1EditModule
-    """
-
-    edge = (indices == concept_or_superclass_id).cumsum(dim = -1)  # [1, 3, 5, 4, 1, 1], 4 -> [0, 0, 0, 1, 0, 0, 0] -> [0, 0, 0, 1, 1, 1]
-    return edge.sum(dim = -1)
-
-@beartype
-def return_text_enc_with_concept_and_superclass(
-    text_ids: IndicesTensor,
-    concept_id: int,
-    superclass_id: int,
-    clip: Optional[OpenClipAdapter] = None
-):
-    batch = text_ids.shape[0]
-    batch_arange = torch.arange(batch, device = text_ids.device)
-    concept_indices = find_first_index(text_ids, concept_id)
-    text_ids_with_superclass = text_ids[batch_arange, concept_indices] = superclass_ids
-
-    if not exists(clip):
-        return text_ids, concept_indices, text_ids_with_superclass
-
-    concept_text_enc = clip.embed_texts(text_ids)
-    superclass_text_enc = clip.embed_texts(text_ids_with_superclass)
-
-    return concept_text_enc, concept_indices, superclass_text_enc
-
 # loss weighted by the mask
 
 @beartype

setup.py

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