Update main.py

main.py

@@ -8,9 +8,9 @@ from pydantic import BaseModel
 from typing import Optional, Dict, Any
 import json
 import re
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-import soundfile as sf
-import io
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from NatureLM.models import NatureLM
+from NatureLM.infer import Pipeline
 
 app = FastAPI(title="NatureLM Audio Analysis API")
 
@@ -23,34 +23,24 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Initialize model and tokenizer
+# Initialize NatureLM model
 model = None
-tokenizer = None
-audio_pipeline = None
+pipeline = None
 
 def load_model():
-    global model, tokenizer, audio_pipeline
+    global model, pipeline
     try:
-        # Load the Llama model that NatureLM uses
-        model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            torch_dtype=torch.float16,
-            device_map="auto",
-            trust_remote_code=True
-        )
-        
-        # Load audio analysis pipeline
-        audio_pipeline = pipeline(
-            "audio-classification",
-            model="microsoft/wavlm-base",
-            return_all_scores=True
-        )
+        # Load NatureLM-audio model from HuggingFace
+        model = NatureLM.from_pretrained("EarthSpeciesProject/NatureLM-audio")
+        model = model.eval()
+        if torch.cuda.is_available():
+            model = model.cuda()
         
-        print("✅ Models loaded successfully")
+        # Initialize pipeline
+        pipeline = Pipeline(model=model)
+        print("✅ NatureLM model loaded successfully")
     except Exception as e:
-        print(f"❌ Error loading models: {e}")
+        print(f"❌ Error loading model: {e}")
         raise e
 
 # Load model on startup
@@ -114,7 +104,10 @@ def extract_species_info(response_text: str) -> Dict[str, str]:
     common_patterns = [
         r"common name[:\s]*([A-Za-z\s]+)",
         r"([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\s+\(common\)",
-        r"species[:\s]*([A-Za-z\s]+)"
+        r"species[:\s]*([A-Za-z\s]+)",
+        r"([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\s+treefrog",
+        r"([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\s+bird",
+        r"([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\s+mammal"
     ]
     
     for pattern in common_patterns:
@@ -144,7 +137,9 @@ def extract_species_info(response_text: str) -> Dict[str, str]:
         r"signal type[:\s]*([A-Za-z\s]+)",
         r"call type[:\s]*([A-Za-z\s]+)",
         r"vocalization[:\s]*([A-Za-z\s]+)",
-        r"sound type[:\s]*([A-Za-z\s]+)"
+        r"sound type[:\s]*([A-Za-z\s]+)",
+        r"([A-Za-z\s]+)\s+call",
+        r"([A-Za-z\s]+)\s+song"
     ]
     
     for pattern in signal_patterns:
@@ -157,7 +152,8 @@ def extract_species_info(response_text: str) -> Dict[str, str]:
     habitat_patterns = [
         r"habitat[:\s]*([A-Za-z\s,]+)",
         r"environment[:\s]*([A-Za-z\s,]+)",
-        r"found in[:\s]*([A-Za-z\s,]+)"
+        r"found in[:\s]*([A-Za-z\s,]+)",
+        r"lives in[:\s]*([A-Za-z\s,]+)"
     ]
     
     for pattern in habitat_patterns:
@@ -170,7 +166,8 @@ def extract_species_info(response_text: str) -> Dict[str, str]:
     behavior_patterns = [
         r"behavior[:\s]*([A-Za-z\s,]+)",
         r"purpose[:\s]*([A-Za-z\s,]+)",
-        r"function[:\s]*([A-Za-z\s,]+)"
+        r"function[:\s]*([A-Za-z\s,]+)",
+        r"used for[:\s]*([A-Za-z\s,]+)"
     ]
     
     for pattern in behavior_patterns:
@@ -181,12 +178,11 @@ def extract_species_info(response_text: str) -> Dict[str, str]:
     
     return info
 
-def analyze_audio_characteristics(audio_data: bytes) -> Dict[str, Any]:
+def analyze_audio_characteristics(audio_path: str) -> Dict[str, Any]:
     """Analyze audio characteristics using librosa"""
     try:
-        # Load audio from bytes
-        audio_bytes = io.BytesIO(audio_data)
-        y, sr = librosa.load(audio_bytes, sr=None)
+        # Load audio file
+        y, sr = librosa.load(audio_path, sr=None)
         
         # Calculate audio features
         duration = librosa.get_duration(y=y, sr=sr)
@@ -227,100 +223,51 @@ def analyze_audio_characteristics(audio_data: bytes) -> Dict[str, Any]:
         print(f"Error analyzing audio characteristics: {e}")
         return {}
 
-def classify_audio_signal(audio_data: bytes) -> Dict[str, Any]:
-    """Classify audio using WavLM model"""
-    try:
-        # Convert audio to the format expected by WavLM
-        audio_bytes = io.BytesIO(audio_data)
-        y, sr = librosa.load(audio_bytes, sr=16000)  # WavLM expects 16kHz
-        
-        # Reshape for the pipeline
-        audio_input = {"array": y, "sampling_rate": sr}
-        
-        # Get classification results
-        results = audio_pipeline(audio_input)
-        
-        # Extract the most likely class and confidence
-        if results and len(results) > 0:
-            top_result = results[0]
-            return {
-                "signal_type": top_result.get("label", "Unknown"),
-                "confidence": top_result.get("score", 0.0) * 100
-            }
-        
-        return {"signal_type": "Unknown", "confidence": 0.0}
-    except Exception as e:
-        print(f"Error in audio classification: {e}")
-        return {"signal_type": "Unknown", "confidence": 0.0}
-
 @app.post("/analyze", response_model=AnalysisResponse)
 async def analyze_audio(file: UploadFile = File(...)):
     try:
-        # Read file content
-        content = await file.read()
+        # Save uploaded file temporarily
+        temp_path = f"/tmp/{file.filename}"
+        with open(temp_path, "wb") as buffer:
+            content = await file.read()
+            buffer.write(content)
         
         # Analyze audio characteristics
-        audio_chars = analyze_audio_characteristics(content)
+        audio_chars = analyze_audio_characteristics(temp_path)
         
-        # Classify audio signal
-        signal_classification = classify_audio_signal(content)
+        # Create multiple queries for comprehensive analysis
+        queries = [
+            "What is the common name for the focal species in the audio? Answer:",
+            "What type of vocalization or call is this? Answer:",
+            "Describe the habitat and behavior context of this species. Answer:",
+            "Provide a detailed analysis of this animal sound including species identification, call type, and behavioral context. Answer:"
+        ]
         
-        # Create enhanced prompt for Llama model
-        enhanced_prompt = f"""
-        You are an expert in animal vocalization analysis. Analyze this audio recording and provide detailed information.
-
-        Audio file: {file.filename}
-        Duration: {audio_chars.get('duration_seconds', 'Unknown')} seconds
-        Sample rate: {audio_chars.get('sample_rate', 'Unknown')} Hz
-        Tempo: {audio_chars.get('tempo_bpm', 'Unknown')} BPM
-        Signal classification: {signal_classification.get('signal_type', 'Unknown')}
-
-        Please provide a comprehensive analysis including:
-        1. Species identification (common name and scientific name if possible)
-        2. Signal type and purpose (mating call, alarm, territorial, etc.)
-        3. Habitat and behavior context
-        4. Confidence level in your assessment (0-100%)
-
-        Format your response with clear sections for each aspect.
-        """
+        # Run NatureLM analysis
+        results = pipeline([temp_path], queries, window_length_seconds=10.0, hop_length_seconds=10.0)
         
-        # Get Llama model prediction
-        with torch.no_grad():
-            inputs = tokenizer(enhanced_prompt, return_tensors="pt", max_length=512, truncation=True)
-            
-            if torch.cuda.is_available():
-                inputs = {k: v.cuda() for k, v in inputs.items()}
-            
-            outputs = model.generate(
-                **inputs,
-                max_length=1024,
-                temperature=0.7,
-                do_sample=True,
-                pad_token_id=tokenizer.eos_token_id,
-                eos_token_id=tokenizer.eos_token_id
-            )
-            
-            response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-            # Remove the input prompt from the response
-            response_text = response_text.replace(enhanced_prompt, "").strip()
+        # Combine results
+        combined_response = " ".join(results)
         
         # Extract information from response
-        confidence_scores = extract_confidence_from_response(response_text)
-        species_info = extract_species_info(response_text)
+        confidence_scores = extract_confidence_from_response(combined_response)
+        species_info = extract_species_info(combined_response)
         
-        # Calculate overall confidence
+        # Calculate overall confidence based on response quality
         overall_confidence = max(
             confidence_scores["model_confidence"],
             confidence_scores["llama_confidence"],
-            signal_classification.get("confidence", 0.0),
-            50.0  # Default fallback
+            75.0 if species_info["common_name"] else 50.0  # Higher confidence if species identified
         )
         
+        # Clean up temp file
+        os.remove(temp_path)
+        
         return AnalysisResponse(
             species=species_info["common_name"] or "Unknown species",
-            interpretation=response_text,
+            interpretation=combined_response,
             confidence=overall_confidence,
-            signal_type=species_info["signal_type"] or signal_classification.get("signal_type", "Vocalization"),
+            signal_type=species_info["signal_type"] or "Vocalization",
             common_name=species_info["common_name"] or "Unknown",
             scientific_name=species_info["scientific_name"] or "Unknown",
             habitat=species_info["habitat"] or "Unknown habitat",
@@ -328,11 +275,15 @@ async def analyze_audio(file: UploadFile = File(...)):
             audio_characteristics=audio_chars,
             model_confidence=confidence_scores["model_confidence"],
             llama_confidence=confidence_scores["llama_confidence"],
-            additional_insights=response_text,
+            additional_insights=combined_response,
             cluster_group="NatureLM Analysis"
         )
         
     except Exception as e:
+        # Clean up temp file if it exists
+        if os.path.exists(temp_path):
+            os.remove(temp_path)
+        
         raise HTTPException(status_code=500, detail=f"Analysis failed: {str(e)}")
 
 @app.get("/health")