model.py

# model.py

from __future__ import annotations

from typing import Optional, Union, Tuple, Dict, Any

import torch
import torch.nn as nn
from torch.nn import CrossEntropyLoss

from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    PretrainedConfig,
    PreTrainedModel,
    GenerationMixin,
)
from transformers.modeling_outputs import CausalLMOutputWithPast


class SpeakerProjector(nn.Module):
    def __init__(
        self,
        input_dim: int,
        output_dim: int,
        hidden_dim: Optional[int] = None,
        dropout: float = 0.0,
    ):
        super().__init__()

        if hidden_dim is None:
            self.net = nn.Sequential(
                nn.Linear(input_dim, output_dim),
                nn.Tanh(),
            )
        else:
            self.net = nn.Sequential(
                nn.Linear(input_dim, hidden_dim),
                nn.GELU(),
                nn.Dropout(dropout),
                nn.Linear(hidden_dim, output_dim),
                nn.Tanh(),
            )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class SpeakerConditionedCausalLMConfig(PretrainedConfig):
    model_type = "speaker_conditioned_causal_lm"

    def __init__(
        self,
        base_model_name_or_path: Optional[str] = None,
        speaker_embedding_dim: int = 512,
        speaker_hidden_dim: Optional[int] = None,
        speaker_dropout: float = 0.0,
        speaker_token_id: Optional[int] = None,
        freeze_base_model: bool = True,
        **kwargs,
    ):
        super().__init__(**kwargs)

        self.base_model_name_or_path = base_model_name_or_path
        self.speaker_embedding_dim = speaker_embedding_dim
        self.speaker_hidden_dim = speaker_hidden_dim
        self.speaker_dropout = speaker_dropout
        self.speaker_token_id = speaker_token_id
        self.freeze_base_model = freeze_base_model


class SpeakerConditionedCausalLM(PreTrainedModel, GenerationMixin):
    config_class = SpeakerConditionedCausalLMConfig
    base_model_prefix = "model"
    main_input_name = "input_ids"
    _no_split_modules = []

    def __init__(
        self,
        config: SpeakerConditionedCausalLMConfig,
        base_lm: Optional[PreTrainedModel] = None,
    ):
        super().__init__(config)

        if config.base_model_name_or_path is None and base_lm is None:
            raise ValueError(
                "You must provide either config.base_model_name_or_path or a preloaded base_lm."
            )

        if base_lm is None:
            self.model = AutoModelForCausalLM.from_pretrained(
                config.base_model_name_or_path
            )
        else:
            self.model = base_lm

        hidden_size = self.model.config.hidden_size

        self.speaker_projector = SpeakerProjector(
            input_dim=config.speaker_embedding_dim,
            output_dim=hidden_size,
            hidden_dim=config.speaker_hidden_dim,
            dropout=config.speaker_dropout,
        )

        # Make sure the wrapper config knows useful generation/training ids.
        for attr in [
            "pad_token_id",
            "bos_token_id",
            "eos_token_id",
            "vocab_size",
            "tie_word_embeddings",
        ]:
            if hasattr(self.model.config, attr):
                setattr(self.config, attr, getattr(self.model.config, attr))

        if config.freeze_base_model:
            self.freeze_base_model()

        # Important for save/load consistency
        self.post_init()

    def freeze_base_model(self) -> None:
        for param in self.model.parameters():
            param.requires_grad = False

    def unfreeze_base_model(self) -> None:
        for param in self.model.parameters():
            param.requires_grad = True

    def get_input_embeddings(self):
        return self.model.get_input_embeddings()

    def set_input_embeddings(self, value):
        self.model.set_input_embeddings(value)

    def get_output_embeddings(self):
        return self.model.get_output_embeddings()

    def set_output_embeddings(self, new_embeddings):
        self.model.set_output_embeddings(new_embeddings)

    def resize_token_embeddings(self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None):
        return self.model.resize_token_embeddings(
            new_num_tokens=new_num_tokens,
            pad_to_multiple_of=pad_to_multiple_of,
        )

    def _inject_speaker_embeddings(
        self,
        input_ids: torch.LongTensor,
        inputs_embeds: torch.FloatTensor,
        speaker_embedding: Optional[torch.FloatTensor],
    ) -> torch.FloatTensor:
        """
        Replace the embedding at speaker_token_id with projected speaker_embedding.

        input_ids:        [B, T]
        inputs_embeds:    [B, T, H]
        speaker_embedding:[B, D_s]
        """
        if speaker_embedding is None:
            return inputs_embeds

        if self.config.speaker_token_id is None:
            raise ValueError("config.speaker_token_id must be set.")

        # Find all marker positions
        speaker_mask = input_ids.eq(self.config.speaker_token_id)  # [B, T]

        if not speaker_mask.any():
            # During cached generation steps this is expected.
            return inputs_embeds

        # Project speaker vectors to LM hidden size
        projected = self.speaker_projector(speaker_embedding)  # [B, H]
        projected = projected.to(dtype=inputs_embeds.dtype, device=inputs_embeds.device)

        # Clone so we do not do in-place ops on shared embeddings
        inputs_embeds = inputs_embeds.clone()

        # Expect exactly one speaker marker per sequence.
        num_markers_per_seq = speaker_mask.sum(dim=1)
        if not torch.all(num_markers_per_seq == 1):
            raise ValueError(
                "Each sequence must contain exactly one speaker token marker. "
                f"Got counts: {num_markers_per_seq.tolist()}"
            )

        batch_idx, time_idx = torch.where(speaker_mask)
        inputs_embeds[batch_idx, time_idx] = projected[batch_idx]

        return inputs_embeds

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.Tensor, ...], ...]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        speaker_embedding: Optional[torch.FloatTensor] = None,
        **kwargs,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        """
        Standard CausalLM forward, but inject projected speaker_embedding at
        the <|SPEAKER_TOKEN_POS|> position.

        Notes:
        - input_ids is required whenever speaker injection is needed.
        - during generation with cache, later decoding steps may not include the
          speaker token anymore; that is fine because the prompt cache already
          contains the injected information from the first step.
        """
        return_dict = return_dict if return_dict is not None else self.config.return_dict

        if inputs_embeds is not None and input_ids is None:
            raise ValueError(
                "This wrapper currently requires input_ids whenever inputs_embeds is passed, "
                "because speaker injection is located via input_ids."
            )

        if inputs_embeds is None:
            if input_ids is None:
                raise ValueError("You must provide input_ids or inputs_embeds.")
            inputs_embeds = self.get_input_embeddings()(input_ids)

        if input_ids is not None:
            inputs_embeds = self._inject_speaker_embeddings(
                input_ids=input_ids,
                inputs_embeds=inputs_embeds,
                speaker_embedding=speaker_embedding,
            )

        outputs = self.model(
            input_ids=None,  # we pass inputs_embeds after modification
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            labels=labels,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            **kwargs,
        )

        return outputs

    def prepare_inputs_for_generation(
        self,
        input_ids: torch.LongTensor,
        past_key_values: Optional[Tuple[Tuple[torch.Tensor, ...], ...]] = None,
        attention_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        speaker_embedding: Optional[torch.FloatTensor] = None,
        **kwargs,
    ) -> Dict[str, Any]:
        """
        Called internally by .generate().

        We must preserve speaker_embedding across generation steps.
        On the first step, the full prompt is passed and the speaker marker exists.
        On later steps with cache, usually only the newest token is passed.
        """
        if past_key_values is not None:
            # Standard generation behavior: only last token is fed after cache exists
            input_ids = input_ids[:, -1:]

        model_inputs = {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "past_key_values": past_key_values,
            "use_cache": kwargs.get("use_cache", True),
            "speaker_embedding": speaker_embedding,
        }

        return model_inputs

    @classmethod
    def from_pretrained_base(
        cls,
        base_model_name_or_path: str,
        speaker_embedding_dim: int,
        speaker_token_id: int,
        speaker_hidden_dim: Optional[int] = None,
        speaker_dropout: float = 0.0,
        freeze_base_model: bool = True,
        **kwargs,
    ) -> "SpeakerConditionedCausalLM":
        """
        Convenience constructor for creating a fresh wrapper around a pretrained LM.

        Example:
            model = SpeakerConditionedCausalLM.from_pretrained_base(
                "meta-llama/Llama-3.2-1B-Instruct",
                speaker_embedding_dim=256,
                speaker_token_id=tokenizer.convert_tokens_to_ids("<|SPEAKER_TOKEN_POS|>")
            )
        """
        config = SpeakerConditionedCausalLMConfig(
            base_model_name_or_path=base_model_name_or_path,
            speaker_embedding_dim=speaker_embedding_dim,
            speaker_hidden_dim=speaker_hidden_dim,
            speaker_dropout=speaker_dropout,
            speaker_token_id=speaker_token_id,
            freeze_base_model=freeze_base_model,
        )
        return cls(config=config, **kwargs)

    def save_pretrained(self, save_directory: str, **kwargs):
        """
        Save:
        - wrapper config
        - wrapper state dict (includes projector + underlying model weights as currently attached)

        If base model is frozen and unchanged, this is still fine; checkpoint size will include it.
        If you want smaller checkpoints, I mention a lighter alternative below.
        """
        return super().save_pretrained(save_directory, **kwargs)
        

# Optional: register with Auto classes in-process
AutoConfig.register(SpeakerConditionedCausalLMConfig.model_type, SpeakerConditionedCausalLMConfig)
AutoModelForCausalLM.register(SpeakerConditionedCausalLMConfig, SpeakerConditionedCausalLM)