Update app.py

app.py

@@ -31,7 +31,7 @@ logging.basicConfig(
 logger = logging.getLogger("cosmic_ai")
 
 # Set a custom NLTK data directory
-nltk_data_dir = os.getenv('NLTK_DATA_DIR', '/tmp/nltk_data')
+nltk_data_dir = os.getenv('NLTK_DATA', '/tmp/nltk_data')
 os.makedirs(nltk_data_dir, exist_ok=True)
 nltk.data.path.append(nltk_data_dir)
 
@@ -131,10 +131,14 @@ def load_model(task: str, model_name: str = None):
             return vqa_function
         
         # Use pipeline for summarization, image-to-text, and file-qa
-        return pipeline(task if task != "file-qa" else "question-answering", model=model_to_load)
+        return pipeline(
+            task if task != "file-qa" else "question-answering",
+            model=model_to_load,
+            tokenizer_kwargs={"clean_up_tokenization_spaces": True}  # Suppress warning
+        )
         
     except Exception as e:
-        logger.error(f"Model load failed: {str(e)}")
+        logger.error(f"Model load failed for {task}: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Model loading failed: {task} - {str(e)}")
 
 def get_gemini_response(user_input: str, is_generation: bool = False):
@@ -170,8 +174,21 @@ def translate_text(text: str, target_language: str):
         
         lang_code = SUPPORTED_LANGUAGES[target_lang]
         
+        # Load translation model on demand if not pre-loaded
         if translation_model is None or translation_tokenizer is None:
-            raise Exception("Translation model not initialized")
+            logger.info("Translation model not pre-loaded, loading on demand...")
+            model_name = MODELS["translation"]
+            translation_model = M2M100ForConditionalGeneration.from_pretrained(
+                model_name,
+                cache_dir=os.getenv("HF_HOME", "/app/cache")
+            )
+            translation_tokenizer = M2M100Tokenizer.from_pretrained(
+                model_name,
+                cache_dir=os.getenv("HF_HOME", "/app/cache")
+            )
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            translation_model.to(device)
+            logger.info("Translation model loaded on demand successfully")
         
         match = re.search(r'how to say\s+(.+?)\s+in\s+(\w+)', text.lower())
         if match:
@@ -829,29 +846,44 @@ async def list_models():
 
 @app.on_event("startup")
 async def startup_event():
-    """Pre-load models at startup with timeout"""
+    """Pre-load models at startup with timeout and fallback"""
     global translation_model, translation_tokenizer
     logger.info("Starting model pre-loading...")
     
-    async def load_model_with_timeout(task):
+    async def load_model_with_timeout(task, model_name=None):
         try:
-            await asyncio.wait_for(asyncio.to_thread(load_model, task), timeout=60.0)
-            logger.info(f"Successfully loaded {task} model")
+            await asyncio.wait_for(
+                asyncio.to_thread(load_model, task, model_name),
+                timeout=60.0
+            )
+            logger.info(f"Successfully pre-loaded {task} model")
         except asyncio.TimeoutError:
             logger.warning(f"Timeout loading {task} model - will load on demand")
         except Exception as e:
             logger.error(f"Error pre-loading {task}: {str(e)}")
     
+    # Load translation model separately with retry mechanism
     try:
         model_name = MODELS["translation"]
-        translation_model = M2M100ForConditionalGeneration.from_pretrained(model_name)
-        translation_tokenizer = M2M100Tokenizer.from_pretrained(model_name)
+        logger.info(f"Attempting to load translation model: {model_name}")
+        translation_model = M2M100ForConditionalGeneration.from_pretrained(
+            model_name,
+            cache_dir=os.getenv("HF_HOME", "/app/cache")
+        )
+        translation_tokenizer = M2M100Tokenizer.from_pretrained(
+            model_name,
+            cache_dir=os.getenv("HF_HOME", "/app/cache")
+        )
         device = "cuda" if torch.cuda.is_available() else "cpu"
         translation_model.to(device)
         logger.info("Translation model pre-loaded successfully")
     except Exception as e:
         logger.error(f"Error pre-loading translation model: {str(e)}")
+        # Fallback: Set to None and load on demand
+        translation_model = None
+        translation_tokenizer = None
     
+    # Pre-load other models concurrently
     await asyncio.gather(
         load_model_with_timeout("summarization"),
         load_model_with_timeout("image-to-text"),
@@ -862,5 +894,4 @@ async def startup_event():
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)
-
+    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)