lea97338/LunaLMV2

import re, torch, gradio as gr
import torch.nn as nn
from huggingface_hub import hf_hub_download

REPO = "lea97338/LunaLMV2"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Téléchargement public : aucun token
path = hf_hub_download(REPO, "luna_model.pth")
ckpt = torch.load(path, map_location="cpu", weights_only=False)

class LunaLM(nn.Module):
    def __init__(self, vocab_size, d_model=128, max_len=64):
        super().__init__()
        self.max_len, self.vocab_size = max_len, vocab_size
        self.wte = nn.Embedding(vocab_size, d_model)
        self.wpe = nn.Embedding(max_len, d_model)
        self.ln1, self.ln2, self.ln_f = (
            nn.LayerNorm(d_model), nn.LayerNorm(d_model),
            nn.LayerNorm(d_model))
        self.q_proj, self.k_proj, self.v_proj, self.out_proj = [
            nn.Linear(d_model, d_model) for _ in range(4)]
        self.mlp = nn.Sequential(
            nn.Linear(d_model, d_model*2), nn.GELU(),
            nn.Linear(d_model*2, d_model))
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, idx):
        _, T = idx.shape
        x = self.wte(idx) + self.wpe(
            torch.arange(T, device=idx.device).unsqueeze(0))
        n = self.ln1(x)
        q, k, v = self.q_proj(n), self.k_proj(n), self.v_proj(n)
        a = q @ k.transpose(-2, -1) / (k.size(-1) ** 0.5)
        mask = torch.tril(torch.ones(T, T, device=idx.device)).bool()
        x = x + self.out_proj(
            torch.softmax(a.masked_fill(~mask, float("-inf")), -1) @ v)
        return self.lm_head(self.ln_f(x + self.mlp(self.ln2(x))))

    @staticmethod
    def clean_tok(text):
        return re.findall(r"\w+|[^\w\s]", text)

model = LunaLM(len(ckpt["w2i"]))
model.load_state_dict(ckpt["model"])
model.w2i, model.i2w = ckpt["w2i"], ckpt["i2w"]
model = model.to(DEVICE).eval()

@torch.inference_mode()
def generate_stream(prompt):
    ids = [model.w2i["<bos>"]] + [
        model.w2i.get(w, 3)
        for w in ["user"] + model.clean_tok(prompt) + ["luna"]]
    ids = ids[:model.max_len]
    result = ""

    while len(ids) < model.max_len:
        x = torch.tensor([ids], device=DEVICE)
        token = int(model(x)[0, -1].argmax().item())
        if token in (model.w2i["<eos>"], model.w2i["<pad>"]):
            break
        word = model.i2w.get(token, "")
        result += ("" if word in ".,!?;:" else " ") + word
        ids.append(token)
        yield result.strip()

demo = gr.ChatInterface(
    fn=generate_stream,
    title="🌙 LunaLMV2",
    description="Inférence seule — génération en streaming."
)
demo.queue().launch(share=True)