Upload modeling_hf.py

0_CollinsSTWrapper/modeling_hf.py

@@ -0,0 +1,212 @@
+"""
+Collins-RoPE 极简 Embedding 模型（HuggingFace 原生实现）
+架构：Hash Embedding (2-Universal + Sign Hash) -> RoPE -> Transformer Encoder -> Mean Pooling
+目标参数量：~2M
+"""
+
+import math
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput
+
+
+class CollinsConfig(PretrainedConfig):
+    model_type = "collins"
+
+    def __init__(
+        self,
+        vocab_size: int = 30522,
+        num_buckets: int = 2048,
+        hidden_size: int = 256,
+        num_hidden_layers: int = 3,
+        num_attention_heads: int = 8,
+        intermediate_size: int = 1024,
+        hidden_dropout_prob: float = 0.1,
+        attention_probs_dropout_prob: float = 0.1,
+        max_position_embeddings: int = 512,
+        # 2-Universal Hash 固定种子（保证 load 后哈希一致）
+        hash_seed: int = 42,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.num_buckets = num_buckets
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.max_position_embeddings = max_position_embeddings
+        self.hash_seed = hash_seed
+
+
+class CollinsHashEmbedding(nn.Module):
+    """
+    2-Universal Hash + Sign Hash 压缩 Embedding。
+    哈希参数从 config.hash_seed 确定性生成，保证 save/load 后一致。
+    """
+
+    def __init__(self, config: CollinsConfig):
+        super().__init__()
+        self.num_buckets = config.num_buckets
+        self.hidden_size = config.hidden_size
+
+        self.hash_table = nn.Parameter(
+            torch.randn(config.num_buckets, config.hidden_size)
+            / math.sqrt(config.hidden_size)
+        )
+
+        prime = 2147483647  # 梅森素数 2^31 - 1
+        rng = torch.Generator()
+        rng.manual_seed(config.hash_seed)
+        a1 = torch.randint(1, prime, (1,), generator=rng, dtype=torch.long)
+        b1 = torch.randint(0, prime, (1,), generator=rng, dtype=torch.long)
+        a2 = torch.randint(1, prime, (1,), generator=rng, dtype=torch.long)
+        b2 = torch.randint(0, prime, (1,), generator=rng, dtype=torch.long)
+
+        self.register_buffer("prime", torch.tensor(prime, dtype=torch.long))
+        self.register_buffer("a1", a1)
+        self.register_buffer("b1", b1)
+        self.register_buffer("a2", a2)
+        self.register_buffer("b2", b2)
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        x = input_ids.long()
+        bucket_idx = ((x * self.a1 + self.b1) % self.prime) % self.num_buckets
+        sign = ((x * self.a2 + self.b2) % self.prime) % 2
+        sign = (sign * 2 - 1).float()
+        return self.hash_table[bucket_idx] * sign.unsqueeze(-1)
+
+
+class CollinsModel(PreTrainedModel):
+    """
+    Collins-RoPE Encoder，输出 last_hidden_state 和 pooler_output。
+    使用 transformers.models.bert 的 BertEncoder + RoPE 替换 BertEmbeddings。
+    """
+
+    config_class = CollinsConfig
+    base_model_prefix = "collins"
+    supports_gradient_checkpointing = True
+
+    def __init__(self, config: CollinsConfig):
+        super().__init__(config)
+        self.config = config
+
+        self.embeddings = CollinsHashEmbedding(config)
+
+        # 直接复用 HF BertEncoder（含 Multi-Head Attention + FFN + LayerNorm）
+        from transformers.models.bert.modeling_bert import BertEncoder, BertConfig
+
+        bert_cfg = BertConfig(
+            hidden_size=config.hidden_size,
+            num_hidden_layers=config.num_hidden_layers,
+            num_attention_heads=config.num_attention_heads,
+            intermediate_size=config.intermediate_size,
+            hidden_dropout_prob=config.hidden_dropout_prob,
+            attention_probs_dropout_prob=config.attention_probs_dropout_prob,
+            max_position_embeddings=config.max_position_embeddings,
+            # 关闭 Bert 自带的位置编码，我们用 RoPE
+            position_embedding_type="relative_key_query",
+        )
+        bert_cfg._attn_implementation = "eager"
+        self.encoder = BertEncoder(bert_cfg)
+
+        # RoPE 频率缓冲（无参数）
+        dim = config.hidden_size
+        inv_freq = 1.0 / (
+            10000 ** (torch.arange(0, dim, 2).float() / dim)
+        )
+        t = torch.arange(config.max_position_embeddings).float()
+        freqs = torch.einsum("i,j->ij", t, inv_freq)
+        self.register_buffer("rope_cos", freqs.cos())
+        self.register_buffer("rope_sin", freqs.sin())
+
+        self.post_init()
+
+    def _apply_rope(self, x: torch.Tensor) -> torch.Tensor:
+        seq_len = x.shape[1]
+        cos = self.rope_cos[:seq_len].unsqueeze(0)
+        sin = self.rope_sin[:seq_len].unsqueeze(0)
+        x1, x2 = x[..., 0::2], x[..., 1::2]
+        return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+
+    def get_extended_attention_mask(self, attention_mask: torch.Tensor) -> torch.Tensor:
+        # BertEncoder 需要 [B, 1, 1, L] 形式的 mask，0 = 保留，-inf = 忽略
+        extended = attention_mask[:, None, None, :]
+        extended = (1.0 - extended.float()) * torch.finfo(torch.float32).min
+        return extended
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ):
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+
+        x = self.embeddings(input_ids)          # [B, L, D]
+        x = self._apply_rope(x)                 # [B, L, D]
+
+        ext_mask = self.get_extended_attention_mask(attention_mask)
+        encoder_out = self.encoder(x, attention_mask=ext_mask)
+        hidden_states = encoder_out.last_hidden_state  # [B, L, D]
+
+        # Mean Pooling
+        mask = attention_mask.unsqueeze(-1).float()
+        pooled = (hidden_states * mask).sum(1) / mask.sum(1).clamp(min=1e-9)
+        pooled = F.normalize(pooled, p=2, dim=-1)
+
+        if not return_dict:
+            return (hidden_states, pooled)
+
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=None,
+            attentions=None,
+        ), pooled
+
+
+class CollinsSTWrapper(nn.Module):
+    """
+    sentence-transformers 5.x 兼容包装层。
+    持有 tokenizer，实现 tokenize() 接口，同时注入 sentence_embedding。
+    """
+
+    def __init__(self, collins_model: CollinsModel, tokenizer_name_or_path: str = "bert-base-uncased", max_seq_length: int = 128):
+        super().__init__()
+        from transformers import AutoTokenizer
+        self.collins_model = collins_model
+        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)
+        self.max_seq_length = max_seq_length
+
+    def tokenize(self, texts: list[str], padding: str | bool = True) -> dict:
+        return self.tokenizer(
+            texts,
+            padding=padding,
+            truncation=True,
+            max_length=self.max_seq_length,
+            return_tensors="pt",
+        )
+
+    def forward(self, features: dict) -> dict:
+        input_ids = features["input_ids"]
+        attention_mask = features.get("attention_mask", None)
+        _, pooled = self.collins_model(input_ids, attention_mask)
+        features["sentence_embedding"] = pooled
+        return features
+
+    def save(self, output_path: str):
+        self.collins_model.save_pretrained(output_path)
+        self.tokenizer.save_pretrained(output_path)
+
+    @staticmethod
+    def load(input_path: str) -> "CollinsSTWrapper":
+        model = CollinsModel.from_pretrained(input_path)
+        return CollinsSTWrapper(model, tokenizer_name_or_path=input_path)