cleanup

README.md

@@ -34,10 +34,10 @@ import torch
 
 def extract_embedding(tensor):
     cls_embedding, patch_embeddings = tensor[:, 0, :], tensor[:, 1:, :]
-    return torch.cat(cls_embedding, patch_embeddings.mean(1), dim=-1)
+    return torch.cat([cls_embedding, patch_embeddings.mean(1)], dim=-1)
 
-image_batch = transform(image).unsqueeze(dim=0)
-embedding = extract_embedding(model(image_batch).last_hidden_state)
+batch = transform(image).unsqueeze(dim=0)
+embedding = extract_embedding(model(batch).last_hidden_state)
 print(f"Embeddings shape: {embedding.shape}")
 ```
 
@@ -46,26 +46,13 @@ print(f"Embeddings shape: {embedding.shape}")
 ```python
 import math
 import torch
-from transformers import modeling_outputs
-from typing_extensions import override
 
 # for segmentation
-class ExtractPatchFeatures:
-    """Extracts the patch features from a model output."""
-    def __call__(self, tensor: torch.Tensor) -> torch.Tensor:
-        """Call method for the transformation.
-
-        Args:
-            tensor: The raw embeddings of the model.
-
-        Returns:
-            A tensor (batch_size, hidden_size, n_patches_height, n_patches_width)
-            representing the model output.
-        """
-        features = tensor[:, 1:, :].permute(0, 2, 1)
-        batch_size, hidden_size, patch_grid = features.shape
-        height = width = int(math.sqrt(patch_grid))
-        return features.view(batch_size, hidden_size, height, width)
+def extract_patch_embeddings(tensor):
+    features = tensor[:, 1:, :].permute(0, 2, 1)
+    batch_size, hidden_size, patch_grid = features.shape
+    height = width = int(math.sqrt(patch_grid))
+    return features.view(batch_size, hidden_size, height, width)
 
 extract_embeddings = ExtractPatchFeatures()
 emb = extract_embeddings(model(transform(image)[None]).last_hidden_state)