File size: 9,691 Bytes
f011463
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
#!/usr/bin/env python3
# ============================================================================
# LEGAL NOTICE & WATERMARK: R08 NON-COMMERCIAL RESEARCH & ETHICAL LICENSE v2.0
# Copyright (c) 2026 Thomas Seri (Sentinel). All Rights Reserved.
# Exclusive Jurisdiction: Court of Milan (Sezione Specializzata Impresa), Italy.
# Clean-room reverse engineering, distillation, commercial or military use
# is strictly prohibited and constitutes a breach of contract.
# ============================================================================
from __future__ import annotations
__author__ = "Thomas Seri (Sentinel)"
__license__ = "R08-Non-Commercial-Ethical-License-v2.0"
__jurisdiction__ = "Court of Milan, Italy"

"""Transformer Baseline 45M — Dense GPT-2 Architecture (14 Attention Layers)
Official Hugging Face / PyTorch implementation for comparative A/B benchmarking vs Sentinel R-08.
Author: Thomas Seri (Sentinel LifeOS)
"""

import json
import math
from dataclasses import dataclass, fields
from pathlib import Path
from typing import List, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from safetensors.torch import load_file

import pezzi_tokenizer as pt
from vocab_stages import load_bundle, StageBundle


@dataclass
class TransformerConfig:
    vocab_size: int = 1251
    d_model: int = 512
    n_layers: int = 14
    n_heads: int = 8
    max_seq_len: int = 256
    dropout: float = 0.0

    @classmethod
    def from_dict(cls, d: dict) -> "TransformerConfig":
        names = {f.name for f in fields(cls)}
        return cls(**{k: v for k, v in d.items() if k in names})


class CausalSelfAttention(nn.Module):
    def __init__(self, cfg: TransformerConfig):
        super().__init__()
        if cfg.d_model % cfg.n_heads != 0:
            raise ValueError(f"d_model={cfg.d_model} non divisibile per n_heads={cfg.n_heads}")
        self.n_heads = cfg.n_heads
        self.head_dim = cfg.d_model // cfg.n_heads
        self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model)
        self.proj = nn.Linear(cfg.d_model, cfg.d_model)
        self.attn_dropout = float(cfg.dropout)
        self.resid_drop = nn.Dropout(cfg.dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, t, d = x.shape
        q, k, v = self.qkv(x).view(b, t, 3, self.n_heads, self.head_dim).permute(2, 0, 3, 1, 4)
        y = F.scaled_dot_product_attention(
            q, k, v, is_causal=True, dropout_p=self.attn_dropout if self.training else 0.0
        )
        return self.resid_drop(self.proj(y.transpose(1, 2).reshape(b, t, d)))


class MLP(nn.Module):
    def __init__(self, cfg: TransformerConfig):
        super().__init__()
        self.fc = nn.Linear(cfg.d_model, 4 * cfg.d_model)
        self.proj = nn.Linear(4 * cfg.d_model, cfg.d_model)
        self.drop = nn.Dropout(cfg.dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.drop(self.proj(F.gelu(self.fc(x))))


class Block(nn.Module):
    def __init__(self, cfg: TransformerConfig):
        super().__init__()
        self.ln1 = nn.LayerNorm(cfg.d_model)
        self.attn = CausalSelfAttention(cfg)
        self.ln2 = nn.LayerNorm(cfg.d_model)
        self.mlp = MLP(cfg)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.ln1(x))
        return x + self.mlp(self.ln2(x))


class BaselineTransformer(nn.Module):
    def __init__(self, cfg: TransformerConfig):
        super().__init__()
        self.cfg = cfg
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.pos_emb = nn.Embedding(cfg.max_seq_len, cfg.d_model)
        self.drop = nn.Dropout(cfg.dropout)
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)])
        self.ln_f = nn.LayerNorm(cfg.d_model)
        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        self.lm_head.weight = self.tok_emb.weight
        self.apply(self._init_weights)
        std_res = 0.02 / math.sqrt(2 * cfg.n_layers)
        for blk in self.blocks:
            nn.init.normal_(blk.attn.proj.weight, mean=0.0, std=std_res)
            nn.init.normal_(blk.mlp.proj.weight, mean=0.0, std=std_res)

    @staticmethod
    def _init_weights(m: nn.Module) -> None:
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)

    def forward(
        self,
        pezzi_ids: torch.Tensor,
        parole_ids: Optional[torch.Tensor] = None,
        unioni_ids: Optional[torch.Tensor] = None,
        record_path: bool = False,
    ) -> Tuple[torch.Tensor, None]:
        b, t = pezzi_ids.shape
        if t > self.cfg.max_seq_len:
            raise ValueError(f"seq_len={t} > max_seq_len={self.cfg.max_seq_len}")
        pos = torch.arange(t, device=pezzi_ids.device)
        h = self.drop(self.tok_emb(pezzi_ids) + self.pos_emb(pos))
        for blk in self.blocks:
            h = blk(h)
        return self.lm_head(self.ln_f(h)), None


class TransformerTokenizer:
    """Tokenizer a subword (1.251 pezzi) identico a Sentinel R-08."""

    def __init__(self, bundle: StageBundle):
        self.bundle = bundle
        self.pezzi = bundle.pezzi
        self.eos_id = int(getattr(bundle.pezzi, "eos_id", 2))
        self.pad_id = int(getattr(bundle.pezzi, "pad_id", 3))

    @classmethod
    def from_pretrained(cls, pretrained_path: Union[str, Path]) -> "TransformerTokenizer":
        root = Path(pretrained_path)
        vocab_dir = root / "vocab_stages" if (root / "vocab_stages").exists() else root
        bundle = load_bundle(vocab_dir)
        return cls(bundle)

    def encode(self, text: str) -> List[int]:
        return pt.encode(text, self.pezzi)

    def decode(self, token_ids: List[int]) -> str:
        return pt.decode(token_ids, self.pezzi)


class TransformerModel(nn.Module):
    """Interfaccia Hugging Face friendly per il modello Transformer Baseline 45M."""

    def __init__(self, config: TransformerConfig, model: BaselineTransformer):
        super().__init__()
        self.config = config
        self.model = model

    @classmethod
    def from_pretrained(
        cls,
        pretrained_path: Union[str, Path],
        weights_file: str = "model.safetensors",
        device: Optional[str] = None,
        torch_dtype: Optional[torch.dtype] = None,
    ) -> "TransformerModel":
        root = Path(pretrained_path)
        cfg_file = root / "config.json"
        if not cfg_file.exists():
            raise FileNotFoundError(f"config.json non trovato in {root}")

        with open(cfg_file, "r", encoding="utf-8") as f:
            cfg_dict = json.load(f)

        cfg = TransformerConfig.from_dict(cfg_dict)
        model = BaselineTransformer(cfg)

        weights_path = root / weights_file
        if not weights_path.exists():
            raise FileNotFoundError(f"{weights_file} non trovato in {root}")

        state_dict = load_file(str(weights_path))
        model.load_state_dict(state_dict)

        if device is None:
            if torch.backends.mps.is_available():
                device = "mps"
            elif torch.cuda.is_available():
                device = "cuda"
            else:
                device = "cpu"

        model.to(torch.device(device))
        if torch_dtype is not None:
            model.to(dtype=torch_dtype)

        model.eval()
        return cls(cfg, model)

    def forward(self, pezzi_ids: torch.Tensor) -> Tuple[torch.Tensor, None]:
        return self.model(pezzi_ids)

    @torch.inference_mode()
    def generate(
        self,
        prompt: str,
        tokenizer: TransformerTokenizer,
        max_new_tokens: int = 48,
        temperature: float = 0.7,
        top_p: float = 0.9,
        repetition_penalty: float = 1.15,
    ) -> str:
        dev = next(self.model.parameters()).device
        ids = tokenizer.encode(prompt)
        if not ids:
            ids = [tokenizer.eos_id]

        new_tokens: List[int] = []
        for _ in range(max_new_tokens):
            x = torch.tensor([ids[-self.config.max_seq_len:]], dtype=torch.long, device=dev)
            logits, _ = self.model(x)
            next_logits = logits[0, -1]

            if repetition_penalty != 1.0 and new_tokens:
                for token_id in set(new_tokens):
                    if next_logits[token_id] > 0:
                        next_logits[token_id] /= repetition_penalty
                    else:
                        next_logits[token_id] *= repetition_penalty

            if temperature > 0:
                probs = torch.softmax(next_logits / temperature, dim=-1)
                if top_p < 1.0:
                    sorted_probs, sorted_indices = torch.sort(probs, descending=True)
                    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
                    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[sorted_indices_to_remove]
                    probs[indices_to_remove] = 0.0
                    probs = probs / probs.sum()
                nxt = int(torch.multinomial(probs, num_samples=1).item())
            else:
                nxt = int(next_logits.argmax().item())

            if nxt in (tokenizer.eos_id, tokenizer.pad_id):
                break

            new_tokens.append(nxt)
            ids.append(nxt)

        return tokenizer.decode(new_tokens)