Optimize performance: use lighter models and pre-cache models in Dockerfile

Dockerfile

@@ -8,6 +8,10 @@ COPY requirements.txt .
 # Install dependencies
 RUN pip install --no-cache-dir -r requirements.txt
 
+# Pre-download models to cache them in the Docker image
+RUN python -c "from transformers import pipeline; pipeline('summarization', model='sshleifer/distilbart-cnn-6-6')"
+RUN python -c "from transformers import pipeline; pipeline('translation', model='facebook/m2m100_418M')"
+
 # Copy the rest of the application
 COPY . .
 

app.py

@@ -27,13 +27,10 @@ async def serve_html():
 
 # Load pre-trained models from Hugging Face Hub
 try:
-    summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
-    translators = {
-        "fr": pipeline("translation", model="Helsinki-NLP/opus-mt-en-fr"),
-        "de": pipeline("translation", model="Helsinki-NLP/opus-mt-en-de"),
-        "es": pipeline("translation", model="Helsinki-NLP/opus-mt-en-es"),
-        "ar": pipeline("translation", model="Helsinki-NLP/opus-mt-en-ar")
-    }
+    # Use a lighter summarization model
+    summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-6-6")
+    # Use a single multilingual translation model
+    translator = pipeline("translation", model="facebook/m2m100_418M")
 except Exception as e:
     print(f"Error loading models: {str(e)}")
     raise e
@@ -60,19 +57,26 @@ async def translate(request: Request):
         data = await request.json()
         text = data["text"]
         lang = data["lang"]
-        if lang not in translators:
+        # Map language codes to M2M100 language codes
+        lang_map = {
+            "fr": "fr",
+            "de": "de",
+            "es": "es",
+            "ar": "ar"
+        }
+        if lang not in lang_map:
             return {"error": "Language not supported"}
         
-        # Perform translation
-        result = translators[lang](text)
+        # Perform translation using M2M100
+        result = translator(text, src_lang="en", tgt_lang=lang_map[lang])
         print(f"Translation result for {lang}: {result}")  # Debugging
         
         # Check if result is a list and has at least one item
         if not isinstance(result, list) or len(result) == 0:
             return {"error": "Translation failed: empty or invalid result"}
         
-        # Extract the translated text (correct key is "translation_text")
-        translation = result[0].get("translation_text")
+        # Extract the translated text
+        translation = result[0]["translation_text"]
         if translation is None:
             return {"error": "Translation failed: 'translation_text' not found in result"}