Update app.py

app.py

@@ -10,15 +10,47 @@ from transformers import (
 )
 from functools import lru_cache
 
-# Initialize models with caching
+# ==================== NEW: PULAR TO FRENCH TRANSLATOR ====================
 @lru_cache(maxsize=1)
-def load_translator():
+def load_pular_to_french():
+    """Load the Pular-to-French translator model"""
+    print("Loading Pular→French translator model...")
+    model_name = "mlamined/pl_fr_104"  # Your new checkpoint
+    
+    try:
+        # Load with NLLB tokenizer for proper language codes
+        tokenizer = NllbTokenizer.from_pretrained(
+            "facebook/nllb-200-distilled-600M", 
+            src_lang="fuv_Latn",  # Pular source
+            tgt_lang="fra_Latn"   # French target
+        )
+        
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+        
+        translator = pipeline(
+            "translation",
+            model=model,
+            tokenizer=tokenizer,
+            src_lang="fuv_Latn",
+            tgt_lang="fra_Latn",
+            max_length=256,
+            num_beams=3,
+            early_stopping=True
+        )
+        print("Pular→French translator model loaded successfully!")
+        return translator
+    except Exception as e:
+        print(f"Error loading Pular→French translator: {e}")
+        return None
+
+# ==================== EXISTING MODELS ====================
+@lru_cache(maxsize=1)
+def load_french_to_pular():
     """Load the French-to-Pular translator model"""
-    print("Loading translator model...")
+    print("Loading French→Pular translator model...")
     model_name = "mlamined/fr_pl_130"
     
     try:
-        # Load with NLLB tokenizer for proper language codes
         tokenizer = NllbTokenizer.from_pretrained(
             "facebook/nllb-200-distilled-600M", 
             src_lang="fra_Latn", 
@@ -37,10 +69,10 @@ def load_translator():
             num_beams=3,
             early_stopping=True
         )
-        print("Translator model loaded successfully!")
+        print("French→Pular translator model loaded successfully!")
         return translator
     except Exception as e:
-        print(f"Error loading translator: {e}")
+        print(f"Error loading French→Pular translator: {e}")
         return None
 
 @lru_cache(maxsize=1)
@@ -74,14 +106,114 @@ def load_llm():
         print(f"Error loading LLM: {e}")
         return None, None
 
-# Load models
-translator = load_translator()
-llm_model, llm_tokenizer = load_llm()
+# ==================== LOAD ALL MODELS ====================
+print("\n" + "="*60)
+print("🚀 LOADING ALL MODELS")
+print("="*60)
+
+translator_pular_to_french = load_pular_to_french()  # NEW
+translator_french_to_pular = load_french_to_pular()  # EXISTING
+llm_model, llm_tokenizer = load_llm()  # EXISTING
 
 # Check if models loaded
 use_llm = llm_model is not None and llm_tokenizer is not None
 
-# Improved system prompt with clearer instructions
+# ==================== TRANSLATION FUNCTIONS ====================
+def translate_pular_to_french(pular_text):
+    """Translate Pular text to French"""
+    if not translator_pular_to_french:
+        return "Erreur: Modèle Pular→Français non disponible."
+    
+    if not pular_text or len(pular_text.strip()) == 0:
+        return ""
+    
+    try:
+        # Clean the Pular text
+        clean_pular = pular_text.strip()
+        clean_pular = re.sub(r'\s+', ' ', clean_pular)
+        clean_pular = clean_pular[:300]  # Limit length
+        
+        print(f"Translating Pular→French: {clean_pular[:100]}...")
+        
+        # Translate
+        result = translator_pular_to_french(clean_pular, max_length=256)
+        
+        # Extract translation
+        if isinstance(result, list) and len(result) > 0:
+            if isinstance(result[0], dict) and "translation_text" in result[0]:
+                french_text = result[0]["translation_text"]
+            elif isinstance(result[0], str):
+                french_text = result[0]
+            else:
+                french_text = str(result[0])
+        elif isinstance(result, dict) and "translation_text" in result:
+            french_text = result["translation_text"]
+        elif isinstance(result, str):
+            french_text = result
+        else:
+            return "Erreur de traduction. Veuillez réessayer."
+        
+        # Clean the French response
+        french_text = re.sub(r'\*.*?\*', '', french_text)
+        french_text = re.sub(r'\[.*?\]|\(.*?\)', '', french_text)
+        french_text = re.sub(r'\s+', ' ', french_text).strip()
+        
+        print(f"Translated to French: {french_text[:100]}...")
+        return french_text
+    
+    except Exception as e:
+        print(f"Pular→French translation error: {e}")
+        return "Erreur technique lors de la traduction."
+
+def translate_french_to_pular(french_text):
+    """Translate French text to Pular"""
+    if not translator_french_to_pular:
+        return "Hakkunde ndee, mi wadataa."
+    
+    if not french_text or len(french_text.strip()) == 0:
+        return ""
+    
+    try:
+        # Clean the French text
+        clean_french = french_text.strip()
+        clean_french = re.sub(r'\*+', '', clean_french)
+        clean_french = re.sub(r'\s+', ' ', clean_french)
+        clean_french = clean_french[:300]  # Limit length
+        
+        print(f"Translating French→Pular: {clean_french[:100]}...")
+        
+        # Translate
+        result = translator_french_to_pular(clean_french, max_length=256)
+        
+        # Extract translation
+        if isinstance(result, list) and len(result) > 0:
+            if isinstance(result[0], dict) and "translation_text" in result[0]:
+                pular_text = result[0]["translation_text"]
+            elif isinstance(result[0], str):
+                pular_text = result[0]
+            else:
+                pular_text = str(result[0])
+        elif isinstance(result, dict) and "translation_text" in result:
+            pular_text = result["translation_text"]
+        elif isinstance(result, str):
+            pular_text = result
+        else:
+            return "Hakkunde ndee, mi wadataa."
+        
+        # Clean the Pular response
+        pular_text = re.sub(r'\*.*?\*', '', pular_text)
+        pular_text = re.sub(r'\bFinsitaare\b.*', '', pular_text)
+        pular_text = re.sub(r'\[.*?\]|\(.*?\)', '', pular_text)
+        pular_text = re.sub(r'\s+', ' ', pular_text).strip()
+        
+        print(f"Translated to Pular: {pular_text[:100]}...")
+        return pular_text
+    
+    except Exception as e:
+        print(f"French→Pular translation error: {e}")
+        return "Hakkunde ndee, tontu kadi."
+
+# ==================== EXISTING FUNCTIONS (UNCHANGED) ====================
 system_prompt = """You are a helpful assistant . Use simple, clear language as if explaining to a young child. Provide accurate and relevant responses. Answer in French, and keep responses short and friendly. Maintenant, réponds aux questions suivantes:"""
 
 def clean_french_response(text):
@@ -90,20 +222,18 @@ def clean_french_response(text):
         return ""
     
     # Remove markdown formatting
-    text = re.sub(r'\*+', '', text)  # Remove asterisks
-    text = re.sub(r'#+\s*', '', text)  # Remove headers
-    text = re.sub(r'`.*?`', '', text)  # Remove code blocks
-    text = re.sub(r'\[.*?\]\(.*?\)', '', text)  # Remove links
+    text = re.sub(r'\*+', '', text)
+    text = re.sub(r'#+\s*', '', text)
+    text = re.sub(r'`.*?`', '', text)
+    text = re.sub(r'\[.*?\]\(.*?\)', '', text)
     
     # Remove any gibberish or repeated patterns
     lines = text.split('\n')
     clean_lines = []
     for line in lines:
         line = line.strip()
-        # Skip empty lines or lines that look like gibberish
         if not line or len(line) < 3:
             continue
-        # Skip lines that are just punctuation or symbols
         if re.match(r'^[^a-zA-Z0-9\s]*$', line):
             continue
         clean_lines.append(line)
@@ -112,7 +242,7 @@ def clean_french_response(text):
     if clean_lines:
         response = clean_lines[0]
     else:
-        response = text[:200]  # Fallback
+        response = text[:200]
     
     # Ensure it ends with proper punctuation
     if response and not response[-1] in '.!?':
@@ -123,7 +253,6 @@ def clean_french_response(text):
 def generate_french_response(user_input, history=None):
     """Generate French response using the actual LLM with improved prompting"""
     if not use_llm:
-        # Fallback responses if LLM is not available
         fallback_responses = [
             "Je comprends votre question. Pouvez-vous la reformuler?",
             "Je vais chercher cette information pour vous.",
@@ -140,7 +269,6 @@ def generate_french_response(user_input, history=None):
         
         # Add conversation history if available (simplified)
         if history and len(history) > 0:
-            # Take only the last exchange for context
             recent = history[-2:] if len(history) >= 2 else history
             for msg in recent:
                 if msg["role"] == "user":
@@ -169,9 +297,9 @@ def generate_french_response(user_input, history=None):
         with torch.no_grad():
             outputs = llm_model.generate(
                 **inputs,
-                max_new_tokens=100,  # Reduced for cleaner responses
+                max_new_tokens=100,
                 do_sample=True,
-                temperature=0.5,  # Lower temperature for more focused responses
+                temperature=0.5,
                 top_p=0.9,
                 top_k=50,
                 pad_token_id=llm_tokenizer.pad_token_id,
@@ -185,11 +313,9 @@ def generate_french_response(user_input, history=None):
         
         # Extract only the assistant's response
         if "Réponse:" in response:
-            # Take everything after the last "Réponse:"
             parts = response.split("Réponse:")
             french_response = parts[-1].strip()
         else:
-            # If no "Réponse:" tag found, take everything after the prompt
             french_response = response[len(prompt):].strip()
         
         # Clean the response
@@ -200,68 +326,12 @@ def generate_french_response(user_input, history=None):
             french_response = "Je ne peux pas répondre à cette question pour le moment."
         
         print(f"Generated French response: {french_response[:150]}...")
-        return french_response[:250]  # Limit length
+        return french_response[:250]
     
     except Exception as e:
         print(f"Error generating French response: {e}")
         return "Je rencontre des difficultés techniques. Pouvez-vous reformuler votre question?"
 
-def translate_french_to_pular(french_text):
-    """Translate French text to Pular with improved cleaning"""
-    if not translator:
-        return "Erreur: Modèle de traduction non disponible."
-    
-    if not french_text or len(french_text.strip()) == 0:
-        return ""
-    
-    try:
-        # Clean the French text before translation
-        clean_french = french_text
-        
-        # Remove any remaining markdown
-        clean_french = re.sub(r'\*+', '', clean_french)
-        clean_french = re.sub(r'\s+', ' ', clean_french).strip()
-        
-        # Limit length
-        clean_french = clean_french[:300]
-        
-        print(f"Translating: {clean_french[:100]}...")
-        
-        # Translate
-        result = translator(clean_french, max_length=256)
-        
-        # Extract translation
-        if isinstance(result, list) and len(result) > 0:
-            if isinstance(result[0], dict) and "translation_text" in result[0]:
-                pular_text = result[0]["translation_text"]
-            elif isinstance(result[0], str):
-                pular_text = result[0]
-            else:
-                pular_text = str(result[0])
-        elif isinstance(result, dict) and "translation_text" in result:
-            pular_text = result["translation_text"]
-        elif isinstance(result, str):
-            pular_text = result
-        else:
-            return "Hakkunde ndee, mi wadataa."
-        
-        # Clean the Pular response
-        pular_text = re.sub(r'\*.*?\*', '', pular_text)
-        pular_text = re.sub(r'\bFinsitaare\b.*', '', pular_text)
-        pular_text = re.sub(r'\[.*?\]|\(.*?\)', '', pular_text)
-        pular_text = re.sub(r'\s+', ' ', pular_text).strip()
-        
-        # If translation seems problematic, provide a fallback
-        if len(pular_text) < 3 or pular_text.lower().startswith("error"):
-            pular_text = "Mi fahmina. Hakkunde ndee, tontu kadi."
-        
-        print(f"Translated to: {pular_text[:100]}...")
-        return pular_text
-    
-    except Exception as e:
-        print(f"Translation error: {e}")
-        return "Hakkunde ndee, tontu kadi."
-
 def chat_function(user_input, chat_history):
     """Main chat function with improved response handling"""
     if not user_input.strip():
@@ -297,16 +367,12 @@ def chat_function(user_input, chat_history):
         details = f"**Erreur technique:** Veuillez réessayer."
         return chat_history, details
 
-def direct_translate(text):
-    """Direct translation function"""
-    return translate_french_to_pular(text)
-
-# Create Gradio interface
+# ==================== GRADIO INTERFACE ====================
 with gr.Blocks(
-    title="🤖 Chatbot Français-Pular avec IA",
+    title="🤖 Chatbot Français-Pular avec IA - BIDIRECTIONNEL",
     theme=gr.themes.Soft(),
     css="""
-    .gradio-container {max-width: 800px; margin: auto;}
+    .gradio-container {max-width: 900px; margin: auto;}
     .chatbot {min-height: 400px;}
     .details-box {
         background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
@@ -327,12 +393,19 @@ with gr.Blocks(
         margin: 2px;
         font-size: 12px;
     }
+    .translation-card {
+        background: #f8f9fa;
+        padding: 15px;
+        border-radius: 10px;
+        border: 1px solid #dee2e6;
+        margin: 10px 0;
+    }
     """
 ) as demo:
     gr.Markdown("""
-    # 🇫🇷 → 🌍 Chatbot Français-Pular avec IA Réelle
+    # 🇫🇷 ↔ 🌍 Chatbot Français-Pular avec IA - BIDIRECTIONNEL
     
-    ### Un assistant intelligent qui comprend le français et répond en pular
+    ### Un assistant intelligent avec traduction dans les deux sens
     """)
     
     # Status indicators
@@ -344,7 +417,10 @@ with gr.Blocks(
                 <strong>🤖 Modèle IA (Gemma-2-2B):</strong> {'<span style="color: green;">✅ Actif</span>' if use_llm else '<span style="color: orange;">⚠️ Basique</span>'}
             </div>
             <div style='background: #e3f2fd; padding: 10px; border-radius: 5px; margin: 5px 0;'>
-                <strong>🔤 Traducteur (mlamined/fr_pl_130):</strong> {'<span style="color: green;">✅ Actif</span>' if translator else '<span style="color: red;">❌ Erreur</span>'}
+                <strong>🔤 Traducteur Pular→Français (mlamined/pl_fr_104):</strong> {'<span style="color: green;">✅ Actif</span>' if translator_pular_to_french else '<span style="color: red;">❌ Erreur</span>'}
+            </div>
+            <div style='background: #e3f2fd; padding: 10px; border-radius: 5px; margin: 5px 0;'>
+                <strong>🔤 Traducteur Français→Pular (mlamined/fr_pl_130):</strong> {'<span style="color: green;">✅ Actif</span>' if translator_french_to_pular else '<span style="color: red;">❌ Erreur</span>'}
             </div>
             <div style='background: #fff3e0; padding: 10px; border-radius: 5px; margin: 5px 0;'>
                 <strong>⚡ Performance:</strong> {'<span style="color: orange;">CPU</span>' if not torch.cuda.is_available() else '<span style="color: green;">GPU</span>'}
@@ -376,7 +452,6 @@ with gr.Blocks(
             with gr.Row():
                 clear_btn = gr.Button("🗑️ Effacer", variant="secondary", size="sm")
                 show_details = gr.Checkbox(label="📋 Afficher les détails", value=True)
-                # Empty column for spacing
                 gr.Column(scale=4, min_width=0)
             
             details_output = gr.Markdown(
@@ -412,10 +487,6 @@ with gr.Blocks(
             def clear_chat():
                 return [], gr.update(value="", visible=False)
             
-            # Simple function to set the example text
-            def set_example_text(example):
-                return example
-            
             # Connect events
             msg.submit(
                 respond,
@@ -433,9 +504,8 @@ with gr.Blocks(
                 [chatbot, details_output]
             )
             
-            # Connect example buttons - SIMPLIFIED
+            # Connect example buttons
             for i, btn in enumerate(example_buttons):
-                # First set the example text
                 btn.click(
                     fn=lambda ex=examples[i]: ex,
                     inputs=None,
@@ -446,70 +516,134 @@ with gr.Blocks(
                     outputs=[msg, chatbot, details_output]
                 )
         
-        with gr.TabItem("🔤 Traducteur", id="translate"):
-            gr.Markdown("### Traduction directe français → pular")
+        with gr.TabItem("🔤 Traducteur Bidirectionnel", id="translate"):
+            gr.Markdown("""
+            ### Traduction dans les deux sens
+            **🇫🇷 Français → 🌍 Pular** et **🌍 Pular → 🇫🇷 Français**
+            """)
+            
             with gr.Row():
+                # French to Pular translation
                 with gr.Column():
-                    french_input = gr.Textbox(
+                    gr.Markdown("#### 🇫🇷 → 🌍 Français vers Pular")
+                    french_input_ftop = gr.Textbox(
                         label="Texte français",
-                        placeholder="Entrez du texte à traduire...",
-                        lines=5
+                        placeholder="Entrez du texte français à traduire en pular...",
+                        lines=4
                     )
                     with gr.Row():
-                        translate_btn = gr.Button("Traduire ➤", variant="primary")
-                        clear_translate = gr.Button("Effacer", variant="secondary")
-                
-                with gr.Column():
+                        translate_fr_to_pl = gr.Button("Traduire 🇫🇷→🌍", variant="primary")
+                        clear_fr_to_pl = gr.Button("Effacer", variant="secondary")
                     pular_output = gr.Textbox(
                         label="Traduction pular",
-                        lines=5,
+                        lines=4,
+                        interactive=False
+                    )
+                
+                # Pular to French translation (NEW)
+                with gr.Column():
+                    gr.Markdown("#### 🌍 → 🇫🇷 Pular vers Français")
+                    pular_input_ptof = gr.Textbox(
+                        label="Texte pular",
+                        placeholder="Entrez du texte pular à traduire en français...",
+                        lines=4
+                    )
+                    with gr.Row():
+                        translate_pl_to_fr = gr.Button("Traduire 🌍→🇫🇷", variant="primary")
+                        clear_pl_to_fr = gr.Button("Effacer", variant="secondary")
+                    french_output = gr.Textbox(
+                        label="Traduction française",
+                        lines=4,
                         interactive=False
                     )
             
-            translate_btn.click(
-                direct_translate,
-                inputs=french_input,
+            # Connect buttons
+            # French to Pular
+            translate_fr_to_pl.click(
+                translate_french_to_pular,
+                inputs=french_input_ftop,
                 outputs=pular_output
             )
-            french_input.submit(
-                direct_translate,
-                inputs=french_input,
+            french_input_ftop.submit(
+                translate_french_to_pular,
+                inputs=french_input_ftop,
                 outputs=pular_output
             )
-            clear_translate.click(
+            clear_fr_to_pl.click(
                 lambda: ("", ""),
                 None,
-                [french_input, pular_output]
+                [french_input_ftop, pular_output]
             )
             
-            gr.Markdown("### 📝 Exemples rapides")
-            translate_examples = gr.Examples(
-                examples=[
-                    ["Bonjour, je m'appelle Mamadou et je suis guinéen."],
-                    ["L'éducation est la clé du développement d'un pays."],
-                    ["La culture guinéenne est riche et diversifiée."],
-                    ["Nous devons protéger notre environnement pour les générations futures."],
-                    ["La paix et la stabilité sont essentielles pour le progrès."]
-                ],
-                inputs=french_input,
-                outputs=pular_output,
-                fn=direct_translate,
-                cache_examples=True,
-                label="Cliquez sur un exemple"
+            # Pular to French (NEW)
+            translate_pl_to_fr.click(
+                translate_pular_to_french,
+                inputs=pular_input_ptof,
+                outputs=french_output
+            )
+            pular_input_ptof.submit(
+                translate_pular_to_french,
+                inputs=pular_input_ptof,
+                outputs=french_output
             )
+            clear_pl_to_fr.click(
+                lambda: ("", ""),
+                None,
+                [pular_input_ptof, french_output]
+            )
+            
+            gr.Markdown("### 📝 Exemples rapides")
+            
+            with gr.Row():
+                # French to Pular examples
+                with gr.Column():
+                    gr.Markdown("**Exemples Français→Pular:**")
+                    fr_to_pl_examples = gr.Examples(
+                        examples=[
+                            ["Bonjour, je m'appelle Mamadou et je suis guinéen."],
+                            ["L'éducation est la clé du développement d'un pays."],
+                            ["La culture guinéenne est riche et diversifiée."]
+                        ],
+                        inputs=french_input_ftop,
+                        outputs=pular_output,
+                        fn=translate_french_to_pular,
+                        cache_examples=True,
+                        label="Cliquez sur un exemple"
+                    )
+                
+                # Pular to French examples (NEW)
+                with gr.Column():
+                    gr.Markdown("**Exemples Pular→Français:**")
+                    pl_to_fr_examples = gr.Examples(
+                        examples=[
+                            ["On jaaraama musee Alpha."],
+                            ["Miɗo weelaa."],
+                            ["Jannde ko saabi fii ɓantal leydi."]
+                        ],
+                        inputs=pular_input_ptof,
+                        outputs=french_output,
+                        fn=translate_pular_to_french,
+                        cache_examples=True,
+                        label="Cliquez sur un exemple"
+                    )
     
     gr.Markdown("---")
     gr.Markdown("""
     ### ℹ️ À propos de ce système
     
+    **Nouveautés:**
+    - ✅ **Traduction Pular→Français** ajoutée (mlamined/pl_fr_104)
+    - 🔄 **Traduction bidirectionnelle** complète
+    - 🚀 **Deux modèles de traduction** indépendants
+    
     **Fonctionnement:**
-    1. Vous écrivez en français
-    4. Vous recevez une réponse naturelle et contextuelle
+    1. Vous écrivez en français ou en pular
+    2. Le système traduit dans la direction choisie
+    3. Pour le chat: français → IA → pular
     
     **Capacités:**
     - Réponses intelligentes et contextuelles
-    - Traduction précise français → pular
-    - Maintien du contexte de conversation
+    - Traduction précise dans les deux sens
     - Interface intuitive et facile à utiliser
     
     **Note:** Les réponses peuvent prendre quelques secondes à générer sur CPU.
@@ -517,9 +651,10 @@ with gr.Blocks(
 
 if __name__ == "__main__":
     print("=" * 60)
-    print("🚀 Démarrage du Chatbot Français-Pular")
+    print("🚀 DÉMARRAGE DU CHATBOT BIDIRECTIONNEL")
     print(f"📊 Statut LLM: {'✅ Prêt' if use_llm else '❌ Échec'}")
-    print(f"📊 Statut traducteur: {'✅ Prêt' if translator else '❌ Échec'}")
+    print(f"📊 Statut traducteur Pular→Français: {'✅ Prêt' if translator_pular_to_french else '❌ Échec'}")
+    print(f"📊 Statut traducteur Français→Pular: {'✅ Prêt' if translator_french_to_pular else '❌ Échec'}")
     print(f"⚡ Matériel: {'GPU' if torch.cuda.is_available() else 'CPU'}")
     print("=" * 60)
     
@@ -529,5 +664,4 @@ if __name__ == "__main__":
         share=True,
         debug=False,
         show_error=True
-    )
-
+    )