import spaces # ⚠️ DOIT ÊTRE LA TOUTE PREMIÈRE LIGNE import gradio as gr import torch import torch.nn as nn import torch.nn.functional as F import json import re from huggingface_hub import hf_hub_download from safetensors.torch import load_file # ========================================== # 1. DÉFINITION DE L'ARCHITECTURE MAISON # ========================================== class LightSelfAttention(nn.Module): def __init__(self, hidden_dim, n_heads=8): super().__init__() assert hidden_dim % n_heads == 0 self.n_heads = n_heads self.head_dim = hidden_dim // n_heads self.q_proj = nn.Linear(hidden_dim, hidden_dim, bias=True) self.k_proj = nn.Linear(hidden_dim, hidden_dim, bias=True) self.v_proj = nn.Linear(hidden_dim, hidden_dim, bias=True) self.o_proj = nn.Linear(hidden_dim, hidden_dim, bias=True) def forward(self, x): B, T, C = x.shape q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) k = self.k_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) v = self.v_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) attn = (q @ k.transpose(-2, -1)) / (self.head_dim ** 0.5) mask = torch.triu(torch.ones(T, T, device=x.device), diagonal=1).bool() attn = attn.masked_fill(mask, float("-inf")) attn = F.softmax(attn, dim=-1) out = (attn @ v).transpose(1, 2).contiguous().view(B, T, C) return self.o_proj(out) class LightMoE(nn.Module): def __init__(self, hidden_dim, n_experts=8, ffn_mult=2): super().__init__() self.router = nn.Linear(hidden_dim, n_experts, bias=True) self.experts = nn.ModuleList([ nn.Sequential( nn.Linear(hidden_dim, hidden_dim * ffn_mult), nn.GELU(), nn.Linear(hidden_dim * ffn_mult, hidden_dim), ) for _ in range(n_experts) ]) def forward(self, x): weights = F.softmax(self.router(x), dim=-1) top1 = torch.argmax(weights, dim=-1) out = torch.zeros_like(x) for e_id, expert in enumerate(self.experts): mask = (top1 == e_id).unsqueeze(-1) if mask.any(): out = out + mask * expert(x) return out class GopuBlockMoE(nn.Module): def __init__(self, hidden_dim, n_heads=8, n_experts=8): super().__init__() self.norm1 = nn.LayerNorm(hidden_dim) self.attn = LightSelfAttention(hidden_dim, n_heads) self.norm2 = nn.LayerNorm(hidden_dim) self.moe = LightMoE(hidden_dim, n_experts) def forward(self, x): x = x + self.attn(self.norm1(x)) x = x + self.moe(self.norm2(x)) return x class GopuTransformerLite(nn.Module): def __init__(self, vocab_size, hidden_dim=256, n_layers=4, n_heads=8, n_experts=8, max_seq_len=64): super().__init__() self.embedding = nn.Embedding(vocab_size, hidden_dim) self.pos_embedding = nn.Embedding(max_seq_len, hidden_dim) self.blocks = nn.ModuleList([ GopuBlockMoE(hidden_dim, n_heads, n_experts) for _ in range(n_layers) ]) self.norm_out = nn.LayerNorm(hidden_dim) self.fc_out = nn.Linear(hidden_dim, vocab_size) def forward(self, x): B, T = x.shape positions = torch.arange(T, device=x.device).unsqueeze(0) h = self.embedding(x) + self.pos_embedding(positions) for block in self.blocks: h = block(h) h = self.norm_out(h) return self.fc_out(h) # ========================================== # 2. CHARGEMENT DU MODÈLE ET DES POIDS # ========================================== REPO_ID = "Mauricio-100/locgi" device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print("Téléchargement des fichiers depuis le Hub...") config_path = hf_hub_download(REPO_ID, "config.json") vocab_path = hf_hub_download(REPO_ID, "vocab.json") weights_path = hf_hub_download(REPO_ID, "gopu_poids.safetensors") with open(config_path, "r", encoding="utf-8") as f: config = json.load(f) with open(vocab_path, "r", encoding="utf-8") as f: word_to_int = json.load(f) # Création du dictionnaire inversé pour le décodage int_to_word = {int(v): k for k, v in word_to_int.items()} vocab_size = config["vocab_size"] SEQ_LEN = config["max_seq_len"] # Instanciation du modèle model = GopuTransformerLite( vocab_size=vocab_size, hidden_dim=config["hidden_dim"], n_layers=config["n_layers"], n_heads=config["n_heads"], n_experts=config["n_experts"], max_seq_len=SEQ_LEN ).to(device) # Chargement des poids safetensors model.load_state_dict(load_file(weights_path)) model.eval() print("Modèle chargé avec succès !") # ========================================== # 3. FONCTION DE GÉNÉRATION (TOKENIZER MAISON) # ========================================== @spaces.GPU # ⚠️ DÉCORATEUR POUR ALLOUER LE GPU PENDANT L'INFÉRENCE def generer(prompt, max_len=30, temperature=0.8): # On garde votre format de prompt d'origine texte_initial = f"question : {prompt} réponse :".lower() tokens = re.findall(r"\w+|[^\w\s]", texte_initial) # Gestion du token unk_id = word_to_int.get("", 1) indices = [word_to_int.get(t, unk_id) for t in tokens] generes = [] with torch.no_grad(): for _ in range(int(max_len)): # Ne garder que les derniers tokens correspondant à max_seq_len (32) contexte_actuel = indices[-SEQ_LEN:] x = torch.tensor([contexte_actuel], dtype=torch.long, device=device) out = model(x) logits = out[0, -1] / max(temperature, 1e-5) # protection division par zero probs = torch.softmax(logits, dim=-1) idx = torch.multinomial(probs, 1).item() mot = int_to_word.get(idx, "") if mot == "": break generes.append(mot) indices.append(idx) # Arrêt basique sur la ponctuation finale if mot in (".", "!", "?") and len(generes) > 2: break # Reconstruire la phrase proprement phrase = "" for tok in generes: phrase += tok if (tok in {".", ",", "!", "?"} or not phrase) else " " + tok return phrase # ========================================== # 4. INTERFACE GRADIO # ========================================== demo = gr.Interface( fn=generer, inputs=[ gr.Textbox(lines=2, label="Votre question", placeholder="Quelle est la capitale de la France ?"), gr.Slider(minimum=10, maximum=50, value=30, step=1, label="Longueur maximale de réponse (tokens)"), gr.Slider(minimum=0.1, maximum=2.0, value=0.8, step=0.1, label="Température") ], outputs=gr.Textbox(label="Réponse de Locgi"), title="Locgi - Modèle Français MoE (36M)", description=("Ce Space utilise **Locgi**, un petit modèle de langage expérimental " "avec une architecture Mixture of Experts (MoE) entraîné from scratch. \n\n" "⚠️ *Limitations : contexte max de 32 tokens, tokenisation mot-à-mot (BPE non utilisé).*") ) if __name__ == "__main__": demo.launch()