""" model.py Architecture du modèle "Hilal Ai" : Transformer Décodeur (style GPT) entraîné from scratch au niveau caractère. Cette architecture DOIT correspondre exactement à celle utilisée pendant l'entraînement sur Colab pour que model.load_state_dict(...) fonctionne. Si tu as modifié l'architecture sur Colab (dropout, bias, etc.), répercute les mêmes changements ici. """ import json import logging import os import torch import torch.nn as nn from torch.nn import functional as F logger = logging.getLogger("hilal_ai.model") # ---------------------------------------------------------------------------- # Hyperparamètres (doivent matcher l'entraînement Colab) # ---------------------------------------------------------------------------- BLOCK_SIZE = 128 N_EMBD = 192 N_HEAD = 6 N_LAYER = 4 DROPOUT = 0.1 # identique à l'entraînement Colab. Sans effet en model.eval() # (le dropout est désactivé automatiquement en inference), # mais gardé identique pour la cohérence de l'architecture. # ---------------------------------------------------------------------------- # Architecture # ---------------------------------------------------------------------------- class Head(nn.Module): """Une tête d'auto-attention causale.""" def __init__(self, head_size): super().__init__() self.key = nn.Linear(N_EMBD, head_size, bias=False) self.query = nn.Linear(N_EMBD, head_size, bias=False) self.value = nn.Linear(N_EMBD, head_size, bias=False) self.register_buffer("tril", torch.tril(torch.ones(BLOCK_SIZE, BLOCK_SIZE))) self.dropout = nn.Dropout(DROPOUT) def forward(self, x): B, T, C = x.shape k = self.key(x) q = self.query(x) wei = q @ k.transpose(-2, -1) * (k.shape[-1] ** -0.5) wei = wei.masked_fill(self.tril[:T, :T] == 0, float("-inf")) wei = F.softmax(wei, dim=-1) wei = self.dropout(wei) v = self.value(x) out = wei @ v return out class MultiHeadAttention(nn.Module): """Plusieurs têtes d'attention en parallèle.""" def __init__(self, num_heads, head_size): super().__init__() self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)]) self.proj = nn.Linear(N_EMBD, N_EMBD) self.dropout = nn.Dropout(DROPOUT) def forward(self, x): out = torch.cat([h(x) for h in self.heads], dim=-1) out = self.dropout(self.proj(out)) return out class FeedForward(nn.Module): """Simple couche feed-forward (MLP) avec activation GELU.""" def __init__(self, n_embd): super().__init__() self.net = nn.Sequential( nn.Linear(n_embd, 4 * n_embd), nn.GELU(), nn.Linear(4 * n_embd, n_embd), nn.Dropout(DROPOUT), ) def forward(self, x): return self.net(x) class Block(nn.Module): """Bloc Transformer : communication (attention) puis calcul (FFN).""" def __init__(self, n_embd, n_head): super().__init__() head_size = n_embd // n_head self.sa = MultiHeadAttention(n_head, head_size) self.ffwd = FeedForward(n_embd) self.ln1 = nn.LayerNorm(n_embd) self.ln2 = nn.LayerNorm(n_embd) def forward(self, x): x = x + self.sa(self.ln1(x)) x = x + self.ffwd(self.ln2(x)) return x class MiniGPT(nn.Module): """Modèle de langage Transformer Décodeur, niveau caractère.""" def __init__(self, vocab_size): super().__init__() self.vocab_size = vocab_size self.token_embedding = nn.Embedding(vocab_size, N_EMBD) self.position_embedding = nn.Embedding(BLOCK_SIZE, N_EMBD) self.blocks = nn.Sequential(*[Block(N_EMBD, N_HEAD) for _ in range(N_LAYER)]) self.ln_f = nn.LayerNorm(N_EMBD) self.lm_head = nn.Linear(N_EMBD, vocab_size) def forward(self, idx, targets=None): B, T = idx.shape tok_emb = self.token_embedding(idx) pos_emb = self.position_embedding( torch.arange(T, device=idx.device) ) x = tok_emb + pos_emb x = self.blocks(x) x = self.ln_f(x) logits = self.lm_head(x) loss = None if targets is not None: B, T, C = logits.shape logits_flat = logits.view(B * T, C) targets_flat = targets.view(B * T) loss = F.cross_entropy(logits_flat, targets_flat) return logits, loss @torch.no_grad() def generate(self, idx, max_new_tokens, temperature=0.7, top_k=None): """ Génère max_new_tokens nouveaux caractères à partir du contexte idx. idx : tensor (B, T) d'indices déjà encodés. Le contexte est tronqué à BLOCK_SIZE à chaque étape (fenêtre glissante). """ self.eval() for _ in range(max_new_tokens): idx_cond = idx[:, -BLOCK_SIZE:] logits, _ = self(idx_cond) logits = logits[:, -1, :] / max(temperature, 1e-6) if top_k is not None: v, _ = torch.topk(logits, min(top_k, logits.size(-1))) logits[logits < v[:, [-1]]] = float("-inf") probs = F.softmax(logits, dim=-1) idx_next = torch.multinomial(probs, num_samples=1) idx = torch.cat((idx, idx_next), dim=1) return idx # ---------------------------------------------------------------------------- # Tokenizer niveau caractère # ---------------------------------------------------------------------------- class CharTokenizer: """ Tokenizer char-level. Le mapping stoi/itos DOIT être identique à celui utilisé pendant l'entraînement sur Colab, sinon le modèle produira du charabia (les poids des embeddings ne correspondront pas aux bons caractères). Stratégie recommandée : 1. Pendant l'entraînement sur Colab, sauvegarde ton vocabulaire : import json with open("vocab.json", "w", encoding="utf-8") as f: json.dump({"chars": sorted(list(set(text)))}, f, ensure_ascii=False) puis télécharge vocab.json à côté de model.pt. 2. Place vocab.json dans le même dossier que ce script. Si vocab.json est absent, un jeu de caractères par défaut est utilisé en secours (lettres FR/EN, chiffres, ponctuation courante). Dans ce cas, la cohérence avec les poids entraînés n'est PAS garantie. """ DEFAULT_CHARS = ( "\n !\"#$%&'()*+,-./0123456789:;<=>?@" "ABCDEFGHIJKLMNOPQRSTUVWXYZ[\\]^_`" "abcdefghijklmnopqrstuvwxyz{|}~" "àâäéèêëïîôöùûüçÀÂÄÉÈÊËÏÎÔÖÙÛÜÇœŒ«»…–—’" ) def __init__(self, vocab_path="vocab.json"): chars = None if vocab_path and os.path.exists(vocab_path): try: with open(vocab_path, "r", encoding="utf-8") as f: data = json.load(f) chars = data["chars"] logger.info("Vocabulaire chargé depuis %s (%d caractères).", vocab_path, len(chars)) except Exception as e: logger.warning("Impossible de lire %s (%s). Utilisation du vocabulaire par défaut.", vocab_path, e) if chars is None: chars = sorted(set(self.DEFAULT_CHARS)) logger.warning( "AUCUN vocab.json trouvé : utilisation d'un vocabulaire par défaut (%d caractères). " "Ceci ne correspondra probablement PAS exactement aux poids entraînés sur Colab.", len(chars), ) self.chars = chars self.vocab_size = len(chars) self.stoi = {ch: i for i, ch in enumerate(chars)} self.itos = {i: ch for i, ch in enumerate(chars)} # Caractère de remplacement pour tout caractère inconnu en entrée self.unk_char = chars[0] def encode(self, s: str): return [self.stoi.get(c, self.stoi.get(self.unk_char, 0)) for c in s] def decode(self, indices) -> str: return "".join(self.itos.get(int(i), "") for i in indices) # ---------------------------------------------------------------------------- # Chargement du modèle complet # ---------------------------------------------------------------------------- def load_model(model_path: str, vocab_path: str = "vocab.json", device: str = "cpu"): """ Instancie le tokenizer + le modèle, charge les poids depuis model_path, et renvoie (model, tokenizer) prêts pour l'inférence. """ tokenizer = CharTokenizer(vocab_path) model = MiniGPT(vocab_size=tokenizer.vocab_size) state_dict = torch.load(model_path, map_location=device) # Selon comment le modèle a été sauvegardé sur Colab (state_dict pur, # ou checkpoint dict contenant 'model_state_dict'), on gère les deux cas. if isinstance(state_dict, dict) and "model_state_dict" in state_dict: state_dict = state_dict["model_state_dict"] missing, unexpected = model.load_state_dict(state_dict, strict=False) if missing: logger.warning("Clés manquantes lors du chargement des poids : %s", missing) if unexpected: logger.warning("Clés inattendues lors du chargement des poids : %s", unexpected) model.to(device) model.eval() return model, tokenizer