Create app.py

app.py

@@ -0,0 +1,413 @@
+# app.py
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+import collections
+from datasets import load_dataset
+from huggingface_hub import PyTorchModelHubMixin, HfApi, login
+import os
+import time
+import json
+import heapq
+from safetensors.torch import save_file as save_safetensors_file
+import gradio as gr
+import sys # Pour la redirection de la sortie
+
+# ==============================================================================
+#                      ARCHITECTURE ARICATE V4 (Intégrée)
+# ==============================================================================
+
+# --- A. WordTokenizer ---
+class WordTokenizer:
+    """Tokenizer simple pour l'architecture Aricate."""
+    def __init__(self, texts):
+        all_words = []
+        for text in texts:
+            # S'assurer que 'text' est une chaîne de caractères avant de l'opérer
+            if isinstance(text, str):
+                words = text.lower().split()
+                all_words.extend(words)
+
+        word_counts = collections.Counter(all_words)
+        sorted_words = [word for word, count in word_counts.most_common()]
+
+        self.special_tokens = {
+            '<pad>': 0,
+            '<unk>': 1,
+            '<eos>': 2,
+            '<sep>': 3,
+        }
+
+        self.word_to_id = self.special_tokens.copy()
+        next_id = len(self.special_tokens)
+
+        for word in sorted_words:
+            if word not in self.word_to_id:
+                self.word_to_id[word] = next_id
+                next_id += 1
+
+        self.id_to_word = {id: word for word, id in self.word_to_id.items()}
+        self.vocab_size = len(self.word_to_id)
+        print(f"Tokenisation effectuée. Taille du vocabulaire : {self.vocab_size}")
+
+    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.word_to_id['<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>', '<sep>'])
+
+# --- B. AricateAttentionLayer ---
+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 ---
+class AricateModel(nn.Module, PyTorchModelHubMixin):
+    """Architecture Aricate V4. Hérite de PyTorchModelHubMixin pour la sauvegarde et la publication."""
+    def __init__(self, vocab_size: int, embedding_dim: int, hidden_dim: int, num_layers: int = 1, config: dict = None):
+        super(AricateModel, self).__init__()
+
+        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 (Simplifiée pour l'espace) ---
+# NOTE: J'ai retiré la fonction de génération pour ne pas alourdir l'application Gradio principale et me concentrer sur l'entraînement/publication.
+# Dans un Space, il est préférable d'avoir une démo séparée après l'entraînement.
+# Je garde le Dataset car c'est nécessaire.
+
+# --- Nouvelle Classe PyTorch Dataset ---
+class AricateDataset(Dataset):
+    """Dataset personnalisé pour PyTorch."""
+    def __init__(self, X_data, Y_data):
+        self.X = X_data
+        self.Y = Y_data
+
+    def __len__(self):
+        return len(self.X)
+
+    def __getitem__(self, idx):
+        return self.X[idx], self.Y[idx]
+
+# ==============================================================================
+#                      FONCTION D'ENTRAÎNEMENT ADAPTÉE POUR GRADIO
+# ==============================================================================
+
+def train_aricate_model(
+    hf_token: str, 
+    hf_user: str, 
+    dataset_name: str, 
+    question_col: str, 
+    response_col: str, 
+    model_name: str, 
+    num_epochs: int
+):
+    """
+    Fonction principale d'entraînement adaptée pour Gradio.
+    
+    Elle prend les entrées de l'utilisateur, configure Aricate v4,
+    lance l'entraînement et publie le modèle sur Hugging Face.
+    """
+    
+    # Rediriger la sortie standard vers la console Gradio
+    sys.stdout.flush() 
+    print(f"\n{'='*50}\n>>> DÉBUT DU PROCESSUS D'ENTRAÎNEMENT Aricat v4 <<<\n{'='*50}")
+
+    try:
+        # --- 0. Configuration & Connexion Hugging Face ---
+        # Paramètres fixes (peuvent être ajustés si nécessaire)
+        EMBEDDING_DIM = 64
+        HIDDEN_DIM = 128
+        NUM_LAYERS = 2
+        BATCH_SIZE = 128 
+        LEARNING_RATE = 0.005
+        
+        # Configuration de l'appareil
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        print(f"Appareil d'entraînement sélectionné: {device}")
+        
+        # Connexion Hugging Face via le token
+        login(token=hf_token, add_to_git_credential=False)
+        REPO_ID = f"{hf_user}/{model_name}"
+        print(f"Connexion Hugging Face établie. Le modèle sera publié sous le dépôt: {REPO_ID}")
+
+        print(f"--- Lancement de l'Entraînement du SLM '{model_name}' (Aricate) ---")
+
+        # 1. Préparation des données
+        DATASET_SPLIT = 'train'
+        print(f"Chargement de la dataset '{dataset_name}' (split '{DATASET_SPLIT}')...")
+        try:
+            dataset = load_dataset(dataset_name, split=DATASET_SPLIT)
+        except Exception as e:
+            raise ValueError(f"Erreur lors du chargement de la dataset '{dataset_name}'. Vérifiez le nom du dépôt. Erreur: {e}")
+
+        # Construction du corpus en utilisant les colonnes spécifiées par l'utilisateur
+        try:
+            corpus_raw = [f"{ex[question_col]} <sep> {ex[response_col]}" for ex in dataset]
+        except KeyError as e:
+            raise KeyError(f"Colonne introuvable dans la dataset. Vérifiez les noms de colonnes : {e}. Les colonnes de votre dataset sont : {dataset.column_names}")
+            
+        tokenizer = WordTokenizer(corpus_raw)
+
+        train_data_X = []
+        train_data_Y = []
+
+        for item in dataset:
+            q = item[question_col]
+            r = item[response_col]
+            full_seq_ids = tokenizer.encode(f"{q} <sep> {r}", add_eos=True)
+            for i in range(1, len(full_seq_ids)):
+                X = full_seq_ids[:i]
+                Y = full_seq_ids[i]
+                train_data_X.append(X)
+                train_data_Y.append(Y)
+
+        max_len = max(len(x) for x in train_data_X)
+        padded_X = []
+        for x in train_data_X:
+            padding_needed = max_len - len(x)
+            # Ajout du padding au DÉBUT de la séquence (convention de certains modèles pour l'alignement)
+            padded_X.append([tokenizer.special_tokens['<pad>']] * padding_needed + x) 
+
+        X_train_tensor = torch.tensor(padded_X)
+        Y_train_tensor = torch.tensor(train_data_Y)
+        VOCAB_SIZE = tokenizer.vocab_size
+
+        print(f"Dataset chargée. Nombre de paires d'entraînement: {len(Y_train_tensor)}")
+        print(f"Taille du vocabulaire total: {VOCAB_SIZE}")
+        print(f"Longueur maximale d'entrée (max_len): {max_len}")
+
+        aricate_dataset = AricateDataset(X_train_tensor, Y_train_tensor)
+        train_loader = DataLoader(
+            dataset=aricate_dataset,
+            batch_size=BATCH_SIZE,
+            shuffle=True,
+            num_workers=0 # Mis à 0 pour éviter des problèmes de multi-processus sur certains environnements HF Space
+        )
+        print(f"Nombre de batches par époque : {len(train_loader)}")
+
+        # 2. Initialisation du Modèle
+        model_config = {
+            "vocab_size": VOCAB_SIZE,
+            "embedding_dim": EMBEDDING_DIM,
+            "hidden_dim": HIDDEN_DIM,
+            "num_layers": NUM_LAYERS
+        }
+        model = AricateModel(**model_config).to(device)
+        loss_function = nn.CrossEntropyLoss()
+        optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
+
+        # 3. Entraînement
+        print(f"\nDébut de l'entraînement pour {num_epochs} époques avec un BATCH_SIZE de {BATCH_SIZE}...")
+        start_time = time.time()
+
+        for epoch in range(num_epochs):
+            model.train()
+            total_loss = 0.0
+            
+            for batch_X, batch_Y in train_loader:
+                batch_X, batch_Y = batch_X.to(device), batch_Y.to(device)
+
+                optimizer.zero_grad()
+                logits = model(batch_X)
+                loss = loss_function(logits, batch_Y)
+                loss.backward()
+                optimizer.step()
+                total_loss += loss.item() * batch_X.size(0)
+
+            avg_loss = total_loss / len(aricate_dataset)
+            
+            # Mise à jour immédiate du statut
+            yield f"Entraînement en cours... Époque [{epoch+1}/{num_epochs}], Perte Moyenne: {avg_loss:.4f}"
+
+        end_time = time.time()
+        yield f"Entraînement terminé ! Durée: {(end_time - start_time):.2f}s. Début de la publication..."
+        print(f"\nEntraînement terminé ! Durée: {(end_time - start_time):.2f}s 🎉")
+
+
+        # 4. Sauvegarde et Publication sur Hugging Face
+        model.to("cpu")
+        print("\n" + "="*50)
+        print(">>> SAUVEGARDE ET PUBLICATION SUR HUGGING FACE <<<")
+        print("="*50)
+
+        save_directory = f"./{model_name}_local_save"
+        os.makedirs(save_directory, exist_ok=True)
+
+        model.save_pretrained(save_directory)
+        print(f"Modèle sauvegardé localement dans: {save_directory}")
+
+        tokenizer_path = os.path.join(save_directory, "aricate_tokenizer.txt")
+        with open(tokenizer_path, 'w', encoding='utf-8') as f:
+            json.dump(tokenizer.word_to_id, f, ensure_ascii=False)
+        print(f"Tokenizer (vocabulaire) sauvegardé dans: {tokenizer_path}")
+
+        # Publication
+        model.push_to_hub(
+            repo_id=REPO_ID,
+            commit_message=f"Modèle entraîné via Aricate v4 Space. Époques: {num_epochs}",
+            config=model_config
+        )
+        HfApi().upload_file(
+            path_or_fileobj=tokenizer_path,
+            path_in_repo="aricate_tokenizer.txt",
+            repo_id=REPO_ID,
+            repo_type="model",
+            commit_message="Update Aricate custom tokenizer vocabulary."
+        )
+
+        final_message = f"\n✅ Publication réussie ! Votre modèle '{model_name}' est disponible sur : https://huggingface.co/{REPO_ID}"
+        print(final_message)
+        yield final_message # Message final pour l'interface Gradio
+
+    except Exception as e:
+        error_message = f"\n❌ ERREUR CRITIQUE. L'entraînement ou la publication a échoué. Détail: {e}"
+        print(error_message)
+        yield error_message # Message d'erreur pour l'interface Gradio
+
+
+# ==============================================================================
+#                             INTERFACE GRADIO
+# ==============================================================================
+
+# Description détaillée pour l'utilisateur
+description = """
+# 🧠 Entraînez votre propre SLM avec Aricate v4 (Clemylia)
+Bienvenue sur l'interface d'entraînement d'Aricate v4 ! Suivez les étapes ci-dessous pour créer et publier votre propre Small Language Model (SLM) basé sur votre dataset personnalisée.
+
+**Étapes à suivre :**
+
+1.  **Authentification :** Entrez votre Token et Nom d'utilisateur Hugging Face. **Le token doit avoir la permission "Write" (Écriture).**
+2.  **Dataset :** Fournissez le nom du dépôt Hugging Face contenant votre dataset.
+3.  **Colonnes :** Indiquez les noms exacts des colonnes pour les questions et les réponses (par défaut : `question` et `reponse`).
+4.  **Nom du Modèle :** Choisissez le nom de votre futur modèle (il sera publié sous `votre_nom_utilisateur/nom_du_modèle`).
+5.  **Hyperparamètres :** Définissez le nombre d'époques.
+6.  **Lancement :** Appuyez sur le bouton et attendez la fin de l'entraînement et de la publication !
+"""
+
+# Création des blocs d'interface
+with gr.Blocks(title="Aricate v4 Trainer") as demo:
+    gr.Markdown(description)
+
+    # --- Section d'Authentification et de Publication ---
+    with gr.Row():
+        hf_token_input = gr.Textbox(
+            label="1. Token d'Accès Hugging Face (avec permission 'Write')",
+            type="password",
+            placeholder="hf_xxxxxxxxxxxxxxxxxxxxxxxxxx",
+            info="Token pour l'authentification et la publication (NE PAS PARTAGER !)"
+        )
+        hf_user_input = gr.Textbox(
+            label="2. Votre Nom d'Utilisateur Hugging Face",
+            placeholder="Clemylia",
+            info="Le modèle sera publié sur ce compte."
+        )
+    
+    # --- Section Dataset ---
+    gr.Markdown("### 🔍 Configuration de la Dataset")
+    with gr.Row():
+        dataset_name_input = gr.Textbox(
+            label="3. Nom du Dépôt Dataset (ex: Clemylia/Melta-revive)",
+            placeholder="le_nom_de_votre_dataset",
+            info="Dépôt public Hugging Face (il doit avoir un split 'train')."
+        )
+        question_col_input = gr.Textbox(
+            label="4. Nom de la Colonne 'Question'",
+            value="question",
+            placeholder="question",
+            info="Nom exact de la colonne contenant les questions."
+        )
+        response_col_input = gr.Textbox(
+            label="5. Nom de la Colonne 'Réponse'",
+            value="reponse",
+            placeholder="reponse",
+            info="Nom exact de la colonne contenant les réponses."
+        )
+
+    # --- Section Modèle et Hyperparamètres ---
+    gr.Markdown("### ⚙️ Configuration du Modèle et Entraînement")
+    with gr.Row():
+        model_name_input = gr.Textbox(
+            label="6. Nom Final du Modèle (sur Hugging Face)",
+            placeholder="mon-super-slm-aricate",
+            info="Sera publié comme 'utilisateur/nom-final'."
+        )
+        num_epochs_input = gr.Slider(
+            label="7. Nombre d'Époques d'Entraînement",
+            minimum=1,
+            maximum=50,
+            step=1,
+            value=10,
+            info="Plus d'époques = plus long, mais peut donner de meilleurs résultats (attention à l'overfitting)."
+        )
+
+    # --- Bouton et Sortie ---
+    train_button = gr.Button("🚀 Entraîner mon propre SLM avec Aricate v4", variant="primary")
+    
+    # Zone de sortie pour afficher la progression et les messages
+    output_log = gr.Textbox(
+        label="Console d'Entraînement et Log de Publication",
+        lines=15,
+        autoscroll=True,
+        interactive=False
+    )
+    
+    # Lien entre le bouton et la fonction Python
+    train_button.click(
+        fn=train_aricate_model, 
+        inputs=[
+            hf_token_input, 
+            hf_user_input, 
+            dataset_name_input, 
+            question_col_input, 
+            response_col_input, 
+            model_name_input, 
+            num_epochs_input
+        ],
+        outputs=output_log
+    )
+
+# Lancement de l'application Gradio
+if __name__ == "__main__":
+    demo.launch()