Upload CoLMbo weights, config, and source code

config.json

@@ -10,8 +10,15 @@
   "n_mels": 80,
   "embedding_dim": 192,
   "channel": 1024,
-  "prefix_length": 10,
-  "gpt_model_name": "gpt2",
+  "map_type": "mlp",
+  "prefix_size": 192,
+  "sid_prefix_length": 40,
+  "sid_prefix_length_clip": 40,
+  "num_layers": 8,
+  "norm_sid_emb": false,
+  "text_decoder": "gpt2",
+  "tok_len": 67,
+  "text_prefix_length": 10,
   "sample_rate": 16000,
   "torch_dtype": "float32"
 }

modeling_colmbo.py

@@ -1,12 +1,20 @@
 """
 modeling_colmbo.py — CoLMbo HuggingFace-compatible model wrapper.
-Registered with AutoModel via trust_remote_code=True.
+
+Faithfully wraps the original ExpWrapper inference pipeline so that
+users can run:
+
+    from transformers import AutoModel
+    model = AutoModel.from_pretrained("cmu-mlsp/CoLMbo", trust_remote_code=True)
+    text  = model.describe_file("audio.wav", "What is the speaker's dialect?")
 """
 
+import os
+import numpy as np
 import torch
 import torch.nn as nn
 import torchaudio
-from transformers import PreTrainedModel, PretrainedConfig
+from transformers import PreTrainedModel, PretrainedConfig, GPT2LMHeadModel, AutoTokenizer
 from transformers.modeling_outputs import BaseModelOutput
 
 
@@ -17,11 +25,22 @@ class CoLMboConfig(PretrainedConfig):
 
     def __init__(
         self,
+        # speaker encoder
         n_mels: int = 80,
         embedding_dim: int = 192,
         channel: int = 1024,
-        prefix_length: int = 10,
-        gpt_model_name: str = "gpt2",
+        # mapper / prefix
+        map_type: str = "mlp",
+        prefix_size: int = 192,          # matches sid embedding dim
+        sid_prefix_length: int = 40,
+        sid_prefix_length_clip: int = 40,
+        num_layers: int = 8,
+        norm_sid_emb: bool = False,
+        # LM
+        text_decoder: str = "gpt2",
+        tok_len: int = 67,
+        text_prefix_length: int = 10,
+        # audio
         sample_rate: int = 16000,
         **kwargs,
     ):
@@ -29,8 +48,15 @@ class CoLMboConfig(PretrainedConfig):
         self.n_mels = n_mels
         self.embedding_dim = embedding_dim
         self.channel = channel
-        self.prefix_length = prefix_length
-        self.gpt_model_name = gpt_model_name
+        self.map_type = map_type
+        self.prefix_size = prefix_size
+        self.sid_prefix_length = sid_prefix_length
+        self.sid_prefix_length_clip = sid_prefix_length_clip
+        self.num_layers = num_layers
+        self.norm_sid_emb = norm_sid_emb
+        self.text_decoder = text_decoder
+        self.tok_len = tok_len
+        self.text_prefix_length = text_prefix_length
         self.sample_rate = sample_rate
 
 
@@ -40,9 +66,14 @@ class CoLMboModel(PreTrainedModel):
     """
     CoLMbo: Speaker Language Model for Descriptive Profiling.
 
-    Usage:
-        model = AutoModel.from_pretrained("cmu-mlsp/CoLMbo", trust_remote_code=True)
-        output = model.describe_file("audio.wav", prompt="What is the speaker's gender?")
+    Architecture:
+        audio → Mel_Spectrogram → ECAPA encoder → sid_mapper → prefix tokens
+        prefix tokens + prompt tokens → GPT-2 LM → natural language description
+
+    Example:
+        >>> from transformers import AutoModel
+        >>> model = AutoModel.from_pretrained("cmu-mlsp/CoLMbo", trust_remote_code=True)
+        >>> print(model.describe_file("speaker.wav", "What is the speaker's dialect?"))
     """
 
     config_class = CoLMboConfig
@@ -51,30 +82,44 @@ class CoLMboModel(PreTrainedModel):
     def __init__(self, config: CoLMboConfig):
         super().__init__(config)
 
+        # Local imports — resolved from files shipped in the HF repo
         from encoder.encoder import Model
         from load_data.extract_fbanks import Mel_Spectrogram
+        from mapper import get_sid_mapper
 
+        # ── Audio frontend ────────────────────────────────────────────
         self.mel_extractor = Mel_Spectrogram()
 
-        # Speaker encoder
+        # ── Speaker encoder (ECAPA-TDNN) ──────────────────────────────
         self.sid_model = Model(
             n_mels=config.n_mels,
             embedding_dim=config.embedding_dim,
             channel=config.channel,
         )
 
-        # Mapper: linear projection from speaker embedding → LM token space
-        self.mapper = nn.Linear(config.embedding_dim, 768 * config.prefix_length)
+        # ── GPT-2 decoder ─────────────────────────────────────────────
+        self.gpt = GPT2LMHeadModel.from_pretrained(config.text_decoder)
+        self.gpt_embedding_size = self.gpt.transformer.wte.weight.shape[1]
+
+        # ── Speaker → prefix mapper ───────────────────────────────────
+        self.sid_mapper = get_sid_mapper(
+            config.map_type,
+            None,
+            config.prefix_size,
+            self.gpt_embedding_size,
+            config.sid_prefix_length,
+            config.sid_prefix_length_clip,
+            config.num_layers,
+        )
 
-        # GPT LM head
-        from transformers import GPT2LMHeadModel
-        self.gpt = GPT2LMHeadModel.from_pretrained(config.gpt_model_name)
+        # ── Tokenizer ───────────────────────────────────────���─────────
+        self.tokenizer = AutoTokenizer.from_pretrained(config.text_decoder)
+        self.tokenizer.add_special_tokens({'pad_token': '!'})
 
-        self.prefix_length = config.prefix_length
         self.post_init()
 
     # ------------------------------------------------------------------
-    # Forward — returns speaker embedding (pipeline compatibility)
+    # HF-standard forward (returns speaker embedding for pipeline compat)
     # ------------------------------------------------------------------
     def forward(self, input_values: torch.Tensor) -> BaseModelOutput:
         mel = self.mel_extractor(input_values)
@@ -82,50 +127,121 @@ class CoLMboModel(PreTrainedModel):
         return BaseModelOutput(last_hidden_state=spk_emb.unsqueeze(1))
 
     # ------------------------------------------------------------------
-    # Internal helpers
+    # Internal helpers — mirror ExpWrapper exactly
     # ------------------------------------------------------------------
-    def _get_sid_prefix(self, spk_emb: torch.Tensor) -> torch.Tensor:
-        batch = spk_emb.size(0)
-        prefix = self.mapper(spk_emb)
-        return prefix.view(batch, self.prefix_length, -1)
-
-    def _get_prompt_prefix(self, prompt: str, device):
-        from transformers import GPT2Tokenizer
-        tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-        tokens = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
-        return self.gpt.transformer.wte(tokens)
+    def _get_sid_prefix(self, sid_embeddings: torch.Tensor) -> torch.Tensor:
+        if self.config.norm_sid_emb:
+            sid_embeddings = sid_embeddings / sid_embeddings.norm(2, -1).reshape(-1, 1)
+        return (
+            self.sid_mapper(sid_embeddings)
+            .contiguous()
+            .view(-1, self.config.sid_prefix_length, self.gpt_embedding_size)
+        )
 
-    @torch.no_grad()
-    def _generate_beam(self, prefix_emb: torch.Tensor, num_beams: int = 5) -> list:
-        from transformers import GPT2Tokenizer
-        tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-        tokenizer.pad_token = tokenizer.eos_token
-        generated = self.gpt.generate(
-            inputs_embeds=prefix_emb,
-            max_new_tokens=100,
-            num_beams=num_beams,
-            early_stopping=True,
-            pad_token_id=tokenizer.eos_token_id,
+    def _preprocess_prompt_single(self, text: str, device) -> dict:
+        tok = self.tokenizer.encode_plus(
+            text=text,
+            add_special_tokens=True,
+            max_length=10,
+            pad_to_max_length=True,
+            return_tensors="pt",
+            truncation=True,
         )
-        return [tokenizer.decode(g, skip_special_tokens=True) for g in generated]
+        return {k: v.reshape(-1).to(device) for k, v in tok.items()}
+
+    def _get_prompt_prefix(self, text: str, device) -> torch.Tensor:
+        preprocessed = self._preprocess_prompt_single(text, device)
+        # Stack to [1, seq_len] then embed
+        input_ids = preprocessed["input_ids"].unsqueeze(0)
+        with torch.no_grad():
+            return self.gpt.transformer.wte(input_ids)   # [1, seq_len, 768]
+
+    def _generate_beam(
+        self,
+        prefix_emb: torch.Tensor,
+        beam_size: int = 1,
+        entry_length: int = 80,
+        temperature: float = 1.0,
+        stop_token: str = " <|endoftext|>",
+    ) -> list:
+        """Exact port of ExpWrapper.generate_beam."""
+        stop_token_index = self.tokenizer.encode(stop_token)[0]
+        tokens = None
+        scores = None
+        device = next(self.gpt.parameters()).device
+        seq_lengths = torch.ones(beam_size, device=device)
+        is_stopped = torch.zeros(beam_size, device=device, dtype=torch.bool)
+
+        with torch.no_grad():
+            generated = prefix_emb
+            for i in range(entry_length):
+                outputs = self.gpt(inputs_embeds=generated)
+                logits = outputs.logits
+                logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
+                logits = logits.softmax(-1).log()
+
+                if scores is None:
+                    scores, next_tokens = logits.topk(beam_size, -1)
+                    generated = generated.expand(beam_size, *generated.shape[1:])
+                    next_tokens, scores = next_tokens.permute(1, 0), scores.squeeze(0)
+                    tokens = next_tokens if tokens is None else torch.cat(
+                        (tokens.expand(beam_size, *tokens.shape[1:]), next_tokens), dim=1
+                    )
+                else:
+                    logits[is_stopped] = -float(np.inf)
+                    logits[is_stopped, 0] = 0
+                    scores_sum = scores[:, None] + logits
+                    seq_lengths[~is_stopped] += 1
+                    scores_sum_average = scores_sum / seq_lengths[:, None]
+                    scores_sum_average, next_tokens = scores_sum_average.view(-1).topk(beam_size, -1)
+                    next_tokens_source = next_tokens // scores_sum.shape[1]
+                    seq_lengths = seq_lengths[next_tokens_source]
+                    next_tokens = next_tokens % scores_sum.shape[1]
+                    next_tokens = next_tokens.unsqueeze(1)
+                    tokens = tokens[next_tokens_source]
+                    tokens = torch.cat((tokens, next_tokens), dim=1)
+                    generated = generated[next_tokens_source]
+                    scores = scores_sum_average * seq_lengths
+                    is_stopped = is_stopped[next_tokens_source]
+
+                next_token_embed = self.gpt.transformer.wte(
+                    next_tokens.squeeze()
+                ).view(generated.shape[0], 1, -1)
+                generated = torch.cat((generated, next_token_embed), dim=1)
+                is_stopped = is_stopped + next_tokens.eq(stop_token_index).squeeze()
+                if is_stopped.all():
+                    break
+
+        scores = scores / seq_lengths
+        output_list = tokens.cpu().numpy()
+        output_texts = [
+            self.tokenizer.decode(output[: int(length)])
+            for output, length in zip(output_list, seq_lengths)
+        ]
+        order = scores.argsort(descending=True)
+        return [output_texts[i] for i in order]
 
     # ------------------------------------------------------------------
-    # Main API
+    # Public API
     # ------------------------------------------------------------------
     @torch.no_grad()
     def describe(
         self,
         waveform: torch.Tensor,
         prompt: str = "Please describe the speaker.",
-        num_beams: int = 5,
+        beam_size: int = 1,
+        entry_length: int = 80,
+        temperature: float = 1.0,
     ) -> str:
         """
         Generate a natural language description of the speaker.
 
         Args:
-            waveform:  raw audio [1, T] at 16 kHz
-            prompt:    e.g. "What is the speaker's dialect?"
-            num_beams: beam search width
+            waveform:     raw audio tensor [1, T] at 16 kHz
+            prompt:       e.g. "What is the speaker's dialect?"
+            beam_size:    beam search width (default 1 = greedy)
+            entry_length: max tokens to generate
+            temperature:  sampling temperature
 
         Returns:
             str: generated description
@@ -135,16 +251,22 @@ class CoLMboModel(PreTrainedModel):
             >>> waveform, sr = torchaudio.load("audio.wav")
             >>> print(model.describe(waveform, "What is the speaker's age?"))
         """
-        device = next(self.parameters()).device
+        device = next(self.gpt.parameters()).device
         self.eval()
 
-        mel     = self.mel_extractor(waveform).to(device)
-        spk_emb = self.sid_model(mel)
-        sid_pfx = self._get_sid_prefix(spk_emb)
-        pmt_pfx = self._get_prompt_prefix(prompt, device)
-        prefix  = torch.cat((sid_pfx, pmt_pfx), dim=1)
+        mel         = self.mel_extractor(waveform).to(device)
+        spk_emb     = self.sid_model(mel)
+        sids_prefix = self._get_sid_prefix(spk_emb)                      # [1, sid_prefix_len, 768]
+        pmt_prefix  = self._get_prompt_prefix(prompt, device)             # [1, tok_len, 768]
+        prefix_emb  = torch.cat((sids_prefix, pmt_prefix), dim=1)        # [1, total_len, 768]
 
-        return self._generate_beam(prefix, num_beams=num_beams)[0]
+        texts = self._generate_beam(
+            prefix_emb,
+            beam_size=beam_size,
+            entry_length=entry_length,
+            temperature=temperature,
+        )
+        return texts[0]
 
     @torch.no_grad()
     def describe_file(

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e0d80efbeffb56f4038bf9d320d15b5377d12b1cb85833e908d9f0f6b5c2bbab
-size 2066033810
+oid sha256:54b52bd0b2c80e0afcddeebf6c30ce4d9645c265b546cdc95c2cf36ba7564b3f
+size 1982720694