Update main.py

main.py

@@ -8,28 +8,48 @@ from pydantic import BaseModel
 from typing import Optional, Dict, Any
 import json
 import re
+from contextlib import asynccontextmanager
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from NatureLM.models import NatureLM
 from NatureLM.infer import Pipeline
 
-app = FastAPI(title="NatureLM Audio Analysis API")
+# Set up cache directories BEFORE importing any HuggingFace modules
+cache_base = "/app/.cache"
+os.environ['HF_HOME'] = cache_base
+os.environ['TRANSFORMERS_CACHE'] = f"{cache_base}/transformers"
+os.environ['HF_DATASETS_CACHE'] = f"{cache_base}/datasets"
+os.environ['HF_HUB_CACHE'] = f"{cache_base}/hub"
 
-# CORS middleware
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+# Ensure cache directories exist with proper permissions
+cache_dirs = [
+    cache_base,
+    os.environ['TRANSFORMERS_CACHE'],
+    os.environ['HF_DATASETS_CACHE'],
+    os.environ['HF_HUB_CACHE']
+]
 
-# Initialize NatureLM model
+for cache_dir in cache_dirs:
+    os.makedirs(cache_dir, exist_ok=True)
+    # Ensure write permissions
+    os.chmod(cache_dir, 0o755)
+
+# Initialize global variables
 model = None
 pipeline = None
 
 def load_model():
+    """Load the NatureLM model with proper error handling"""
     global model, pipeline
     try:
+        print("🔄 Loading NatureLM model...")
+        print(f"📁 Using cache directory: {os.environ.get('HF_HOME', cache_base)}")
+        
+        # Verify cache directories are writable
+        for cache_dir in cache_dirs:
+            if not os.access(cache_dir, os.W_OK):
+                raise PermissionError(f"Cache directory {cache_dir} is not writable")
+            print(f"✅ Cache directory {cache_dir} is writable")
+        
         # Load NatureLM-audio model from HuggingFace
         model = NatureLM.from_pretrained("EarthSpeciesProject/NatureLM-audio")
         model = model.eval()
@@ -41,12 +61,41 @@ def load_model():
         print("✅ NatureLM model loaded successfully")
     except Exception as e:
         print(f"❌ Error loading model: {e}")
+        print(f"🔍 Cache directory status:")
+        for cache_dir in cache_dirs:
+            exists = os.path.exists(cache_dir)
+            writable = os.access(cache_dir, os.W_OK) if exists else False
+            print(f"   {cache_dir}: {'✅' if exists and writable else '❌'} (exists: {exists}, writable: {writable})")
         raise e
 
-# Load model on startup
-@app.on_event("startup")
-async def startup_event():
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Lifespan context manager for FastAPI"""
+    # Startup
+    print("🚀 Starting up Animal Whisper AI Decoder...")
     load_model()
+    print("✅ Application startup complete")
+    
+    yield
+    
+    # Shutdown
+    print("🛑 Shutting down Animal Whisper AI Decoder...")
+
+app = FastAPI(
+    title="NatureLM Audio Analysis API",
+    description="AI-powered animal sound analysis using NatureLM",
+    version="1.0.0",
+    lifespan=lifespan
+)
+
+# CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
 class AnalysisResponse(BaseModel):
     species: str
@@ -249,47 +298,50 @@ async def analyze_audio(file: UploadFile = File(...)):
         # Combine results
         combined_response = " ".join(results)
         
-        # Extract information from response
+        # Extract information
         confidence_scores = extract_confidence_from_response(combined_response)
         species_info = extract_species_info(combined_response)
         
-        # Calculate overall confidence based on response quality
-        overall_confidence = max(
-            confidence_scores["model_confidence"],
-            confidence_scores["llama_confidence"],
-            75.0 if species_info["common_name"] else 50.0  # Higher confidence if species identified
-        )
-        
-        # Clean up temp file
+        # Clean up temporary file
         os.remove(temp_path)
         
         return AnalysisResponse(
             species=species_info["common_name"] or "Unknown species",
             interpretation=combined_response,
-            confidence=overall_confidence,
-            signal_type=species_info["signal_type"] or "Vocalization",
+            confidence=confidence_scores["model_confidence"] / 100.0,
+            signal_type=species_info["signal_type"] or "Unknown",
             common_name=species_info["common_name"] or "Unknown",
             scientific_name=species_info["scientific_name"] or "Unknown",
-            habitat=species_info["habitat"] or "Unknown habitat",
-            behavior=species_info["behavior"] or "Unknown behavior",
+            habitat=species_info["habitat"] or "Unknown",
+            behavior=species_info["behavior"] or "Unknown",
             audio_characteristics=audio_chars,
-            model_confidence=confidence_scores["model_confidence"],
-            llama_confidence=confidence_scores["llama_confidence"],
-            additional_insights=combined_response,
-            cluster_group="NatureLM Analysis"
+            model_confidence=confidence_scores["model_confidence"] / 100.0,
+            llama_confidence=confidence_scores["llama_confidence"] / 100.0,
+            additional_insights="Analysis completed successfully",
+            cluster_group="NatureLM"
         )
         
     except Exception as e:
-        # Clean up temp file if it exists
-        if os.path.exists(temp_path):
+        # Clean up temporary file if it exists
+        if 'temp_path' in locals() and os.path.exists(temp_path):
             os.remove(temp_path)
-        
         raise HTTPException(status_code=500, detail=f"Analysis failed: {str(e)}")
 
 @app.get("/health")
 async def health_check():
-    return {"status": "healthy", "model_loaded": model is not None}
+    """Health check endpoint"""
+    return {
+        "status": "healthy",
+        "model_loaded": model is not None,
+        "pipeline_ready": pipeline is not None,
+        "cache_directories": {
+            "hf_home": os.environ.get('HF_HOME'),
+            "transformers_cache": os.environ.get('TRANSFORMERS_CACHE'),
+            "datasets_cache": os.environ.get('HF_DATASETS_CACHE'),
+            "hub_cache": os.environ.get('HF_HUB_CACHE')
+        }
+    }
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=8000)