Remote code: batched embed_text_batch / embed_audio_batch entry points

modeling_fusion_embedding.py

@@ -439,6 +439,74 @@ class FusionEmbeddingModel(PreTrainedModel):
         pooled = last_token_pool(h, inputs["attention_mask"])
         return self._finish(pooled, dim)
 
+    # ------------------------------------------------------------- batched
+    @torch.no_grad()
+    def embed_text_batch(self, texts, instruction: str = DEFAULT_QUERY_INSTRUCTION,
+                         dim: Optional[int] = None,
+                         max_tokens: Optional[int] = None) -> torch.Tensor:
+        """Batch text embedding [B, dim] (right-padded, mask-aware last-token pooling)."""
+        self._ensure_backbones()
+        if self._rt["gate"].active:
+            raise RuntimeError("adapter gate is open during a text encode — "
+                               "non-audio inputs must run with the gate closed")
+        cfg, tok = self.config, self._rt["tok"]
+        max_tokens = max_tokens or cfg.max_text_tokens
+        seqs = [tok.encode(_chat(instruction, t), add_special_tokens=False)[:max_tokens]
+                for t in texts]
+        L = max(len(s) for s in seqs)
+        ids = torch.full((len(seqs), L), cfg.pad_id, dtype=torch.long, device=self._device)
+        mask = torch.zeros(len(seqs), L, dtype=torch.long, device=self._device)
+        for b, s in enumerate(seqs):
+            ids[b, : len(s)] = torch.tensor(s, device=self._device)
+            mask[b, : len(s)] = 1
+        full = self._rt["full"]
+        out = full.language_model(inputs_embeds=full.get_input_embeddings()(ids),
+                                  attention_mask=mask)
+        hidden = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0]
+        pooled = self.text_whitening(last_token_pool(hidden, mask))
+        return mrl_truncate_normalize(pooled.float(), dim or cfg.mrl_default).cpu()
+
+    @torch.no_grad()
+    def embed_audio_batch(self, wavs, sr: int, dim: Optional[int] = None) -> torch.Tensor:
+        """Batch audio embedding [B, dim] from raw waveform arrays at a common rate."""
+        import librosa
+        import numpy as np
+
+        self._ensure_backbones()
+        cfg, fe_audio = self.config, self._rt["fe_audio"]
+        target_sr = fe_audio.sampling_rate
+        prepped = []
+        for wav in wavs:
+            wav = np.asarray(wav, dtype=np.float32)
+            if wav.ndim > 1:
+                wav = wav.mean(axis=-1)
+            if sr != target_sr:
+                wav = librosa.resample(wav, orig_sr=sr, target_sr=target_sr)
+            prepped.append(wav)
+        feats = fe_audio(prepped, sampling_rate=target_sr, return_tensors="pt",
+                         return_attention_mask=True, padding="max_length", truncation=True)
+        mel, am = feats["input_features"], feats.get("attention_mask")
+        if am is not None:
+            tmax = int(am.sum(dim=1).max().item())
+            mel, am = mel[:, :, :tmax], am[:, :tmax]
+        fmask = (am.bool() if am is not None
+                 else torch.ones(mel.shape[0], mel.shape[2], dtype=torch.bool))
+        frames, frame_mask = self._rt["tower"](mel.to(self._device),
+                                               fmask.to(self._device))
+        audio_tok = self.resampler(frames, frame_mask)
+        ids = torch.tensor([[cfg.audio_pad_id] * cfg.n_query + [cfg.eos_id]] * mel.shape[0],
+                           device=self._device)
+        attention_mask = torch.ones_like(ids)
+        full = self._rt["full"]
+        embeds = full.get_input_embeddings()(ids).clone()
+        embeds[ids == cfg.audio_pad_id] = (
+            audio_tok.reshape(-1, audio_tok.size(-1)).to(embeds.dtype))
+        with self._rt["gate"]:
+            out = full.language_model(inputs_embeds=embeds, attention_mask=attention_mask)
+        hidden = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0]
+        pooled = last_token_pool(hidden, attention_mask)
+        return mrl_truncate_normalize(pooled.float(), dim or cfg.mrl_default).cpu()
+
     # ------------------------------------------------------------- read-out
     @staticmethod
     def center(embs: torch.Tensor) -> torch.Tensor: