geevec-ai

/

geevec-embeddings-1.0-lite

Like 4

Follow

geevec-ai 2

Please try transformers==4.57.1. You can refer to this: https://github.com/embeddings-benchmark/results/pull/530#issuecomment-4419041562.

@lnxtree that's a lame response, everyone is eventually moving to Transformers 5+

@christian-winkler-th try using this replacement code for "modeling_qwen3_pseudo_moe.py" and I'll do a PR on this HF repo as well, which they should be grateful for...

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, Tuple, Union

from transformers import AutoModel, AutoConfig
from transformers import Qwen3Config, Qwen3PreTrainedModel, Qwen3Model
from transformers.modeling_outputs import BaseModelOutputWithPast


# ============================================================================
#  Config
# ============================================================================

class Qwen3PseudoMoEConfig(Qwen3Config):
    model_type = "qwen3_pseudo_moe"

    def __init__(
        self,
        proj_dim: int = 4096,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        lora_target_modules: list = None,
        domain_names: list = None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.proj_dim = proj_dim
        self.lora_rank = lora_rank
        self.lora_alpha = lora_alpha
        self.lora_target_modules = lora_target_modules or [
            "q_proj", "k_proj", "v_proj", "o_proj",
            "gate_proj", "up_proj", "down_proj",
            "proj_linear",
        ]
        self.domain_names = domain_names or ["coding", "reasoning"]


# ============================================================================
#  MultiLoRALinear
# ============================================================================

class MultiLoRALinear(nn.Module):

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = False,
        domain_names: list = None,
        rank: int = 16,
        alpha: float = 32.0,
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = rank
        self.scaling = alpha / rank
        self.active_domain: Optional[str] = None

        self.weight = nn.Parameter(torch.empty(out_features, in_features))
        if bias:
            self.bias = nn.Parameter(torch.empty(out_features))
        else:
            self.register_parameter("bias", None)

        self.lora_A = nn.ModuleDict()
        self.lora_B = nn.ModuleDict()
        for name in (domain_names or []):
            self.lora_A[name] = nn.Linear(in_features, rank, bias=False)
            self.lora_B[name] = nn.Linear(rank, out_features, bias=False)
            nn.init.zeros_(self.lora_B[name].weight)
            nn.init.kaiming_uniform_(self.lora_A[name].weight)

        nn.init.kaiming_uniform_(self.weight)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = F.linear(x, self.weight, self.bias)
        if self.active_domain and self.active_domain != "general":
            if self.active_domain in self.lora_A:
                lora_out = self.lora_B[self.active_domain](
                    self.lora_A[self.active_domain](x)
                )
                out = out + lora_out * self.scaling
        return out


# ============================================================================
#  Qwen3PseudoMoEModel
# ============================================================================

class Qwen3PseudoMoEModelModel(Qwen3PreTrainedModel):
    config_class = Qwen3PseudoMoEConfig
    _no_split_modules = ["Qwen3DecoderLayer"]
    # Nothing is tied (the model has no LM head). None is the PreTrainedModel default in every Transformers
    # version; the former [] crashes Transformers 5, which expects a dict here when tie_word_embeddings is true.
    _tied_weights_keys = None

    def __init__(self, config: Qwen3PseudoMoEConfig):
        super().__init__(config)

        self._current_domain = "general"

        self.model = Qwen3Model(config)

        self.proj_linear = nn.Linear(config.hidden_size, config.proj_dim, bias=False)

        self._inject_lora(config)

        self.post_init()

    def _inject_lora(self, config: Qwen3PseudoMoEConfig):

        target_set = set(config.lora_target_modules)

        to_replace = []
        for name, module in self.named_modules():
            if type(module) is not nn.Linear:
                continue
            short_name = name.split(".")[-1]
            if short_name not in target_set:
                continue
            parts = name.split(".")
            parent = self
            for p in parts[:-1]:
                parent = getattr(parent, p)
            if isinstance(parent, MultiLoRALinear):
                continue
            to_replace.append((name, module, parent, parts[-1]))

        for name, module, parent, attr_name in to_replace:
            multi_lora = MultiLoRALinear(
                in_features=module.in_features,
                out_features=module.out_features,
                bias=module.bias is not None,
                domain_names=config.domain_names,
                rank=config.lora_rank,
                alpha=config.lora_alpha,
            )
            multi_lora.weight = module.weight
            if module.bias is not None:
                multi_lora.bias = module.bias
            setattr(parent, attr_name, multi_lora)

    @property
    def domain(self) -> str:
        return self._current_domain

    def set_domain(self, domain: str):
        valid = {"general"} | set(self.config.domain_names)
        assert domain in valid, f"domain must be one of {valid}, got '{domain}'"
        self._current_domain = domain
        for module in self.modules():
            if isinstance(module, MultiLoRALinear):
                module.active_domain = domain

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values=None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        domain: Optional[str] = None,
        **kwargs,
    ) -> Union[Tuple, BaseModelOutputWithPast]:
        if domain is not None:
            self.set_domain(domain)

        # Same result as config.use_return_dict, which Transformers 5 deprecates (5.x also drops torchscript).
        return_dict = return_dict if return_dict is not None else (
            getattr(self.config, "return_dict", True) and not getattr(self.config, "torchscript", False)
        )

        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=True,
            **kwargs,
        )

        hidden_states = outputs.last_hidden_state
        projected = self.proj_linear(hidden_states)

        if not return_dict:
            return (projected,) + tuple(
                v for v in [
                    outputs.past_key_values,
                    outputs.hidden_states,
                    outputs.attentions,
                ] if v is not None
            )

        return BaseModelOutputWithPast(
            last_hidden_state=projected,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    def encode(
        self,
        input_ids: torch.LongTensor,
        attention_mask: torch.Tensor,
        domain: str = "general",
    ) -> torch.Tensor:
        self.set_domain(domain)
        outputs = self.forward(
            input_ids=input_ids,
            attention_mask=attention_mask,
            return_dict=True,
        )
        seq_lengths = attention_mask.sum(dim=1) - 1
        batch_idx = torch.arange(input_ids.size(0), device=input_ids.device)
        embeddings = outputs.last_hidden_state[batch_idx, seq_lengths]
        return F.normalize(embeddings, p=2, dim=-1)


AutoConfig.register("qwen3_pseudo_moe", Qwen3PseudoMoEConfig)
AutoModel.register(Qwen3PseudoMoEConfig, Qwen3PseudoMoEModelModel)

See my PR here: https://huggingface.co/geevec-ai/geevec-embeddings-1.0-lite/discussions/2

Correction, this should be the good one: https://huggingface.co/geevec-ai/geevec-embeddings-1.0-lite/discussions/4