Create compile_model.py

compile_model.py

@@ -0,0 +1,73 @@
+# compile_model.py
+"""
+With transformers, model downloading and caching is automatic.
+This script just performs a warmup to:
+1. Download model during Docker build
+2. Compile CUDA kernels
+3. Verify installation
+"""
+import os
+import logging
+from datetime import datetime
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Import shared model manager
+from shared_models import get_shared_llama
+
+HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
+
+def warmup_model():
+    """
+    Warmup Llama-3.2-3B model:
+    - Downloads from HuggingFace Hub
+    - Loads with 4-bit quantization
+    - Runs test inference
+    """
+    logger.info("="*60)
+    logger.info("LLAMA-3.2-3B WARMUP")
+    logger.info("="*60)
+    
+    try:
+        # Get shared model instance
+        llama = get_shared_llama()
+        
+        # This triggers model download and loading
+        logger.info("Running warmup inference...")
+        
+        test_response = llama.generate(
+            system_prompt="You are a helpful educational assistant.",
+            user_message="Hello, this is a test warmup message.",
+            max_tokens=20,
+            temperature=0.7,
+        )
+        
+        logger.info(f"✅ Warmup successful")
+        logger.info(f"   Response preview: {test_response[:80]}...")
+        
+        # Get model info
+        info = llama.get_model_info()
+        logger.info("="*60)
+        logger.info("MODEL INFO")
+        logger.info("="*60)
+        for key, value in info.items():
+            logger.info(f"  {key}: {value}")
+        
+        logger.info("="*60)
+        logger.info("✅ MODEL READY FOR PRODUCTION")
+        logger.info("="*60)
+        
+        return True
+        
+    except Exception as e:
+        logger.error(f"❌ Warmup failed: {e}")
+        import traceback
+        traceback.print_exc()
+        return False
+
+
+if __name__ == "__main__":
+    success = warmup_model()
+    if not success:
+        exit(1)