Finisha-LLM/

Nac-1

# Installation des bibliothèques nécessaires
!pip install huggingface_hub

import torch
import torch.nn as nn
import json
import os
from huggingface_hub import hf_hub_download

# ==============================================================================
# 0. CONFIGURATION GLOBALE
# ==============================================================================

# Identifiant du dépôt (à remplacer par le vôtre si différent)
REPO_ID = "Clemylia/Nac-1" 
FINAL_FILENAME = 'nac1_final_weights.pth'
VOCAB_FILENAME = "nac1_vocab.json"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
HIDDEN_SIZE = 256 # Doit correspondre à la taille utilisée lors de l'entraînement

# Tokens Spéciaux (doivent être les mêmes que ceux utilisés dans le vocabulaire)
SOS_token = 0 
EOS_token = 1

# ==============================================================================
# 1. CLASSES D'ARCHITECTURE DU MODÈLE (L'utilisateur en a besoin)
# ==============================================================================

# Note : Les utilisateurs doivent copier ces définitions exactes pour charger les poids.

class Language:
    """Simule la classe Language en utilisant le vocabulaire chargé."""
    def __init__(self, name, word2index, index2word):
        self.name = name
        self.word2index = word2index
        self.index2word = index2word
        self.n_words = len(word2index)

class EncoderRNN(nn.Module):
    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.hidden_size = hidden_size
        self.embedding = nn.Embedding(input_size, hidden_size)
        self.lstm = nn.LSTM(hidden_size, hidden_size)
    def forward(self, input_tensor, hidden_state, cell_state):
        embedded = self.embedding(input_tensor).view(1, 1, -1)
        output, (hidden_state, cell_state) = self.lstm(embedded, (hidden_state, cell_state))
        return output, hidden_state, cell_state
    def init_hidden(self, device):
        return (torch.zeros(1, 1, self.hidden_size, device=device),
                torch.zeros(1, 1, self.hidden_size, device=device))

class DecoderRNN(nn.Module):
    def __init__(self, hidden_size, output_size):
        super().__init__()
        self.hidden_size = hidden_size
        self.embedding = nn.Embedding(output_size, hidden_size)
        self.lstm = nn.LSTM(hidden_size, hidden_size)
        self.out = nn.Linear(hidden_size, output_size)
        self.softmax = nn.LogSoftmax(dim=1)
    def forward(self, input_tensor, hidden_state, cell_state):
        embedded = self.embedding(input_tensor).view(1, 1, -1)
        output, (hidden_state, cell_state) = self.lstm(embedded, (hidden_state, cell_state))
        output = self.softmax(self.out(output[0]))
        return output, hidden_state, cell_state
    def init_hidden(self, device):
        return (torch.zeros(1, 1, self.hidden_size, device=device),
                torch.zeros(1, 1, self.hidden_size, device=device))


# ==============================================================================
# 2. FONCTIONS DE TÉLÉCHARGEMENT ET UTILITAIRES
# ==============================================================================

def load_vocab_and_models(repo_id, vocab_filename, weights_filename, hidden_size, device):
    """Télécharge les fichiers et initialise les modèles."""
    print(f"Téléchargement des fichiers depuis {repo_id}...")
    
    # 1. Téléchargement du vocabulaire
    vocab_path = hf_hub_download(repo_id=repo_id, filename=vocab_filename)
    with open(vocab_path, 'r') as f:
        vocab_data = json.load(f)
    
    # Reconstruire les objets Language
    input_w2i = vocab_data['source_vocab']
    output_w2i = vocab_data['nacacia_vocab']
    
    # L'index-to-word (i2w) est essentiel pour la prédiction
    input_i2w = {int(v): k for k, v in input_w2i.items()}
    output_i2w = {int(v): k for k, v in output_w2i.items()}

    input_lang = Language("Source", input_w2i, input_i2w)
    output_lang = Language("Nacacia", output_w2i, output_i2w)
    
    # 2. Initialisation des modèles
    encoder = EncoderRNN(input_lang.n_words, hidden_size).to(device)
    decoder = DecoderRNN(hidden_size, output_lang.n_words).to(device)
    
    # 3. Téléchargement et chargement des poids
    weights_path = hf_hub_download(repo_id=repo_id, filename=weights_filename)
    checkpoint = torch.load(weights_path, map_location=device)
    
    encoder.load_state_dict(checkpoint['encoder_state_dict'])
    decoder.load_state_dict(checkpoint['decoder_state_dict'])
    
    encoder.eval() # Mode évaluation
    decoder.eval()
    
    print("✅ Modèle Nac-1 chargé et prêt pour l'inférence.")
    return encoder, decoder, input_lang, output_lang

def tensor_from_sentence(lang, sentence):
    """Convertit une phrase en un tenseur d'indices de mots."""
    # S'assurer que tous les mots sont connus
    try:
        indexes = [lang.word2index[word] for word in sentence.split(' ')]
    except KeyError as e:
        raise ValueError(f"Le mot '{e.args[0]}' n'est pas dans le vocabulaire connu du modèle. Échec.")
        
    indexes.append(EOS_token)
    return torch.tensor(indexes, dtype=torch.long, device=device).view(-1, 1)

def evaluate(encoder, decoder, sentence, input_lang, output_lang, max_length=15):
    """Génère la phrase Nacacia à partir d'une phrase source."""
    with torch.no_grad():
        try:
            input_tensor = tensor_from_sentence(input_lang, sentence.lower().strip())
        except ValueError as e:
            return str(e) # Retourne l'erreur de vocabulaire

        input_length = input_tensor.size(0)
        encoder_hidden, encoder_cell = encoder.init_hidden(device)

        # ENCODEUR
        for ei in range(input_length):
            _, encoder_hidden, encoder_cell = encoder(input_tensor[ei], encoder_hidden, encoder_cell)

        # DÉCODEUR
        decoder_input = torch.tensor([[SOS_token]], device=device)
        decoder_hidden, decoder_cell = encoder_hidden, encoder_cell
        decoded_words = []

        for _ in range(max_length):
            decoder_output, decoder_hidden, decoder_cell = decoder(decoder_input, decoder_hidden, decoder_cell)
            topv, topi = decoder_output.data.topk(1)
            predicted_index = topi.item()
            
            if predicted_index == EOS_token: break
            
            decoded_words.append(output_lang.index2word[predicted_index])
            decoder_input = torch.tensor([[predicted_index]], device=device)
        
        return ' '.join(decoded_words)

# ==============================================================================
# 3. EXÉCUTION POUR L'UTILISATEUR
# ==============================================================================

if __name__ == "__main__":
    
    # 1. Chargement du modèle depuis Hugging Face
    encoder, decoder, input_lang, output_lang = load_vocab_and_models(
        REPO_ID, VOCAB_FILENAME, FINAL_FILENAME, HIDDEN_SIZE, device
    )

    print("\n--- Inférence sur de Nouvelles Phrases ---")
    
    # Phrases de test utilisant le vocabulaire connu :
    test_sentences = [
        "il fait beau aujourd'hui", 
        "la lune est claire", 
        "j'ai besoin d'aide",
        "le chat dort sur le canapé"
    ]
    
    # Phrase avec un mot inconnu pour tester la robustesse :
    test_sentences.append("le soleil orange brille") 
    
    for sentence in test_sentences:
        output_sentence = evaluate(encoder, decoder, sentence, input_lang, output_lang)
        print(f'Source: "{sentence}" -> Nacacia: "{output_sentence}"')