Upload folder using huggingface_hub

src/embeddings/audio_embedder.py

@@ -25,21 +25,28 @@ class AudioEmbedder:
         self.model.to(self.device)
         self.model.eval()
 
-    def _squeeze_features(self, feats: torch.Tensor, projection: str) -> torch.Tensor:
-        """Ensure features are 1-D projected embeddings (512-d).
+    def _extract_features(self, output, projection: str) -> torch.Tensor:
+        """Extract 1-D projected embedding (512-d) from model output.
 
-        Some transformers versions return raw hidden states (batch, seq, hidden)
-        instead of projected features (batch, proj_dim). Detect and fix.
+        Handles both raw tensors and BaseModelOutputWithPooling objects
+        across different transformers versions.
         """
-        if feats.dim() == 3:
-            pooled = feats[:, 0, :]  # CLS token
+        if not isinstance(output, torch.Tensor):
+            # BaseModelOutputWithPooling — extract pooled features and project
+            pooled = output.pooler_output
             proj = getattr(self.model, projection, None)
             if proj is not None:
                 pooled = proj(pooled)
-            feats = pooled
-        if feats.dim() == 2:
-            feats = feats[0]
-        return feats
+            output = pooled
+        if output.dim() == 3:
+            pooled = output[:, 0, :]
+            proj = getattr(self.model, projection, None)
+            if proj is not None:
+                pooled = proj(pooled)
+            output = pooled
+        if output.dim() == 2:
+            output = output[0]
+        return output
 
     @torch.no_grad()
     def embed(self, audio_path: str) -> np.ndarray:
@@ -52,7 +59,7 @@ class AudioEmbedder:
         ).to(self.device)
 
         outputs = self.model.get_audio_features(**inputs)
-        emb = self._squeeze_features(outputs, "audio_projection")
+        emb = self._extract_features(outputs, "audio_projection")
         return emb.cpu().numpy().astype("float32")
 
     @torch.no_grad()
@@ -64,5 +71,5 @@ class AudioEmbedder:
             padding=True,
         ).to(self.device)
         feats = self.model.get_text_features(**inputs)
-        feats = self._squeeze_features(feats, "text_projection")
+        feats = self._extract_features(feats, "text_projection")
         return feats.cpu().numpy().astype("float32")

src/embeddings/image_embedder.py

@@ -26,7 +26,14 @@ class ImageEmbedder:
         image = Image.open(image_path).convert("RGB")
         inputs = self.processor(images=image, return_tensors="pt").to(self.device)
         feats = self.model.get_image_features(**inputs)
-        # Handle different transformers versions (some return 3-D hidden states)
+        # Handle different transformers versions
+        if not isinstance(feats, torch.Tensor):
+            # BaseModelOutputWithPooling — extract and project
+            pooled = feats.pooler_output
+            proj = getattr(self.model, "visual_projection", None)
+            if proj is not None:
+                pooled = proj(pooled)
+            feats = pooled
         if feats.dim() == 3:
             pooled = feats[:, 0, :]
             proj = getattr(self.model, "visual_projection", None)

src/embeddings/text_embedder.py

@@ -29,7 +29,14 @@ class TextEmbedder:
             truncation=True,
         ).to(self.device)
         feats = self.model.get_text_features(**inputs)
-        # Handle different transformers versions (some return 3-D hidden states)
+        # Handle different transformers versions
+        if not isinstance(feats, torch.Tensor):
+            # BaseModelOutputWithPooling — extract and project
+            pooled = feats.pooler_output
+            proj = getattr(self.model, "text_projection", None)
+            if proj is not None:
+                pooled = proj(pooled)
+            feats = pooled
         if feats.dim() == 3:
             pooled = feats[:, 0, :]
             proj = getattr(self.model, "text_projection", None)