Update app.py

app.py

@@ -1,97 +1,155 @@
 import gradio as gr
-import spaces
 from transformers import pipeline
+import torch
+import logging
+import spaces
 from typing import Union, List
 
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+if torch.cuda.is_available():
+    device = "cuda"
+    logger.info("Using CUDA for inference.")
+elif torch.backends.mps.is_available():
+    device = "mps"
+    logger.info("Using MPS for inference.")
+else:
+    device = "cpu"
+    logger.info("Using CPU for inference.")
+
 class BambaraTranslator:
     def __init__(self, model_name: str = "sudoping01/nllb-bambara-v2"):
         self.translator = pipeline(
             "translation",
             model=model_name,
+            device=device,
             max_length=512,
-            truncation=True,
-            device=0 if spaces.config.Hardware.current == "cuda" else -1  
+            truncation=True
         )
         self.flores_codes = {
             "French": "fra_Latn",
             "English": "eng_Latn",
             "Bambara": "bam_Latn"
         }
+        logger.info("Translation pipeline initialized successfully.")
 
-    @spaces.GPU
     def translate(self, text: Union[str, List[str]], src_lang: str, tgt_lang: str) -> Union[str, List[str]]:
         source_lang = self.flores_codes[src_lang]
         target_lang = self.flores_codes[tgt_lang]
-        if isinstance(text, str):
-            translation = self.translator(text, src_lang=source_lang, tgt_lang=target_lang)
-            return str(translation[0]['translation_text'])
-        translations = self.translator(text, src_lang=source_lang, tgt_lang=target_lang)
-        return [str(t['translation_text']) for t in translations]
+        
+        logger.info(f"Translating text from {source_lang} to {target_lang}.")
+        
+        try:
+            if isinstance(text, str):
+                translation = self.translator(text, src_lang=source_lang, tgt_lang=target_lang, num_beams=2)
+                return str(translation[0]['translation_text'])
+            else:
+                translations = self.translator(text, src_lang=source_lang, tgt_lang=target_lang, num_beams=2)
+                return [str(t['translation_text']) for t in translations]
+        except Exception as e:
+            logger.error(f"Translation failed: {e}")
+            return "An error occurred during translation."
+
 
-# Initialize translator
 translator = BambaraTranslator()
 
-# Three examples (Bambara, French, English)
 examples = [
-    "An filɛ ni ye yɔrɔ minna ni an ye an sigi ka a layɛ yala an bɛ ka baara min kɛ ɛsike a kɛlen don ka Ɲɛ wa ? Bɛɛ ka kan ka i jɔyɔrɔ fa walasa an ka se ka taa Ɲɛ",
-    "Le Mali est un pays riche en culture mais confronté à de nombreux défis.",
-    "The sun rises every morning to bring light to the world."
+    ["An filɛ ni ye yɔrɔ minna ni an ye an sigi ka a layɛ yala an bɛ ka baara min kɛ ɛsike a kɛlen don ka Ɲɛ wa ? Bɛɛ ka kan ka i jɔyɔrɔ fa walasa an ka se ka taa Ɲɛ", "Bambara", "French"],
+    ["Le Mali est un pays riche en culture mais confronté à de nombreux défis.", "French", "Bambara"],
+    ["The sun rises every morning to bring light to the world.", "English", "Bambara"],
+    ["Good morning", "English", "Bambara"],
 ]
 
-@spaces.GPU
-def translate_text(src_lang, tgt_lang, text):
+@spaces.GPU()
+def translate_text(text: str, src_lang: str, tgt_lang: str) -> str:
+    """
+    Translate the input text from the source language to the target language.
+    """
     if not text.strip():
         return "Please enter text to translate."
+    
     if src_lang == tgt_lang:
         return "Source and target languages must be different."
+    
     try:
         result = translator.translate(text, src_lang, tgt_lang)
+        logger.info("Translation successful.")
         return result
     except Exception as e:
+        logger.error(f"Translation failed: {e}")
         return f"Error: {str(e)}"
 
-with gr.Blocks(title="Bambara Translator") as demo:
-    gr.Markdown(
-        """
-        # 🇲🇱 Bambara Translator
-        Translate between Bambara, French, and English instantly.
-        ## How to Use
-        1. Pick source and target languages
-        2. Enter text or select an example
-        3. Click "Translate" to see the result
-        """
-    )
-    
-    with gr.Row():
-        src_lang = gr.Dropdown(choices=["Bambara", "French", "English"], label="Source Language", value="Bambara")
-        tgt_lang = gr.Dropdown(choices=["Bambara", "French", "English"], label="Target Language", value="French")
-        with gr.Column():
-            text_input = gr.Textbox(label="Input Text", lines=5, placeholder="Type your text here...")
-            example_dropdown = gr.Dropdown(
-                choices=examples,
-                label="Examples",
-                value=examples[0]
-            )
-    
-    translate_btn = gr.Button("Translate", variant="primary")
-    output = gr.Textbox(label="Translation", lines=3)
-    
-    gr.Markdown(
-        """
-        By [sudoping01](https://huggingface.co/sudoping01), from [sudoping01/nllb-bambara-v2](https://huggingface.co/sudoping01/nllb-bambara-v2). Fine-tuned on Meta's NLLB, CC BY-NC 4.0, non-commercial.
-        """
-    )
+def build_interface():
+    """
+    Builds the Gradio interface for translating text between supported languages.
+    """
+    with gr.Blocks(title="Bambara Translator") as demo:
+        gr.Markdown(
+            """
+            # 🇲🇱 Bambara Translator
+            Translate between Bambara, French, and English instantly using NLLB model.
+            
+            ## How to Use
+            1. Select source and target languages from the dropdowns
+            2. Enter your text or choose from examples
+            3. Click "Translate" to see the result
+            """
+        )
+        
+        with gr.Row():
+            with gr.Column():
+                text_input = gr.Textbox(
+                    lines=5, 
+                    label="Text to Translate", 
+                    placeholder="Enter text here..."
+                )
+                
+                with gr.Row():
+                    src_lang = gr.Dropdown(
+                        choices=["Bambara", "French", "English"], 
+                        label="Source Language", 
+                        value="Bambara"
+                    )
+                    tgt_lang = gr.Dropdown(
+                        choices=["Bambara", "French", "English"], 
+                        label="Target Language", 
+                        value="French"
+                    )
+                
+                translate_btn = gr.Button("Translate", variant="primary")
+            
+            with gr.Column():
+                output = gr.Textbox(label="Translation", lines=5, interactive=False)
+        
+        # Examples section
+        gr.Examples(
+            examples=examples,
+            inputs=[text_input, src_lang, tgt_lang],
+            outputs=output,
+            fn=translate_text,
+            cache_examples=False
+        )
+        
+        gr.Markdown(
+            """
+            **Model:** [sudoping01/nllb-bambara-v2](https://huggingface.co/sudoping01/nllb-bambara-v2)  
+            **License:** CC BY-NC 4.0 (Non-commercial use)  
+            **Based on:** Meta's NLLB (No Language Left Behind)
+            """
+        )
+        
+        translate_btn.click(
+            fn=translate_text,
+            inputs=[text_input, src_lang, tgt_lang],
+            outputs=output
+        )
     
-    # Connect functions
-    translate_btn.click(
-        fn=translate_text,
-        inputs=[src_lang, tgt_lang, text_input],
-        outputs=output
-    )
-    example_dropdown.change(
-        fn=lambda x: x,
-        inputs=example_dropdown,
-        outputs=text_input
-    )
+    return demo
 
-demo.launch()
+if __name__ == "__main__":
+    logger.info("Starting the Gradio interface for the Bambara translator.")
+    interface = build_interface()
+    interface.launch()
+    logger.info("Gradio interface running.")