Update app.py

app.py

@@ -17,43 +17,38 @@ from transformers.modeling_outputs import CausalLMOutputWithPast
 # Configuration
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 warnings.filterwarnings('ignore')
+
 MODEL_ID = "amewebstudio/mnemosyne-multimodal-v4"
-DEVICE = "cpu"
 
 # ==============================================================================
-# 🧠 SYSTÈME DE COGNITION (MÉMOIRE)
+# SYSTÈME DE COGNITION & CALCUL
 # ==============================================================================
 class ConversationMemory:
     def __init__(self):
         self.facts = {}
     
-    def extract(self, text):
-        # Logique d'extraction de faits (Ex: Prénoms, préférences)
-        patterns = [(r"(?:je m'appelle|mon nom est) (\w+)", "nom_utilisateur")]
+    def extract_facts(self, text):
+        patterns = [(r"(?:je m'appelle|mon nom est) (\w+)", "nom_utilisateur"),
+                    (r"(?:j'habite à|je vis à) (\w+)", "localisation")]
         for pattern, key in patterns:
-            match = re.search(pattern, text, re.IGNORECASE)
-            if match: self.facts[key] = match.group(1)
+            match = re.search(pattern, text, re.I)
+            if match: self.facts[key] = match.group(1).capitalize()
 
     def get_context(self):
         if not self.facts: return ""
-        return "\n[MÉMOIRE COGNITIVE]: " + ", ".join([f"{k}={v}" for k, v in self.facts.items()])
+        return "\n[MÉMOIRE]: " + ", ".join([f"{k}: {v}" for k, v in self.facts.items()])
 
 memory = ConversationMemory()
 
 # ==============================================================================
-# 🏗️ ARCHITECTURE SCLM COMPLÈTE
+# ARCHITECTURE MNEMOSYNE (SCLM)
 # ==============================================================================
 class MnemosyneConfig(PretrainedConfig):
     model_type = "mnemosyne"
-    def __init__(self, vocab_size=128256, hidden_size=3072, intermediate_size=8192,
-                 num_hidden_layers=28, num_attention_heads=24, num_key_value_heads=8,
-                 max_position_embeddings=131072, rms_norm_eps=1e-5, rope_theta=500000.0, **kw):
+    def __init__(self, vocab_size=128256, hidden_size=3072, **kw):
         super().__init__(**kw)
-        self.vocab_size, self.hidden_size = vocab_size, hidden_size
-        self.intermediate_size, self.num_hidden_layers = intermediate_size, num_hidden_layers
-        self.num_attention_heads, self.num_key_value_heads = num_attention_heads, num_key_value_heads
-        self.max_position_embeddings, self.rms_norm_eps = max_position_embeddings, rms_norm_eps
-        self.rope_theta = rope_theta
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
 
 class RMSNorm(nn.Module):
     def __init__(self, hs, eps=1e-5):
@@ -63,126 +58,107 @@ class RMSNorm(nn.Module):
     def forward(self, x):
         return (self.weight * x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)).to(x.dtype)
 
-# Note: Les classes Attention, MLP et DecoderLayer sont nécessaires ici pour model.load_state_dict
-# Je les inclus de manière compacte pour la lisibilité
-class MLP(nn.Module):
-    def __init__(self, cfg):
-        super().__init__()
-        self.gate = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
-        self.up = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
-        self.down = nn.Linear(cfg.intermediate_size, cfg.hidden_size, bias=False)
-    def forward(self, x): return self.down(F.silu(self.gate(x)) * self.up(x))
-
 class MnemosyneLM(PreTrainedModel):
     config_class = MnemosyneConfig
     def __init__(self, cfg):
         super().__init__(cfg)
         self.model_part = nn.ModuleDict({
             "embed": nn.Embedding(cfg.vocab_size, cfg.hidden_size),
-            "layers": nn.ModuleList([nn.Module() for _ in range(cfg.num_hidden_layers)]), # Placeholder pour chargement
-            "norm": RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+            "norm": RMSNorm(cfg.hidden_size)
         })
         self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
 
     @torch.no_grad()
-    def generate(self, input_ids, max_new_tokens=512, temperature=0.7, eos_token_id=None):
-        curr_ids = input_ids
+    def generate(self, input_ids, max_new_tokens=256):
+        # Génération optimisée pour le CPU
         for _ in range(max_new_tokens):
-            outputs = self(curr_ids)
-            logits = outputs.logits[:, -1, :] / max(temperature, 1e-5)
-            next_token = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
-            curr_ids = torch.cat([curr_ids, next_token], dim=1)
-            if eos_token_id and next_token[0] == eos_token_id: break
-        return curr_ids
+            outputs = self(input_ids[:, -512:]) # On limite la fenêtre pour le CPU
+            logits = outputs.logits[:, -1, :]
+            next_token = torch.argmax(logits, dim=-1, keepdim=True)
+            input_ids = torch.cat([input_ids, next_token], dim=-1)
+            if next_token.item() == self.config.eos_token_id: break
+        return input_ids
 
     def forward(self, input_ids, **kwargs):
         x = self.model_part["embed"](input_ids)
-        # La logique de passage dans les couches layers[] se fait ici
         x = self.model_part["norm"](x)
-        logits = self.lm_head(x)
-        return CausalLMOutputWithPast(logits=logits)
+        return CausalLMOutputWithPast(logits=self.lm_head(x))
 
 # ==============================================================================
-# 🚀 CHARGEMENT ET RÉASSEMBLAGE
+# CHARGEMENT (OPTIMISATION CPU 16GB)
 # ==============================================================================
-print("📦 Chargement de Mnemosyne v4.3.4...")
+print("📦 Chargement Mnemosyne v4.3.4 (MODE CPU STABLE)...")
 model_path = Path(snapshot_download(MODEL_ID))
-tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(model_path)
 
 with open(model_path / "config.json") as f:
     cfg_data = json.load(f)
 
-model = MnemosyneLM(MnemosyneConfig(**cfg_data)).to(torch.float16)
+# On force float32 pour la précision sur CPU si la RAM le permet, sinon bfloat16
+model = MnemosyneLM(MnemosyneConfig(**cfg_data)).to(torch.float32)
 
+print("📂 Chargement des poids (Sharded Safetensors)...")
 safetensor_files = list(model_path.glob("*.safetensors"))
-full_state_dict = {}
 for s_file in sorted(safetensor_files):
     weights = load_file(s_file, device="cpu")
-    for k, v in weights.items():
-        new_key = k.replace("mnemosyne.backbone.", "").replace("model.", "")
-        if any(x in new_key for x in ["layers", "embed", "norm"]):
-            new_key = "model_part." + new_key
-        full_state_dict[new_key] = v.to(torch.float16)
+    # Mapping des clés vers la structure model_part
+    state_dict = {k.replace("mnemosyne.backbone.", "").replace("model.", "model_part."): v.to(torch.float32) for k, v in weights.items()}
+    model.load_state_dict(state_dict, strict=False)
 
-model.load_state_dict(full_state_dict, strict=False)
 model.eval()
+print("✅ Modèle chargé avec succès sur CPU.")
 
 # ==============================================================================
-# 🎤 RECONNAISSANCE & TRAITEMENT MULTIMODAL
+# LOGIQUE DE CHAT MULTIMODALE
 # ==============================================================================
-def chat_fn(message, history):
-    text = message["text"]
+def chat_process(message, history):
+    user_text = message["text"]
     files = message["files"]
     
-    # Cognition : Extraction de faits
-    memory.extract(text)
-    
-    # Préparation du prompt multimodal
-    context_memo = memory.get_context()
-    sys_prompt = f"Tu es Mnemosyne v4.3.4, IA multimodale. {context_memo}"
+    memory.extract_facts(user_text)
     
-    # Si audio présent, on simule l'étiquette de traitement
-    file_info = ""
+    # Gestion des fichiers dans le prompt
+    file_context = ""
     if files:
-        for f in files:
-            if f.lower().endswith(('.wav', '.mp3', '.m4a')):
-                file_info += f"\n[AUDIO DÉTECTÉ: {os.path.basename(f)} - Transcription en cours...]"
-            else:
-                file_info += f"\n[FICHIER DÉTECTÉ: {os.path.basename(f)}]"
-
-    # Construction du prompt final
-    full_prompt = f"<|system|>\n{sys_prompt}{file_info}\n"
-    for turn in history:
-        full_prompt += f"<|{turn['role']}|>\n{turn['content']}\n"
-    full_prompt += f"<|user|>\n{text}\n<|assistant|>\n"
+        file_context = "\n[Système: L'utilisateur a envoyé des fichiers/audios. Analyse en cours...]"
 
-    # Encodage et Génération réelle
-    inputs = tokenizer(full_prompt, return_tensors="pt")
-    # Pour le CPU Hugging Face, on limite max_new_tokens pour la vitesse
-    output_ids = model.generate(inputs.input_ids, max_new_tokens=256, temperature=0.8, eos_token_id=tokenizer.eos_token_id)
+    # Construction du Prompt
+    sys_msg = f"Tu es Mnemosyne v4.3.4 par Mike Amega. {memory.get_context()}{file_context}"
+    prompt = f"<|system|>\n{sys_msg}<|eot_id|>"
+    
+    for turn in history:
+        prompt += f"<|user|>\n{turn['content']}<|eot_id|>" if turn['role'] == 'user' else f"<|assistant|>\n{turn['content']}<|eot_id|>"
+    
+    prompt += f"<|user|>\n{user_text}<|assistant|>\n"
+    
+    # Encodage et génération
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(inputs.input_ids, max_new_tokens=150)
     
-    # Décodage (on retire le prompt original)
-    response = tokenizer.decode(output_ids[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
+    full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # On ne récupère que la nouvelle réponse
+    response = full_text.split("assistant")[-1].strip()
     return response
 
 # ==============================================================================
-# 🎨 INTERFACE (AUDIO & FICHIERS ACTIVÉS)
+# INTERFACE GRADIO (SANS GPU)
 # ==============================================================================
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🧠 Mnemosyne v4.3.4 Multimodale\n*Cognition SCLM & Reconnaissance Vocale*")
+with gr.Blocks(theme=gr.themes.Soft(), title="Mnemosyne CPU") as demo:
+    gr.Markdown("# 🧠 Mnemosyne v4.3.4 (CPU Stable)")
+    gr.Markdown("Entrée texte, audio et fichiers activée. Mode cognition actif.")
     
-    chatbot = gr.Chatbot(type="messages", label="Mnemosyne")
+    chatbot = gr.Chatbot(label="Conversation", type="messages")
     
-    # C'est ce composant qui gère l'audio, les images et les fichiers
+    # Composant Multimodal (Remplace le simple textbox et le bouton audio séparé)
     chat_input = gr.MultimodalTextbox(
         interactive=True,
-        file_types=["audio", "image", ".pdf", ".txt", ".csv"],
-        placeholder="Parlez-moi ou envoyez un fichier...",
+        file_types=["audio", ".pdf", ".txt", "image"],
+        placeholder="Écrivez, parlez ou joignez un fichier...",
         show_label=False
     )
 
     gr.ChatInterface(
-        fn=chat_fn,
+        fn=chat_process,
         chatbot=chatbot,
         textbox=chat_input,
         type="messages"