Mini-Groutouille

# ==============================================================================
# 🚀 Inférence Aricate V4 - Chargement depuis Hugging Face
# Ce code charge le modèle et le tokenizer customisés pour la génération de texte.
# ==============================================================================

import torch
import torch.nn as nn
import torch.nn.functional as F
import collections
from huggingface_hub import hf_hub_download, PyTorchModelHubMixin
import os
import json

# --- 1. Définitions des Classes (Réplique de l'architecture d'entraînement) ---

# --- A. WordTokenizer (Customisé pour le chargement) ---
class WordTokenizer:
    """Tokenizer simple pour l'architecture Aricate, chargé depuis le vocabulaire publié."""
    def __init__(self, vocab_data):
        self.word_to_id = vocab_data["word_to_id"]
        # Récupérer max_len_input pour le padding lors de l'inférence
        self.max_len_input = vocab_data["max_len_input"]
        self.id_to_word = {id: word for word, id in self.word_to_id.items()}
        self.vocab_size = len(self.word_to_id)

        # Définition des tokens spéciaux
        self.special_tokens = {
            '<pad>': self.word_to_id.get('<pad>', 0),
            '<unk>': self.word_to_id.get('<unk>', 1),
            '<eos>': self.word_to_id.get('<eos>', 2),
        }

        print(f"Tokenizer chargé. Taille du vocabulaire : {self.vocab_size}")
        print(f"Longueur maximale d'entrée pour le padding : {self.max_len_input}")

    def encode(self, text, add_eos=False):
        words = text.lower().split()
        if add_eos:
            words.append('<eos>')
        ids = [self.word_to_id.get(word, self.special_tokens['<unk>']) for word in words]
        return ids

    def decode(self, ids):
        words = [self.id_to_word.get(id, '<unk>') for id in ids]
        return " ".join(word for word in words if word not in ['<pad>', '<unk>', '<eos>'])

# --- B. AricateAttentionLayer (Inchangé) ---
class AricateAttentionLayer(nn.Module):
    """Couche d'Attention Additive (Bahdanau)."""
    def __init__(self, hidden_dim):
        super(AricateAttentionLayer, self).__init__()
        self.W = nn.Linear(hidden_dim, hidden_dim)
        self.U = nn.Linear(hidden_dim, hidden_dim)
        self.V = nn.Linear(hidden_dim, 1, bias=False)
    def forward(self, rnn_outputs, last_hidden):
        last_hidden_expanded = last_hidden.unsqueeze(1)
        energy = torch.tanh(self.W(rnn_outputs) + self.U(last_hidden_expanded))
        attention_weights_raw = self.V(energy).squeeze(2)
        attention_weights = F.softmax(attention_weights_raw, dim=1)
        context_vector = torch.sum(rnn_outputs * attention_weights.unsqueeze(2), dim=1)
        return context_vector

# --- C. AricateModel V4 (Doit hériter de PyTorchModelHubMixin) ---
class AricateModel(nn.Module, PyTorchModelHubMixin):
    """Architecture Aricate V4 pour Modélisation de Langage."""
    def __init__(self, vocab_size: int, embedding_dim: int, hidden_dim: int, num_layers: int = 1, config: dict = None, **kwargs):
        super(AricateModel, self).__init__()

        # Chargement de la configuration si disponible (utilisé par from_pretrained_fast)
        if config is not None:
             vocab_size = config.get("vocab_size", vocab_size)
             embedding_dim = config.get("embedding_dim", embedding_dim)
             hidden_dim = config.get("hidden_dim", hidden_dim)
             num_layers = config.get("num_layers", num_layers)

        self.vocab_size = vocab_size
        self.embedding_dim = embedding_dim
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers

        self.word_embeddings = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embedding_dim, padding_idx=0)
        self.rnn = nn.GRU(input_size=embedding_dim, hidden_size=hidden_dim, num_layers=num_layers, batch_first=True)
        self.attention = AricateAttentionLayer(hidden_dim)
        self.hidden_to_vocab = nn.Linear(hidden_dim * 2, vocab_size)

    def forward(self, input_words):
        embeds = self.word_embeddings(input_words)
        rnn_out, hn = self.rnn(embeds)
        last_hidden = hn[-1]
        context_vector = self.attention(rnn_out, last_hidden)
        combined_features = torch.cat((context_vector, last_hidden), dim=1)
        logits = self.hidden_to_vocab(combined_features)
        return logits

# --- D. Fonction de Génération (Adaptée) ---
def generate_sequence(model, tokenizer, prompt, max_length, temperature=1.0, do_sample=False):
    """Génère la continuation en utilisant le décodage glouton ou l'échantillonnage avec température."""
    model.eval()

    # Récupérer la longueur maximale d'entrée pour le padding
    max_len_input = tokenizer.max_len_input
    eos_id = tokenizer.special_tokens['<eos>']

    current_ids = tokenizer.encode(prompt, add_eos=False)

    print(f"\n--- Génération ({'Échantillonnée' if do_sample else 'Gloutonne'}) ---")
    print(f"Amorce: '{prompt}'")
    if do_sample:
        print(f"Température: {temperature}")

    with torch.no_grad():
        for _ in range(max_length):
            # 1. Préparer l'entrée (Longueur fixe, padding à gauche)
            input_ids_to_pad = current_ids[-max_len_input:] if len(current_ids) > max_len_input else current_ids
            padding_needed = max_len_input - len(input_ids_to_pad)
            input_ids_padded = [tokenizer.special_tokens['<pad>']] * padding_needed + input_ids_to_pad

            input_tensor = torch.tensor(input_ids_padded).unsqueeze(0)

            # Assurez-vous que le modèle est sur le bon appareil (CPU)
            device = next(model.parameters()).device
            input_tensor = input_tensor.to(device)

            # 2. Forward Pass et Prédiction
            logits = model(input_tensor)

            if do_sample:
                # Appliquer la température et échantillonner
                probabilities = F.softmax(logits / temperature, dim=-1)
                predicted_id = torch.multinomial(probabilities, num_samples=1).item()
            else:
                # Décoding glouton (choix du mot le plus probable)
                predicted_id = torch.argmax(logits, dim=-1).item()

            # 3. Arrêt
            if predicted_id == eos_id:
                break

            # 4. Ajouter le nouveau mot
            current_ids.append(predicted_id)

    # Décodage de la partie générée seulement
    prompt_length = len(tokenizer.encode(prompt))
    generated_ids = current_ids[prompt_length:]
    final_response = tokenizer.decode(generated_ids)

    print(f"Continuation générée: '{final_response}'")
    print("-" * 40)

    return final_response


# --- 2. Code de Chargement et de Test ---
def load_and_test_aricate_model(repo_id: str):
    """Charge le modèle Aricate V4 et son tokenizer depuis Hugging Face et lance l'inférence."""

    # Fichiers à télécharger
    MODEL_FILE = "model.safetensors"
    TOKENIZER_FILE = "aricate_tokenizer.json"

    # Force l'utilisation du CPU
    device = torch.device("cpu")
    print(f"Appareil d'inférence sélectionné: {device}")

    try:
        print(f"\n↻ 1. Chargement du modèle {repo_id}...")

        # ⚠️ Utilisez model_file_name pour spécifier le fichier de poids du modèle
        model = AricateModel.from_pretrained(repo_id, model_file_name=MODEL_FILE)
        model.to(device)
        model.eval()

        print("✅ Modèle Aricate V4 chargé avec succès.")

        print(f"\n↻ 2. Chargement du tokenizer custom...")
        # Téléchargement du fichier de vocabulaire
        tokenizer_path = hf_hub_download(repo_id=repo_id, filename=TOKENIZER_FILE)

        # Chargement des données du tokenizer
        with open(tokenizer_path, 'r', encoding='utf-8') as f:
            tokenizer_data = json.load(f)

        tokenizer = WordTokenizer(tokenizer_data)
        print("✅ Tokenizer custom chargé.")


        # 3. Test d'Inférence
        print("\n--- 3. TEST D'INFÉRENCE ---")

        # Longueur maximale de la continuation
        MAX_GENERATION_LENGTH = 10

        # --- Test en mode Glouton (Greedy) ---
        print("\n--- Test en mode Glouton (Greedy) ---")
        prompts_to_test_greedy = [
            "le boulanger donnait",
            "Marie allait chercher du",
            "il aimait cuisiner"
        ]
        for prompt in prompts_to_test_greedy:
            generate_sequence(model, tokenizer, prompt, MAX_GENERATION_LENGTH, do_sample=False)

        # --- Test en mode Échantillonnage (Sampling) ---
        print(f"\n--- Test en mode Échantillonnage (Sampling, Température=0.7) ---")
        prompts_to_test_sampled = [
            "le boulanger cuisinait des gâteaux qui",
            "Marie et sa mère allaient voir",
            "un met extrêmement délicieux pour"
        ]
        TEMPERATURE_VALUE = 0.7 # Une valeur courante pour l'échantillonnage
        for prompt in prompts_to_test_sampled:
            generate_sequence(model, tokenizer, prompt, MAX_GENERATION_LENGTH, temperature=TEMPERATURE_VALUE, do_sample=True)

    except Exception as e:
        print(f"\n❌ ÉCHEC DU CHARGEMENT OU DE L'INFÉRENCE :")
        print("Ceci peut se produire si le modèle custom n'est pas correctement publié ou si les classes ne sont pas définies.")
        print(f"Détail de l'erreur: {e}")


# --- Lancement de la fonction principale ---
if __name__ == '__main__':
    # Le REPO_ID que vous avez fourni
    ARICATE_REPO_ID = "Clemylia/Mini-Groutouille"
    load_and_test_aricate_model(ARICATE_REPO_ID)