R-08-Transformer-Baseline / modeling_transformer.py

Sync corrected public source from GitHub (2026-11-07)

f011463 verified 3 days ago

9.69 kB

#!/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)

	#!/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)