Upload 9 files

config.json

@@ -0,0 +1,28 @@
+{
+  "model_type": "rne_tiny_gpt",
+  "architectures": [
+    "RNETinyGPTModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_rne_tiny_gpt.RNETinyGPTConfig",
+    "AutoModel": "modeling_rne_tiny_gpt.RNETinyGPTModel"
+  },
+  "vocab_size": 32768,
+  "ctx_len": 4096,
+  "max_position_embeddings": 4096,
+  "n_layer": 4,
+  "num_hidden_layers": 4,
+  "n_head": 4,
+  "num_attention_heads": 4,
+  "n_embd": 384,
+  "hidden_size": 384,
+  "dropout": 0.0,
+  "pad_token_id": 0,
+  "sep_token_id": 3,
+  "pooling": "mean",
+  "normalize_embeddings": true,
+  "attention_backend": "torch",
+  "torch_fallback": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "custom"
+}

configuration_rne_tiny_gpt.py

@@ -0,0 +1,45 @@
+from transformers import PretrainedConfig
+
+
+class RNETinyGPTConfig(PretrainedConfig):
+    model_type = "rne_tiny_gpt"
+
+    def __init__(
+        self,
+        vocab_size=32768,
+        ctx_len=4096,
+        n_layer=4,
+        n_head=4,
+        n_embd=384,
+        dropout=0.0,
+        pad_token_id=0,
+        sep_token_id=3,
+        pooling="mean",
+        normalize_embeddings=True,
+        attention_backend="sage",
+        torch_fallback=False,
+        **kwargs,
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            sep_token_id=sep_token_id,
+            **kwargs,
+        )
+
+        self.vocab_size = int(vocab_size)
+        self.ctx_len = int(ctx_len)
+        self.max_position_embeddings = int(ctx_len)
+
+        self.n_layer = int(n_layer)
+        self.n_head = int(n_head)
+        self.n_embd = int(n_embd)
+
+        self.num_hidden_layers = int(n_layer)
+        self.num_attention_heads = int(n_head)
+        self.hidden_size = int(n_embd)
+
+        self.dropout = float(dropout)
+        self.pooling = str(pooling)
+        self.normalize_embeddings = bool(normalize_embeddings)
+        self.attention_backend = str(attention_backend)
+        self.torch_fallback = bool(torch_fallback)

generation_config.json

@@ -0,0 +1,4 @@
+{
+  "pad_token_id": 0,
+  "sep_token_id": 3
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9d761488ad84788bd5f9f5ca3e4c10661c08d76c43616e7fac1e1b6832258e81
+size 67650336

modeling_rne_tiny_gpt.py

@@ -0,0 +1,212 @@
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+
+from .configuration_rne_tiny_gpt import RNETinyGPTConfig
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config: RNETinyGPTConfig):
+        super().__init__()
+
+        if config.n_embd % config.n_head != 0:
+            raise ValueError("n_embd must be divisible by n_head")
+
+        self.n_head = config.n_head
+        self.head_dim = config.n_embd // config.n_head
+        self.attention_backend = getattr(config, "attention_backend", "sage")
+        self.torch_fallback = bool(getattr(config, "torch_fallback", False))
+
+        if self.attention_backend not in ("sage", "torch"):
+            raise ValueError("attention_backend must be 'sage' or 'torch'")
+
+        if self.attention_backend == "sage" and self.head_dim not in (64, 96, 128):
+            raise ValueError(f"SageAttention requires head_dim in [64, 96, 128], got {self.head_dim}.")
+
+        if self.attention_backend == "sage" and config.dropout != 0.0:
+            raise ValueError("SageAttention strict mode requires dropout=0.0")
+
+        self.qkv = nn.Linear(config.n_embd, 3 * config.n_embd, bias=False)
+        self.proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
+        self.dropout = nn.Dropout(config.dropout)
+
+        mask = torch.tril(torch.ones(config.ctx_len, config.ctx_len, dtype=torch.bool))
+        self.register_buffer("mask", mask.view(1, 1, config.ctx_len, config.ctx_len), persistent=False)
+
+        self.sageattn = None
+        if self.attention_backend == "sage":
+            try:
+                from sageattention import sageattn
+                self.sageattn = sageattn
+            except Exception as exc:
+                if self.torch_fallback:
+                    self.attention_backend = "torch"
+                    self.sageattn = None
+                else:
+                    raise RuntimeError(
+                        "Ce modèle a été entraîné avec SageAttention. "
+                        "Installe sageattention: pip install sageattention"
+                    ) from exc
+
+    def _torch_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, t: int) -> torch.Tensor:
+        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        scores = scores.masked_fill(self.mask[:, :, :t, :t] == 0, float("-inf"))
+        att = F.softmax(scores.float(), dim=-1).to(q.dtype)
+        att = self.dropout(att)
+        y = att @ v
+        return y
+
+    def _sage_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
+        if self.sageattn is None:
+            raise RuntimeError("SageAttention demandé mais sageattn est None")
+        if not q.is_cuda:
+            if self.torch_fallback:
+                return None
+            raise RuntimeError(
+                "SageAttention exige CUDA. Passe le modèle sur CUDA avec model.cuda(), "
+                "ou active torch_fallback dans config.json."
+            )
+        q = q.contiguous()
+        k = k.contiguous()
+        v = v.contiguous()
+        return self.sageattn(q, k, v, tensor_layout="HND", is_causal=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        b, t, c = x.shape
+        qkv = self.qkv(x)
+        q, k, v = qkv.chunk(3, dim=-1)
+        q = q.view(b, t, self.n_head, self.head_dim).transpose(1, 2).contiguous()
+        k = k.view(b, t, self.n_head, self.head_dim).transpose(1, 2).contiguous()
+        v = v.view(b, t, self.n_head, self.head_dim).transpose(1, 2).contiguous()
+
+        if self.attention_backend == "sage":
+            y = self._sage_attention(q, k, v)
+            if y is None:
+                y = self._torch_attention(q, k, v, t)
+        else:
+            y = self._torch_attention(q, k, v, t)
+
+        y = y.transpose(1, 2).contiguous().view(b, t, c)
+        y = self.proj(y)
+        return y
+
+
+class MLP(nn.Module):
+    def __init__(self, config: RNETinyGPTConfig):
+        super().__init__()
+        self.fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
+        self.proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc(x)
+        x = F.gelu(x)
+        x = self.proj(x)
+        x = self.dropout(x)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(self, config: RNETinyGPTConfig):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(config.n_embd)
+        self.attn = CausalSelfAttention(config)
+        self.ln2 = nn.LayerNorm(config.n_embd)
+        self.mlp = MLP(config)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.ln1(x))
+        x = x + self.mlp(self.ln2(x))
+        return x
+
+
+class RNETinyGPTPreTrainedModel(PreTrainedModel):
+    config_class = RNETinyGPTConfig
+    base_model_prefix = "rne_tiny_gpt"
+    supports_gradient_checkpointing = False
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        if isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+
+class RNETinyGPTModel(RNETinyGPTPreTrainedModel):
+    def __init__(self, config: RNETinyGPTConfig):
+        super().__init__(config)
+        self.tok_emb = nn.Embedding(config.vocab_size, config.n_embd)
+        self.pos_emb = nn.Embedding(config.ctx_len, config.n_embd)
+        self.drop = nn.Dropout(config.dropout)
+        self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layer)])
+        self.ln_f = nn.LayerNorm(config.n_embd)
+        self.post_init()
+
+    def _mean_pool(self, hidden: torch.Tensor, attention_mask: Optional[torch.Tensor], input_ids: torch.Tensor) -> torch.Tensor:
+        if attention_mask is None:
+            mask = input_ids.ne(self.config.pad_token_id)
+        else:
+            mask = attention_mask.bool()
+        mask = mask.unsqueeze(-1).to(hidden.dtype)
+        summed = (hidden * mask).sum(dim=1)
+        denom = mask.sum(dim=1).clamp(min=1.0)
+        return summed / denom
+
+    def _last_pool(self, hidden: torch.Tensor, attention_mask: Optional[torch.Tensor], input_ids: torch.Tensor) -> torch.Tensor:
+        if attention_mask is None:
+            mask = input_ids.ne(self.config.pad_token_id)
+        else:
+            mask = attention_mask.bool()
+        lengths = mask.sum(dim=1).clamp(min=1)
+        last_pos = lengths - 1
+        batch_idx = torch.arange(input_ids.size(0), device=input_ids.device)
+        return hidden[batch_idx, last_pos, :]
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = True,
+        **kwargs,
+    ):
+        b, t = input_ids.shape
+        if t > self.config.ctx_len:
+            raise ValueError(f"Input length {t} > ctx_len {self.config.ctx_len}. Truncate before calling the model.")
+
+        pos = torch.arange(0, t, dtype=torch.long, device=input_ids.device).unsqueeze(0)
+        x = self.tok_emb(input_ids) + self.pos_emb(pos)
+        x = self.drop(x)
+        for block in self.blocks:
+            x = block(x)
+        hidden = self.ln_f(x)
+
+        if self.config.pooling == "last":
+            pooled = self._last_pool(hidden, attention_mask, input_ids)
+        else:
+            pooled = self._mean_pool(hidden, attention_mask, input_ids)
+
+        pooled = pooled.float()
+        if self.config.normalize_embeddings:
+            pooled = F.normalize(pooled, p=2, dim=-1)
+
+        if not return_dict:
+            return (hidden, pooled)
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=hidden,
+            pooler_output=pooled,
+            hidden_states=None,
+            attentions=None,
+        )
+
+    @torch.no_grad()
+    def encode(self, input_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        out = self.forward(input_ids=input_ids, attention_mask=attention_mask, return_dict=True)
+        return out.pooler_output

requirements.txt

@@ -0,0 +1,5 @@
+torch
+transformers
+safetensors
+tokenizers
+sageattention

special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "unk_token": "[UNK]",
+  "pad_token": "[PAD]",
+  "cls_token": "[CLS]",
+  "sep_token": "[SEP]",
+  "mask_token": "[MASK]"
+}

tokenizer_config.json

@@ -0,0 +1,12 @@
+{
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "model_max_length": 4096,
+  "padding_side": "right",
+  "truncation_side": "right",
+  "clean_up_tokenization_spaces": false,
+  "unk_token": "[UNK]",
+  "pad_token": "[PAD]",
+  "cls_token": "[CLS]",
+  "sep_token": "[SEP]",
+  "mask_token": "[MASK]"
+}