export_torchscript.py

#!/usr/bin/env python3
"""
Export trained models to TorchScript for fast inference

TorchScript models can be loaded and run without Python dependencies,
providing significant speedup for inference.
"""

import torch
import argparse
from pathlib import Path
from rich.console import Console
from rich.progress import Progress
import json

from src.ml.advanced_classifier import AdvancedMetalClassifier, MetalDetectorCNN, MetalDetectorTransformer
from src.ml.deep_model import Wav2VecClassifier

console = Console()

def export_advanced_models(model_path: Path, output_dir: Path):
    """Export the advanced ensemble models to TorchScript."""
    console.print("🔧 Exporting Advanced Ensemble Models to TorchScript...")
    
    # Load the classifier
    classifier = AdvancedMetalClassifier(model_path=model_path)
    
    if not classifier._model_exists():
        console.print("❌ No trained models found! Train models first.")
        return False
    
    classifier.load_model()
    n_classes = len(classifier.label_encoder.classes_)
    sample_rate = classifier.sample_rate
    max_length = int(classifier.max_audio_length * sample_rate)
    
    console.print(f"✅ Loaded models with {n_classes} classes")
    
    # Create dummy input for tracing
    dummy_input = torch.randn(1, max_length)  # Batch size 1, audio length
    
    with Progress() as progress:
        task = progress.add_task("Exporting models...", total=3)
        
        # 1. Export CNN model
        try:
            classifier.cnn_model.eval()
            traced_cnn = torch.jit.trace(classifier.cnn_model, dummy_input)
            cnn_path = output_dir / "cnn_model_traced.pt"
            torch.jit.save(traced_cnn, cnn_path)
            console.print(f"✅ CNN model exported to: {cnn_path}")
            progress.update(task, advance=1)
        except Exception as e:
            console.print(f"❌ Failed to export CNN: {e}")
            return False
        
        # 2. Export Transformer model
        try:
            classifier.transformer_model.eval()
            traced_transformer = torch.jit.trace(classifier.transformer_model, dummy_input)
            transformer_path = output_dir / "transformer_model_traced.pt"
            torch.jit.save(traced_transformer, transformer_path)
            console.print(f"✅ Transformer model exported to: {transformer_path}")
            progress.update(task, advance=1)
        except Exception as e:
            console.print(f"❌ Failed to export Transformer: {e}")
            return False
        
        # 3. Export preprocessing info
        metadata = {
            "sample_rate": sample_rate,
            "max_audio_length": classifier.max_audio_length,
            "max_samples": max_length,
            "n_classes": n_classes,
            "classes": classifier.label_encoder.classes_.tolist(),
            "ensemble_weights": {
                "cnn": 0.4,
                "transformer": 0.4,
                "traditional": 0.2
            }
        }
        
        metadata_path = output_dir / "model_metadata.json"
        with open(metadata_path, 'w') as f:
            json.dump(metadata, f, indent=2)
        console.print(f"✅ Model metadata saved to: {metadata_path}")
        progress.update(task, advance=1)
    
    return True

def export_wav2vec_model(checkpoint_path: Path, output_dir: Path):
    """Export Wav2Vec2 model to TorchScript."""
    console.print("🔧 Exporting Wav2Vec2 Model to TorchScript...")
    
    if not checkpoint_path.exists():
        console.print(f"❌ Checkpoint not found: {checkpoint_path}")
        return False
    
    # Load checkpoint
    checkpoint = torch.load(checkpoint_path, map_location='cpu')
    
    # Extract hyperparameters
    hparams = checkpoint.get('hyper_parameters', {})
    num_classes = hparams.get('num_classes', 3)
    max_length = hparams.get('max_length', 5 * 16000)  # 5 seconds at 16kHz
    
    # Create model
    model = Wav2VecClassifier(num_classes=num_classes, max_length=max_length)
    model.load_state_dict(checkpoint['state_dict'])
    model.eval()
    
    console.print(f"✅ Loaded Wav2Vec2 model with {num_classes} classes")
    
    # Create dummy input
    dummy_input = torch.randn(1, max_length)  # Batch size 1
    
    try:
        # Use scripting instead of tracing for Wav2Vec2 (more complex model)
        scripted_model = torch.jit.script(model)
        output_path = output_dir / "wav2vec_model_scripted.pt"
        torch.jit.save(scripted_model, output_path)
        console.print(f"✅ Wav2Vec2 model exported to: {output_path}")
        
        # Save metadata
        metadata = {
            "model_type": "wav2vec2",
            "num_classes": num_classes,
            "max_length": max_length,
            "sample_rate": 16000,  # Wav2Vec2 expects 16kHz
            "checkpoint_source": str(checkpoint_path)
        }
        
        metadata_path = output_dir / "wav2vec_metadata.json"
        with open(metadata_path, 'w') as f:
            json.dump(metadata, f, indent=2)
        console.print(f"✅ Metadata saved to: {metadata_path}")
        
        return True
        
    except Exception as e:
        console.print(f"❌ Failed to export Wav2Vec2: {e}")
        return False

def create_fast_inference_script(output_dir: Path):
    """Create a standalone inference script that uses TorchScript models."""
    script_content = '''#!/usr/bin/env python3
"""Fast inference using TorchScript models - minimal dependencies"""

import torch
import torch.nn.functional as F
import numpy as np
import json
from pathlib import Path

class FastInference:
    def __init__(self, model_dir: Path):
        self.model_dir = Path(model_dir)
        
        # Load metadata
        with open(self.model_dir / "model_metadata.json", 'r') as f:
            self.metadata = json.load(f)
        
        # Load TorchScript models
        self.cnn_model = torch.jit.load(self.model_dir / "cnn_model_traced.pt")
        self.transformer_model = torch.jit.load(self.model_dir / "transformer_model_traced.pt")
        
        # Set to eval mode
        self.cnn_model.eval()
        self.transformer_model.eval()
        
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.cnn_model = self.cnn_model.to(self.device)
        self.transformer_model = self.transformer_model.to(self.device)
    
    def preprocess_audio(self, audio: np.ndarray) -> torch.Tensor:
        """Preprocess audio to match training format."""
        # Ensure correct length
        target_length = self.metadata['max_samples']
        if len(audio) > target_length:
            audio = audio[:target_length]
        else:
            audio = np.pad(audio, (0, target_length - len(audio)), mode='constant')
        
        # Convert to tensor
        audio_tensor = torch.FloatTensor(audio).unsqueeze(0)  # Add batch dimension
        return audio_tensor.to(self.device)
    
    def classify(self, audio: np.ndarray) -> dict:
        """Classify audio using ensemble of models."""
        # Preprocess
        audio_tensor = self.preprocess_audio(audio)
        
        with torch.no_grad():
            # Get predictions from each model
            cnn_logits = self.cnn_model(audio_tensor)
            trans_logits = self.transformer_model(audio_tensor)
            
            # Convert to probabilities
            cnn_probs = F.softmax(cnn_logits, dim=1)[0].cpu().numpy()
            trans_probs = F.softmax(trans_logits, dim=1)[0].cpu().numpy()
            
            # Ensemble (excluding traditional ML for pure TorchScript)
            weights = self.metadata['ensemble_weights']
            ensemble_probs = (cnn_probs * weights['cnn'] + trans_probs * weights['transformer'])
            ensemble_probs = ensemble_probs / (weights['cnn'] + weights['transformer'])
            
            # Get prediction
            predicted_idx = int(np.argmax(ensemble_probs))
            confidence = float(ensemble_probs[predicted_idx])
            
            return {
                'predicted_label': self.metadata['classes'][predicted_idx],
                'confidence': confidence,
                'all_probabilities': {
                    cls: float(prob) for cls, prob in zip(self.metadata['classes'], ensemble_probs)
                }
            }

# Usage example
if __name__ == "__main__":
    import sys
    
    if len(sys.argv) < 2:
        print("Usage: python fast_inference.py <model_dir>")
        sys.exit(1)
    
    # Initialize
    model_dir = Path(sys.argv[1])
    inferencer = FastInference(model_dir)
    
    # Example: classify random audio
    sample_audio = np.random.randn(inferencer.metadata['max_samples']).astype(np.float32)
    result = inferencer.classify(sample_audio)
    
    print(f"Predicted: {result['predicted_label']} (confidence: {result['confidence']:.3f})")
    print(f"All probabilities: {result['all_probabilities']}")
'''
    
    script_path = output_dir / "fast_inference.py"
    script_path.write_text(script_content)
    script_path.chmod(0o755)  # Make executable
    
    console.print(f"✅ Fast inference script created: {script_path}")

def main():
    parser = argparse.ArgumentParser(description="Export models to TorchScript")
    parser.add_argument("--model-dir", type=str, default="models/advanced",
                       help="Directory containing trained models")
    parser.add_argument("--output-dir", type=str, default="models/torchscript",
                       help="Output directory for TorchScript models")
    parser.add_argument("--wav2vec-checkpoint", type=str, 
                       help="Path to Wav2Vec2 checkpoint file")
    
    args = parser.parse_args()
    
    model_dir = Path(args.model_dir)
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    console.print(f"📦 Exporting models to TorchScript format...")
    console.print(f"📁 Output directory: {output_dir}")
    
    # Export advanced ensemble models
    if model_dir.exists():
        success = export_advanced_models(model_dir, output_dir)
        if success:
            create_fast_inference_script(output_dir)
    else:
        console.print(f"❌ Model directory not found: {model_dir}")
        return
    
    # Export Wav2Vec2 if checkpoint provided
    if args.wav2vec_checkpoint:
        checkpoint_path = Path(args.wav2vec_checkpoint)
        export_wav2vec_model(checkpoint_path, output_dir)
    
    console.print("\n🎉 Export complete!")
    console.print(f"📍 TorchScript models saved to: {output_dir}")
    console.print("\n💡 Usage:")
    console.print(f"   python {output_dir}/fast_inference.py {output_dir}")

if __name__ == "__main__":
    main()