TuKoResearch
/

GSLM-HuBERT200

+"""
+GSLM Model Configuration
+"""
+import json
+import os
+from typing import Optional
+class GSLMConfig:
+    """
+    Configuration class for GSLM (Generative Spoken Language Model).
+    This configuration class stores all parameters needed to initialize a GSLMModel.
+    """
+    model_type = "gslm"
+    def __init__(
+        self,
+        vocab_size: int = 204,
+        d_model: int = 1024,
+        nhead: int = 16,
+        num_layers: int = 12,
+        dim_feedforward: int = 4096,
+        dropout: float = 0.1,
+        attention_dropout: float = 0.1,
+        max_seq_length: int = 3072,
+        pad_idx: int = 0,
+        share_input_output_embed: bool = True,
+        activation: str = "relu",
+        architecture: str = "transformer_lm_big",
+        **kwargs
+    ):
+        """
+        Initialize GSLM configuration.
+        Args:
+            vocab_size: Size of the vocabulary
+            d_model: Dimensionality of the embeddings and hidden states
+            nhead: Number of attention heads
+            num_layers: Number of transformer layers
+            dim_feedforward: Dimensionality of the feedforward network
+            dropout: Dropout probability
+            attention_dropout: Dropout probability for attention weights
+            max_seq_length: Maximum sequence length
+            pad_idx: Padding token index
+            share_input_output_embed: Whether to share input and output embeddings
+            activation: Activation function ("relu" or "gelu")
+            architecture: Model architecture name
+        """
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.nhead = nhead
+        self.num_layers = num_layers
+        self.dim_feedforward = dim_feedforward
+        self.dropout = dropout
+        self.attention_dropout = attention_dropout
+        self.max_seq_length = max_seq_length
+        self.pad_idx = pad_idx
+        self.share_input_output_embed = share_input_output_embed
+        self.activation = activation
+        self.architecture = architecture
+        # Handle any extra kwargs
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+    def to_dict(self):
+        """Convert configuration to dictionary."""
+        output = {}
+        for key, value in self.__dict__.items():
+            if not key.startswith('_'):
+                output[key] = value
+        output['model_type'] = self.model_type
+        return output
+    def to_json_string(self):
+        """Convert configuration to JSON string."""
+        return json.dumps(self.to_dict(), indent=2, sort_keys=True)
+    def save_pretrained(self, save_directory: str):
+        """Save configuration to directory."""
+        if not os.path.exists(save_directory):
+            os.makedirs(save_directory)
+        config_file = os.path.join(save_directory, "config.json")
+        with open(config_file, 'w') as f:
+            f.write(self.to_json_string())
+    @classmethod
+    def from_dict(cls, config_dict: dict):
+        """Create configuration from dictionary."""
+        return cls(**config_dict)
+    @classmethod
+    def from_json_file(cls, json_file: str):
+        """Create configuration from JSON file."""
+        with open(json_file, 'r') as f:
+            config_dict = json.load(f)
+        return cls.from_dict(config_dict)
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
+        """
+        Load configuration from pretrained model.
+        Args:
+            pretrained_model_name_or_path: Path to pretrained model or model identifier
+            **kwargs: Additional configuration parameters to override
+        Returns:
+            GSLMConfig instance
+        """
+        if os.path.isdir(pretrained_model_name_or_path):
+            config_file = os.path.join(pretrained_model_name_or_path, "config.json")
+        else:
+            config_file = pretrained_model_name_or_path
+        # Load config from file
+        config = cls.from_json_file(config_file)
+        # Override with any provided kwargs
+        for key, value in kwargs.items():
+            setattr(config, key, value)
+        return config
+    def __repr__(self):
+        return f"{self.__class__.__name__} {self.to_json_string()}"

modeling_gslm_ulm.py DELETED Viewed

@@ -1,457 +0,0 @@
-"""
-GSLM ULM model definition for HuggingFace.
-"""
-import math
-import torch
-import torch.nn as nn
-from torch.nn import functional as F
-from typing import Optional, Tuple, Dict, List, Union
-from transformers import PreTrainedModel
-from transformers.modeling_outputs import BaseModelOutput, CausalLMOutput
-from .configuration_gslm_ulm import GSLMULMConfig
-class GSLMULM(PreTrainedModel):
-    """
-    GSLM Unit Language Model - Transformer Language Model for discrete speech units.
-    """
-    config_class = GSLMULMConfig
-    base_model_prefix = "transformer"
-    supports_gradient_checkpointing = True
-    def __init__(self, config):
-        super().__init__(config)
-        self.config = config
-        self.transformer = nn.ModuleDict(dict(
-            wte = nn.Embedding(config.vocab_size, config.n_embd),
-            drop = nn.Dropout(config.dropout),
-            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
-            ln_f = RMSNorm(config.n_embd, bias=config.bias),
-        ))
-        # Sinusoidal positional encoding
-        if getattr(config, 'use_sinusoidal_embeddings', True):
-            self.pos_encoder = SinusoidalPositionalEncoding(
-                config.n_embd,
-                config.max_position_embeddings
-            )
-        else:
-            self.transformer.wpe = nn.Embedding(config.max_position_embeddings, config.n_embd)
-            self.pos_encoder = None
-        # Language modeling head
-        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
-        # Weight tying
-        self.lm_head.weight = self.transformer.wte.weight
-        # Initialize weights
-        self.post_init()
-    def get_input_embeddings(self):
-        return self.transformer.wte
-    def set_input_embeddings(self, new_embeddings):
-        self.transformer.wte = new_embeddings
-    def get_output_embeddings(self):
-        return self.lm_head
-    def set_output_embeddings(self, new_embeddings):
-        self.lm_head = new_embeddings
-    def _init_weights(self, module):
-        if isinstance(module, nn.Linear):
-            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
-            if module.bias is not None:
-                torch.nn.init.zeros_(module.bias)
-        elif isinstance(module, nn.Embedding):
-            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
-            if hasattr(module, 'padding_idx') and module.padding_idx is not None:
-                module.weight.data[module.padding_idx].zero_()
-        # Special scaled init for residual projections
-        if hasattr(module, "c_proj") and isinstance(module.c_proj, nn.Linear):
-            torch.nn.init.normal_(module.c_proj.weight, mean=0.0, std=0.02/math.sqrt(2 * self.config.n_layer))
-    def forward(
-        self,
-        input_ids: Optional[torch.LongTensor] = None,
-        attention_mask: Optional[torch.FloatTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        **kwargs
-    ) -> Union[Tuple, CausalLMOutput]:
-        """
-        Forward pass of the model.
-        Args:
-            input_ids: Input token IDs of shape (batch_size, sequence_length)
-            attention_mask: Attention mask (not used, kept for compatibility)
-            labels: Labels for language modeling loss
-            output_attentions: Whether to return attention weights
-            output_hidden_states: Whether to return hidden states
-            return_dict: Whether to return a dictionary
-        Returns:
-            CausalLMOutput or tuple
-        """
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        # Get embeddings
-        tok_emb = self.transformer.wte(input_ids)
-        tok_emb = tok_emb * math.sqrt(self.config.n_embd)  # Scale embeddings
-        # Add positional encoding
-        if self.pos_encoder is not None:
-            x = self.pos_encoder(tok_emb)
-        else:
-            pos = torch.arange(0, input_ids.size(1), dtype=torch.long, device=input_ids.device)
-            pos_emb = self.transformer.wpe(pos)
-            x = tok_emb + pos_emb
-        x = self.transformer.drop(x)
-        # Pass through transformer blocks
-        all_hidden_states = () if output_hidden_states else None
-        all_attentions = () if output_attentions else None
-        for i, block in enumerate(self.transformer.h):
-            if output_hidden_states:
-                all_hidden_states = all_hidden_states + (x,)
-            outputs = block(x, output_attentions=output_attentions)
-            x = outputs[0]
-            if output_attentions:
-                all_attentions = all_attentions + (outputs[1],)
-        # Final layer norm
-        x = self.transformer.ln_f(x)
-        if output_hidden_states:
-            all_hidden_states = all_hidden_states + (x,)
-        # Language modeling head
-        logits = self.lm_head(x)
-        loss = None
-        if labels is not None:
-            # Shift so that tokens < n predict n
-            shift_logits = logits[..., :-1, :].contiguous()
-            shift_labels = labels[..., 1:].contiguous()
-            # Flatten the tokens
-            loss_fct = nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id)
-            loss = loss_fct(
-                shift_logits.view(-1, shift_logits.size(-1)),
-                shift_labels.view(-1)
-            )
-        if not return_dict:
-            output = (logits,) + (all_hidden_states,) if output_hidden_states else ()
-            output = output + (all_attentions,) if output_attentions else ()
-            return ((loss,) + output) if loss is not None else output
-        return CausalLMOutput(
-            loss=loss,
-            logits=logits,
-            hidden_states=all_hidden_states,
-            attentions=all_attentions,
-        )
-    def sample_logits(
-        self,
-        logits: torch.FloatTensor,
-        temperature: float = 1.0,
-        top_k: Optional[int] = None,
-        top_p: Optional[float] = None
-    ) -> torch.LongTensor:
-        """
-        Sample from logits with temperature, top-k, and top-p (nucleus) sampling.
-        """
-        # If temperature is 0.0, use argmax
-        if temperature == 0.0:
-            return torch.argmax(logits, dim=-1)
-        # Apply temperature
-        logits = logits / temperature
-        # Apply top-k filtering if specified
-        if top_k is not None:
-            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
-            logits[logits < v[..., [-1]]] = -float('Inf')
-        # Apply top-p (nucleus) filtering if specified
-        if top_p is not None:
-            sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
-            sorted_probs = F.softmax(sorted_logits, dim=-1)
-            cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
-            # Remove tokens with cumulative probability above the threshold
-            sorted_indices_to_remove = cumulative_probs > top_p
-            sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
-            sorted_indices_to_remove[..., 0] = 0
-            indices_to_remove = sorted_indices_to_remove.scatter(
-                dim=-1, index=sorted_indices, src=sorted_indices_to_remove
-            )
-            logits[indices_to_remove] = -float('Inf')
-        # Compute softmax probabilities
-        probs = F.softmax(logits, dim=-1)
-        # Sample from the distribution
-        flat_probs = probs.view(-1, probs.size(-1))
-        sampled = torch.multinomial(flat_probs, num_samples=1)
-        sampled = sampled.view(*logits.shape[:-1])
-        return sampled
-    @torch.no_grad()
-    def generate(
-        self,
-        input_ids: torch.LongTensor,
-        max_new_tokens: int = 100,
-        temperature: float = 1.0,
-        top_k: Optional[int] = None,
-        top_p: Optional[float] = None,
-        eos_token_id: Optional[int] = None,
-        pad_token_id: Optional[int] = None,
-        attention_mask: Optional[torch.FloatTensor] = None,
-        **kwargs
-    ) -> torch.LongTensor:
-        """
-        Generate sequences autoregressively.
-        Args:
-            input_ids: Input token IDs of shape (batch_size, sequence_length)
-            max_new_tokens: Maximum number of tokens to generate
-            temperature: Sampling temperature
-            top_k: Top-k filtering parameter
-            top_p: Top-p (nucleus) filtering parameter
-            eos_token_id: End-of-sequence token ID
-            pad_token_id: Padding token ID
-            attention_mask: Attention mask (not used, kept for compatibility)
-        Returns:
-            Generated token IDs
-        """
-        if eos_token_id is None:
-            eos_token_id = self.config.eos_token_id
-        if pad_token_id is None:
-            pad_token_id = self.config.pad_token_id
-        batch_size = input_ids.shape[0]
-        device = input_ids.device
-        # Keep track of which sequences are done
-        unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=device)
-        # Cache for key-value pairs (more efficient generation)
-        past_key_values = None
-        for _ in range(max_new_tokens):
-            # Forward pass
-            outputs = self.forward(input_ids)
-            next_token_logits = outputs.logits[:, -1, :]
-            # Sample next tokens
-            next_tokens = self.sample_logits(
-                next_token_logits,
-                temperature=temperature,
-                top_k=top_k,
-                top_p=top_p
-            )
-            # Update sequences
-            if eos_token_id is not None:
-                tokens_to_add = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences)
-                unfinished_sequences = unfinished_sequences * (next_tokens != eos_token_id).long()
-            else:
-                tokens_to_add = next_tokens
-            # Concatenate tokens
-            input_ids = torch.cat([input_ids, tokens_to_add.unsqueeze(-1)], dim=-1)
-            # Stop if all sequences are finished
-            if eos_token_id is not None and unfinished_sequences.sum() == 0:
-                break
-        return input_ids
-    def prepare_inputs_for_generation(
-        self,
-        input_ids,
-        past_key_values=None,
-        attention_mask=None,
-        **kwargs
-    ):
-        """Prepare inputs for generation."""
-        # If past_key_values is used, only use the last token
-        if past_key_values:
-            input_ids = input_ids[:, -1:]
-        return {
-            "input_ids": input_ids,
-            "past_key_values": past_key_values,
-            "attention_mask": attention_mask,
-        }
-class SinusoidalPositionalEncoding(nn.Module):
-    """Sinusoidal positional encoding."""
-    def __init__(self, d_model: int, max_len: int = 5000):
-        super().__init__()
-        pe = torch.zeros(max_len, d_model)
-        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
-        div_term = torch.exp(torch.arange(0, d_model, 2).float() *
-                           (-math.log(10000.0) / d_model))
-        pe[:, 0::2] = torch.sin(position * div_term)
-        pe[:, 1::2] = torch.cos(position * div_term)
-        self.register_buffer('pe', pe.unsqueeze(0))
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Add positional encoding to input tensor."""
-        return x + self.pe[:, :x.size(1)]
-class Block(nn.Module):
-    """Transformer block with pre-normalization."""
-    def __init__(self, config):
-        super().__init__()
-        self.attn = CausalSelfAttention(config)
-        self.mlp = MLP(config)
-        self.norm1 = RMSNorm(config.n_embd, bias=config.bias)
-        self.norm2 = RMSNorm(config.n_embd, bias=config.bias)
-        self.config = config
-    def forward(self, x, output_attentions=False):
-        # Pre-norm attention block
-        attn_output = self.attn(self.norm1(x), output_attentions=output_attentions)
-        if output_attentions:
-            attn_output, attn_weights = attn_output
-            x = x + attn_output
-            x = x + self.mlp(self.norm2(x))
-            return x, attn_weights
-        else:
-            x = x + attn_output
-            x = x + self.mlp(self.norm2(x))
-            return (x,)
-class CausalSelfAttention(nn.Module):
-    """Multi-head causal self-attention."""
-    def __init__(self, config):
-        super().__init__()
-        assert config.n_embd % config.n_head == 0
-        self.n_head = config.n_head
-        self.n_embd = config.n_embd
-        self.head_dim = config.n_embd // config.n_head
-        self.scaling = self.head_dim ** -0.5
-        # Key, query, value projections for all heads
-        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
-        # Output projection
-        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
-        # Attention dropout
-        self.attn_dropout = nn.Dropout(config.attention_dropout)
-        self.resid_dropout = nn.Dropout(config.dropout)
-    def forward(self, x, output_attentions=False):
-        B, T, C = x.size()  # batch size, sequence length, embedding dimensionality
-        # Calculate query, key, values for all heads
-        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
-        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)  # (B, nh, T, hs)
-        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)  # (B, nh, T, hs)
-        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)  # (B, nh, T, hs)
-        # Scale query
-        q = q * self.scaling
-        # Causal self-attention
-        if not output_attentions:
-            # Use flash attention when available
-            y = F.scaled_dot_product_attention(
-                q, k, v,
-                attn_mask=None,
-                dropout_p=self.attn_dropout.p if self.training else 0.0,
-                is_causal=True
-            )
-        else:
-            # Manual implementation for attention weights
-            att = torch.matmul(q, k.transpose(-2, -1))  # (B, nh, T, T)
-            # Causal mask
-            causal_mask = torch.triu(
-                torch.ones(T, T, dtype=torch.bool, device=x.device),
-                diagonal=1
-            )
-            att = att.masked_fill(causal_mask, float('-inf'))
-            # Softmax
-            att = F.softmax(att, dim=-1, dtype=torch.float32).to(q.dtype)
-            att = self.attn_dropout(att)
-            # Apply attention to values
-            y = torch.matmul(att, v)  # (B, nh, T, hs)
-        # Re-assemble all head outputs
-        y = y.transpose(1, 2).contiguous().view(B, T, C)
-        # Output projection
-        y = self.resid_dropout(self.c_proj(y))
-        if output_attentions:
-            return y, att
-        else:
-            return y
-class MLP(nn.Module):
-    """Position-wise feed-forward network."""
-    def __init__(self, config):
-        super().__init__()
-        self.c_fc = nn.Linear(config.n_embd, config.n_inner, bias=config.bias)
-        self.c_proj = nn.Linear(config.n_inner, config.n_embd, bias=config.bias)
-        self.act = nn.GELU()
-        self.dropout = nn.Dropout(config.dropout)
-    def forward(self, x):
-        x = self.c_fc(x)
-        x = self.act(x)
-        x = self.c_proj(x)
-        x = self.dropout(x)
-        return x
-class RMSNorm(nn.Module):
-    """Root Mean Square Layer Normalization."""
-    def __init__(self, dim: int, eps: float = 1e-6, bias: bool = False):
-        super().__init__()
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(dim))
-        self.bias = nn.Parameter(torch.zeros(dim)) if bias else None
-    def forward(self, x):
-        norm_x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
-        if self.bias is not None:
-            return self.weight * norm_x + self.bias
-        return self.weight * norm_x
-# Register the model
-from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
-AutoConfig.register("gslm_ulm", GSLMULMConfig)
-AutoModel.register(GSLMULMConfig, GSLMULM)
-AutoModelForCausalLM.register(GSLMULMConfig, GSLMULM)