app.py

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>
    unk_id = word_to_int.get("<UNK>", 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, "<UNK>")
            
            if mot == "<PAD>":
                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()