push atlas first

Dockerfile

@@ -42,6 +42,10 @@ RUN python -c "from transformers import pipeline; pipeline('text-to-speech', mod
  && python -c "from transformers import pipeline; pipeline('text-to-speech', model='facebook/mms-tts-eng')" \
  && python -c "from transformers import pipeline; pipeline('text-to-speech', model='facebook/mms-tts-yor')"
 
+# Pre-load N-ATLaS model during build
+RUN python -c "from huggingface_hub import snapshot_download; snapshot_download(repo_id='NCAIR1/N-ATLaS')" \
+ && python -c "from transformers import AutoTokenizer, AutoModelForCausalLM; import torch; tokenizer = AutoTokenizer.from_pretrained('NCAIR1/N-ATLaS'); model = AutoModelForCausalLM.from_pretrained('NCAIR1/N-ATLaS', torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, device_map='auto' if torch.cuda.is_available() else None, trust_remote_code=True, low_cpu_mem_usage=True, use_cache=True); print('N-ATLaS model loaded successfully')"
+
 # Copy project files
 COPY . .
 

app.py

@@ -59,6 +59,10 @@ def load_models():
         logger.warning("Please set HF_TOKEN environment variable to access restricted models.")
     else:
         logger.info("HF_TOKEN is set and ready for authenticated model access.")
+    
+    logger.info("Loading N-ATLaS language identification model...")
+    _load_natlas()
+    
     logger.info("Loading TTS models...")
     try:
         tts_ha = pipeline("text-to-speech", model="facebook/mms-tts-hau", device=device)
@@ -82,8 +86,6 @@ def load_models():
     tts_ig = None
     logger.info("Igbo TTS model disabled - will return text responses for Igbo language")
 
-    logger.info("N-ATLaS language identification model will be lazy-loaded on first use")
-
     logger.info("Deferred ASR model loads: will lazy-load per language on first use")
 
 def _get_asr(lang_code: str):
@@ -199,6 +201,7 @@ IGBO_WORDS = [
 def _load_natlas():
     global natlas_tokenizer, natlas_model
     if natlas_tokenizer is not None and natlas_model is not None:
+        logger.info("N-ATLaS model already loaded")
         return True
     
     hf_token = os.getenv("HF_TOKEN")
@@ -210,8 +213,8 @@ def _load_natlas():
         return False
     
     try:
-        logger.info("Lazy-loading N-ATLaS language identification model...")
-        logger.info("This may take a few minutes as the model loads its shards...")
+        logger.info("Loading N-ATLaS language identification model...")
+        logger.info("Model files are pre-cached from Docker build, loading should be faster...")
         
         natlas_tokenizer = AutoTokenizer.from_pretrained("NCAIR1/N-ATLaS", token=hf_token)
         natlas_model = AutoModelForCausalLM.from_pretrained(