refactored classification embedding logic

README.md

@@ -31,18 +31,14 @@ model = AutoModel.from_pretrained('kaiko-ai/midnight')
 ### Extract embeddings for classification
 ```python
 import torch
-from transformers import modeling_outputs
-from typing_extensions import override
 
-# for classification
-class ExtractConcatToken:
-    """Extracts the CLS with Mean Patch tokens from a model output."""
-    def __call__(self, tensor: torch.Tensor) -> torch.Tensor:
-        return torch.cat([tensor[:, 0, :], tensor[:, 1:, :].mean(1)], dim=-1)
+def extract_embedding(tensor):
+    cls_embedding, patch_embeddings = tensor[:, 0, :], tensor[:, 1:, :]
+    return torch.cat(cls_embedding, patch_embeddings.mean(1), dim=-1)
 
-extract_embeddings = ExtractConcatToken()
-emb = extract_embeddings(model(transform(image)[None]).last_hidden_state)
-print(f"Embeddings shape: {emb.shape}")
+image_batch = transform(image).unsqueeze(dim=0)
+embedding = extract_embedding(model(image_batch).last_hidden_state)
+print(f"Embeddings shape: {embedding.shape}")
 ```