Convert to safetensors format and add model architecture

.gitignore

@@ -0,0 +1 @@
+__pycache__

README.md

@@ -18,10 +18,10 @@ A ResNet-34 fine-tuned on a personally-curated dataset to classify images into o
 - Incoherent
 - Semi-Incoherent
 
-**Key Feature**: This model includes a production-ready loading wrapper that handles FastAI's `AdaptiveConcatPool2d` and safe/unsafe weight loading.
+**Key Feature**: This model is provided in safetensors format with a production-ready loading wrapper (`model_architecture.py`) that handles FastAI's `AdaptiveConcatPool2d` layer automatically.
 
 ## Installation and Usage
-**Important**: This model requires custom loading code due to FastAI architecture.
+**Important**: This model is provided in safetensors format and requires the `model_architecture.py` module for proper loading.
 
 1. **Install**
 ```bash
@@ -38,7 +38,7 @@ pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
 pip install -r requirements.txt
 ```
 
-2. **Download and run the example**:
+2. **Download entire repo and run the example**:
 ```bash
 python example_usage.py
 ```
@@ -47,16 +47,13 @@ python example_usage.py
 ```python
 from example_usage import CoherenceClassifier
 # Initialize with your model
-classifier = CoherenceClassifier("coherence_model.pth")
+classifier = CoherenceClassifier("coherence_model.safetensors")
 # Predict on an image
 result = classifier.predict("your_image.jpg", return_probs=True)
 print(result)  # {'coherent': 0.85, 'incoherent': 0.05, 'semi-incoherent': 0.10}
 ```
 
-Note on `weights_only=False`: This model contains FastAI's AdaptiveConcatPool2d layer. You must either:
-
-- Use `weights_only=False` (only if you trust the source), OR
-- Register `AdaptiveConcatPool2d` in `__main__` before loading (see example_usage.py).
+**Note**: The model uses FastAI's `AdaptiveConcatPool2d` layer. Import and use `model_architecture.py` which handles this automatically. The `example_usage.py` script demonstrates the proper import pattern.
 
 ## Model Architecture
 Backbone: ResNet-34 (via FastAI's default)
@@ -68,7 +65,7 @@ For version 1.0, a small dataset was used of ~20k images in coherent category, a
 
 ## Limitations
 - PyTorch environment
-- `weights_only=False` loading has security implications; use only with trusted models.
+- Requires `model_architecture.py` module for proper loading (handles FastAI-specific layers automatically)
 
 Note especially that while an attempt was made at classifying the less obvious but definitely noticeable generation failures like the following into the "semi-incoherent" category, these are much harder to detect and in the current iteration of this model there is no expectation that they will be observed in any particular coherence category, especially when these particular errors occur in a small part of the image.
 - Extra or missing limbs, fingers, or facial features
@@ -83,7 +80,7 @@ Note especially that while an attempt was made at classifying the less obvious b
 ## Comment
 Given the low coherence rate of results produced by early image generation models, it was very surprising that a model was not found for this purpose, necessitating the creation of this one for high-volume review scenarios.
 
-Perhaps models such as this one are avoided or seen as improper due to the perceived danger they pose in introduction of bias to image analysis, however it is highly likely that image generators would rather have at least some bias towards coherence and a somewhat clear mind when reviewing their image output than no bias and a mind littered with the psychologically damaging results of obviously-failed generations which have little to do with the prompter's intent.
+Perhaps models such as this one are avoided or seen as improper due to the perceived danger they pose in introduction of bias to image analysis, however it is highly likely that image generators would rather have at least some bias towards coherence and a somewhat clear mind when reviewing their image output than no bias and a mind littered with the psychologically-damaging results of obviously-failed generations which have little to do with the prompter's intent.
 
 ## Model Card Authors
 Tom Hall

coherence_model.pth → coherence_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9086f65a0853e6c75751c1ae552ecf2ad213e7af92d72619aeee17a31e4c9be6
-size 87437751
+oid sha256:b2f937854cc4ae58ebbe2e6b451006a5ee8dc65fabe5c32d33ef937cae24724b
+size 87356700

example_usage.py

@@ -1,326 +1,287 @@
-"""
-Minimal example to load and use the FastAI-based Coherence Detection model.
-"""
-
-import torch
-import torch.nn as nn
-from torchvision import transforms
-from PIL import Image
-import numpy as np
-import sys
-import os
-
-# ============================================================================
-# 1. CRITICAL: Register AdaptiveConcatPool2d for FastAI compatibility
-# ============================================================================
-def _register_adaptive_concat_pool():
-    """Register AdaptiveConcatPool2d in __main__ for unpickling compatibility"""
-    # Define the class
-    class AdaptiveConcatPool2d(nn.Module):
-        def __init__(self, sz=None):
-            super().__init__()
-            self.ap = nn.AdaptiveAvgPool2d(sz or 1)
-            self.mp = nn.AdaptiveMaxPool2d(sz or 1)
-            
-        def forward(self, x):
-            return torch.cat([self.mp(x), self.ap(x)], 1)
-    
-    # Register in the current module
-    current_module = sys.modules[__name__]
-    if not hasattr(current_module, 'AdaptiveConcatPool2d'):
-        setattr(current_module, 'AdaptiveConcatPool2d', AdaptiveConcatPool2d)
-    
-    # Also register in __main__ if it exists (required for unpickling)
-    if '__main__' in sys.modules:
-        main_module = sys.modules['__main__']
-        if not hasattr(main_module, 'AdaptiveConcatPool2d'):
-            setattr(main_module, 'AdaptiveConcatPool2d', AdaptiveConcatPool2d)
-    
-    return AdaptiveConcatPool2d
-
-# Execute registration immediately
-AdaptiveConcatPool2d = _register_adaptive_concat_pool()
-
-
-# ============================================================================
-# 2. Simplified Model Loader
-# ============================================================================
-class CoherenceClassifier:
-    """Simple loader for the FastAI-based coherence detection model."""
-    
-    # Model categories (update if your model uses different names)
-    CATEGORIES = ["coherent", "incoherent", "semi-incoherent"]
-    
-    def __init__(self, model_path, device="auto"):
-        """
-        Args:
-            model_path: Path to the .pth model file
-            device: "auto", "cuda", or "cpu"
-        """
-        self.model_path = model_path
-        self.device = self._get_device(device)
-        self.model = None
-        self.transform = None
-        
-        # ImageNet normalization (standard for ResNet)
-        self.normalize_mean = [0.485, 0.456, 0.406]
-        self.normalize_std = [0.229, 0.224, 0.225]
-        
-        self.load_model()
-    
-    def _get_device(self, device_str):
-        """Determine the best available device."""
-        if device_str == "auto":
-            return torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        return torch.device(device_str)
-    
-    def load_model(self):
-        """Load the FastAI model with proper handling of AdaptiveConcatPool2d."""
-        print(f"Loading model from {self.model_path}...")
-        print(f"Using device: {self.device}")
-        
-        # IMPORTANT: weights_only=False is REQUIRED for FastAI models
-        # because they contain custom layers like AdaptiveConcatPool2d
-        print("Note: Using weights_only=False for FastAI compatibility")
-        
-        try:
-            # Load the entire model (not just state_dict)
-            self.model = torch.load(
-                self.model_path, 
-                map_location=self.device, 
-                weights_only=False  # Must be False for FastAI models
-            )
-            
-            # Ensure model is in evaluation mode
-            self.model.eval()
-            print("Model loaded successfully!")
-            
-        except Exception as e:
-            print(f"Error loading model: {e}")
-            print("\nTroubleshooting tips:")
-            print("1. Ensure AdaptiveConcatPool2d is registered before loading")
-            print("2. Make sure the model file is not corrupted")
-            print("3. Verify PyTorch and torchvision are installed")
-            raise
-        
-        # Setup image transformations
-        self._setup_transforms()
-    
-    def _setup_transforms(self):
-        """Create image preprocessing pipeline."""
-        self.transform = transforms.Compose([
-            transforms.Resize((224, 224)),  # ResNet standard size
-            transforms.ToTensor(),
-            transforms.Normalize(mean=self.normalize_mean, std=self.normalize_std)
-        ])
-    
-    def preprocess_image(self, image_path):
-        """Load and preprocess an image for the model."""
-        try:
-            # Open and convert to RGB
-            image = Image.open(image_path).convert('RGB')
-            
-            # Apply transformations
-            tensor = self.transform(image)
-            
-            # Add batch dimension [1, 3, 224, 224]
-            tensor = tensor.unsqueeze(0).to(self.device)
-            
-            return tensor
-            
-        except Exception as e:
-            print(f"Error processing image {image_path}: {e}")
-            raise
-    
-    def predict(self, image_path, return_probs=False):
-        """
-        Make a prediction on an image.
-        
-        Args:
-            image_path: Path to the image file
-            return_probs: If True, return all probabilities; if False, return only the top category
-        
-        Returns:
-            Dictionary with predictions or string with top category
-        """
-        if self.model is None:
-            raise ValueError("Model not loaded. Call load_model() first.")
-        
-        # Preprocess the image
-        input_tensor = self.preprocess_image(image_path)
-        
-        # Run inference
-        with torch.no_grad():
-            output = self.model(input_tensor)
-            
-            # Convert to probabilities using softmax
-            probabilities = torch.nn.functional.softmax(output, dim=1)
-            probs = probabilities[0].cpu().numpy()
-        
-        # Create results dictionary
-        results = {self.CATEGORIES[i]: float(probs[i]) for i in range(len(self.CATEGORIES))}
-        
-        if return_probs:
-            return results
-        else:
-            # Return the category with highest probability
-            top_idx = np.argmax(probs)
-            return self.CATEGORIES[top_idx]
-    
-    def predict_batch(self, image_paths):
-        """Make predictions for multiple images."""
-        return [self.predict(img_path, return_probs=True) for img_path in image_paths]
-
-
-# ============================================================================
-# 3. Example Usage
-# ============================================================================
-def main():
-    """Example demonstrating how to use the classifier with all three test images."""
-    
-    # Update this path to your actual model file
-    MODEL_PATH = "coherence_model.pth"  # Change to your model filename
-    
-    # All three example images you've provided
-    DEMO_IMAGES = [
-        "example_coherent.jpg",
-        "example_semi_incoherent.jpg", 
-        "example_incoherent.jpg"
-    ]
-    
-    # Check if model file exists
-    if not os.path.exists(MODEL_PATH):
-        print(f"Error: Model file not found at {MODEL_PATH}")
-        print("Please download the model from the Hugging Face repository:")
-        print("https://huggingface.co/your-username/your-model-name")
-        return
-    
-    # Check which demo images exist
-    available_images = [img for img in DEMO_IMAGES if os.path.exists(img)]
-    
-    if not available_images:
-        print("Error: No example images found.")
-        print(f"Please add one or more of these images to the directory: {DEMO_IMAGES}")
-        return
-    
-    print("=" * 60)
-    print("Coherence Detection Model - Complete Demo")
-    print("=" * 60)
-    
-    # Initialize classifier
-    print(f"Loading model from: {MODEL_PATH}")
-    classifier = CoherenceClassifier(MODEL_PATH, device="auto")
-    
-    print(f"\nFound {len(available_images)} example image(s) for demonstration.")
-    print("-" * 60)
-    
-    # ========================================================================
-    # PART 1: Detailed single image analysis for each available example
-    # ========================================================================
-    print("\n1. SINGLE IMAGE ANALYSIS")
-    print("-" * 40)
-    
-    for img_file in available_images:
-        print(f"\nAnalyzing: {img_file}")
-        print("-" * 30)
-        
-        try:
-            # Get full probability distribution
-            probs = classifier.predict(img_file, return_probs=True)
-            
-            # Display all probabilities
-            print("Prediction probabilities:")
-            for category in classifier.CATEGORIES:
-                prob = probs[category]
-                # Visual indicator for high confidence (>70%)
-                indicator = " ★" if prob > 0.7 else ""
-                print(f"  {category:20} {prob:.4f}{indicator}")
-            
-            # Get and display top category
-            top_category = classifier.predict(img_file, return_probs=False)
-            top_prob = probs[top_category]
-            
-            print(f"\nTop prediction: '{top_category}' ({top_prob:.4f})")
-            
-            # Add interpretation note based on image name
-            if "coherent" in img_file:
-                print("Note: This example should ideally show high 'coherent' probability.")
-            elif "semi_incoherent" in img_file:
-                print("Note: This example demonstrates borderline/partial coherence issues.")
-            elif "incoherent" in img_file:
-                print("Note: This example should show obvious generation failures.")
-            
-        except Exception as e:
-            print(f"Error analyzing {img_file}: {e}")
-            continue
-    
-    # ========================================================================
-    # PART 2: Batch prediction comparison (if multiple images available)
-    # ========================================================================
-    if len(available_images) > 1:
-        print("\n" + "=" * 60)
-        print("2. BATCH PREDICTION COMPARISON")
-        print("-" * 40)
-        
-        try:
-            print(f"Running batch prediction on {len(available_images)} images...")
-            batch_results = classifier.predict_batch(available_images)
-            
-            # Create a comparison table
-            print(f"\n{'Image':30} {'Top Prediction':20} {'Confidence':12}")
-            print("-" * 65)
-            
-            for img_path, result in zip(available_images, batch_results):
-                top_cat = max(result, key=result.get)
-                confidence = result[top_cat]
-                
-                # Shorten filename if too long
-                display_name = os.path.basename(img_path)
-                if len(display_name) > 28:
-                    display_name = display_name[:25] + "..."
-                
-                # Color code high confidence predictions
-                if confidence > 0.8:
-                    confidence_str = f"{confidence:.4f} (HIGH)"
-                elif confidence > 0.6:
-                    confidence_str = f"{confidence:.4f} (MED)"
-                else:
-                    confidence_str = f"{confidence:.4f} (LOW)"
-                
-                print(f"{display_name:30} {top_cat:20} {confidence_str:12}")
-            
-            print("\nBatch processing complete!")
-            
-        except Exception as e:
-            print(f"Error in batch prediction: {e}")
-    
-    # ========================================================================
-    # PART 3: Quick summary
-    # ========================================================================
-    print("\n" + "=" * 60)
-    print("DEMO SUMMARY")
-    print("-" * 40)
-    
-    print(f"✓ Model loaded successfully on: {classifier.device}")
-    print(f"✓ Analyzed {len(available_images)} example image(s)")
-    print(f"✓ Example categories: {classifier.CATEGORIES}")
-    
-    missing_images = [img for img in DEMO_IMAGES if img not in available_images]
-    if missing_images:
-        print(f"\nNote: Missing example images: {missing_images}")
-        print("To complete the demo, add these images to the directory.")
-    
-    print("\n" + "=" * 60)
-    print("Demo completed successfully!")
-    print("\nNext steps:")
-    print("1. Try your own images by modifying the DEMO_IMAGES list")
-    print("2. Use the classifier in your own code:")
-    print("   ```python")
-    print("   from example_usage import CoherenceClassifier")
-    print("   classifier = CoherenceClassifier('your_model.pth')")
-    print("   result = classifier.predict('your_image.jpg')")
-    print("   ```")
-
-
-if __name__ == "__main__":
-    main()
+"""
+Minimal example to load and use the Coherence Detection model.
+Requires safetensors format with exact architecture.
+"""
+
+import torch
+from torchvision import transforms
+from PIL import Image
+import numpy as np
+import sys
+import os
+
+# ============================================================================
+# Import the exact architecture
+# ============================================================================
+try:
+    from model_architecture import load_coherence_model
+    print("✓ Imported exact model architecture")
+except ImportError as e:
+    print(f"Error: model_architecture.py not found or has issues: {e}")
+    print("Please download it from the repository.")
+    sys.exit(1)
+
+# ============================================================================
+# CoherenceClassifier for safetensors
+# ============================================================================
+class CoherenceClassifier:
+    """Loader for coherence detection model (exact architecture)."""
+    
+    # Categories in alphabetical order (as per training)
+    CATEGORIES = ["coherent", "incoherent", "semi-incoherent"]
+    
+    def __init__(self, model_path, device="auto"):
+        """
+        Args:
+            model_path: Path to .safetensors file
+            device: "auto", "cuda", or "cpu"
+        """
+        self.model_path = model_path
+        self.device = self._get_device(device)
+        self.model = None
+        self.transform = None
+        
+        # ImageNet normalization (standard for ResNet)
+        self.normalize_mean = [0.485, 0.456, 0.406]
+        self.normalize_std = [0.229, 0.224, 0.225]
+        
+        self._setup_transforms()
+        self.load_model()
+    
+    def _get_device(self, device_str):
+        """Determine the best available device."""
+        if device_str == "auto":
+            return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        return torch.device(device_str)
+    
+    def load_model(self):
+        """Load model using safetensors and exact architecture."""
+        print(f"Loading coherence model...")
+        print(f"  File: {os.path.basename(self.model_path)}")
+        print(f"  Device: {self.device}")
+        
+        # Verify file type
+        if not self.model_path.lower().endswith('.safetensors'):
+            print("⚠️  Warning: Expected .safetensors file for secure loading.")
+        
+        # Load using our helper function
+        self.model = load_coherence_model(self.model_path, str(self.device))
+    
+    def _setup_transforms(self):
+        """Create image preprocessing pipeline."""
+        self.transform = transforms.Compose([
+            transforms.Resize((224, 224)),  # ResNet standard size
+            transforms.ToTensor(),
+            transforms.Normalize(mean=self.normalize_mean, std=self.normalize_std)
+        ])
+    
+    def preprocess_image(self, image_path):
+        """Load and preprocess an image for the model."""
+        try:
+            # Open and convert to RGB
+            image = Image.open(image_path).convert('RGB')
+            
+            # Apply transformations
+            tensor = self.transform(image)
+            
+            # Add batch dimension [1, 3, 224, 224]
+            tensor = tensor.unsqueeze(0).to(self.device)
+            
+            return tensor
+            
+        except Exception as e:
+            print(f"Error processing image {image_path}: {e}")
+            raise
+    
+    def predict(self, image_path, return_probs=False):
+        """
+        Make a prediction on an image.
+        
+        Args:
+            image_path: Path to the image file
+            return_probs: If True, return all probabilities; if False, return only the top category
+        
+        Returns:
+            Dictionary with predictions or string with top category
+        """
+        if self.model is None:
+            raise ValueError("Model not loaded. Call load_model() first.")
+        
+        # Preprocess the image
+        input_tensor = self.preprocess_image(image_path)
+        
+        # Run inference
+        with torch.no_grad():
+            output = self.model(input_tensor)
+            
+            # Convert to probabilities using softmax
+            probabilities = torch.nn.functional.softmax(output, dim=1)
+            probs = probabilities[0].cpu().numpy()
+        
+        # Create results dictionary
+        results = {self.CATEGORIES[i]: float(probs[i]) for i in range(len(self.CATEGORIES))}
+        
+        if return_probs:
+            return results
+        else:
+            # Return the category with highest probability
+            top_idx = np.argmax(probs)
+            return self.CATEGORIES[top_idx]
+    
+    def predict_batch(self, image_paths):
+        """Make predictions for multiple images."""
+        return [self.predict(img_path, return_probs=True) for img_path in image_paths]
+
+
+# ============================================================================
+# 3. Example Usage
+# ============================================================================
+def main():
+    """Example demonstrating how to use the classifier with all three test images."""
+    
+    # Update this path to your actual model file
+    MODEL_PATH = "coherence_model.safetensors"  # Change to your model filename
+    
+    # All three example images you've provided
+    DEMO_IMAGES = [
+        "example_coherent.jpg",
+        "example_semi_incoherent.jpg", 
+        "example_incoherent.jpg"
+    ]
+    
+    # Check if model file exists
+    if not os.path.exists(MODEL_PATH):
+        print(f"Error: Model file not found at {MODEL_PATH}")
+        print("Please download the model from the Hugging Face repository:")
+        print("https://huggingface.co/your-username/your-model-name")
+        return
+    
+    # Check which demo images exist
+    available_images = [img for img in DEMO_IMAGES if os.path.exists(img)]
+    
+    if not available_images:
+        print("Error: No example images found.")
+        print(f"Please add one or more of these images to the directory: {DEMO_IMAGES}")
+        return
+    
+    print("=" * 60)
+    print("Coherence Detection Model - Complete Demo")
+    print("=" * 60)
+    
+    # Initialize classifier
+    print(f"Loading model from: {MODEL_PATH}")
+    classifier = CoherenceClassifier(MODEL_PATH, device="auto")
+    
+    print(f"\nFound {len(available_images)} example image(s) for demonstration.")
+    print("-" * 60)
+    
+    # ========================================================================
+    # PART 1: Detailed single image analysis for each available example
+    # ========================================================================
+    print("\n1. SINGLE IMAGE ANALYSIS")
+    print("-" * 40)
+    
+    for img_file in available_images:
+        print(f"\nAnalyzing: {img_file}")
+        print("-" * 30)
+        
+        try:
+            # Get full probability distribution
+            probs = classifier.predict(img_file, return_probs=True)
+            
+            # Display all probabilities
+            print("Prediction probabilities:")
+            for category in classifier.CATEGORIES:
+                prob = probs[category]
+                # Visual indicator for high confidence (>70%)
+                indicator = " ★" if prob > 0.7 else ""
+                print(f"  {category:20} {prob:.4f}{indicator}")
+            
+            # Get and display top category
+            top_category = classifier.predict(img_file, return_probs=False)
+            top_prob = probs[top_category]
+            
+            print(f"\nTop prediction: '{top_category}' ({top_prob:.4f})")
+            
+            # Add interpretation note based on image name
+            if img_file == "example_coherent.jpg":
+                print("Note: This example should ideally show high 'coherent' probability.")
+            elif img_file == "example_semi_incoherent.jpg":
+                print("Note: This example demonstrates borderline/partial coherence issues.")
+            elif img_file == "example_incoherent.jpg":
+                print("Note: This example should show obvious generation failures.")
+            
+        except Exception as e:
+            print(f"Error analyzing {img_file}: {e}")
+            continue
+    
+    # ========================================================================
+    # PART 2: Batch prediction comparison (if multiple images available)
+    # ========================================================================
+    if len(available_images) > 1:
+        print("\n" + "=" * 60)
+        print("2. BATCH PREDICTION COMPARISON")
+        print("-" * 40)
+        
+        try:
+            print(f"Running batch prediction on {len(available_images)} images...")
+            batch_results = classifier.predict_batch(available_images)
+            
+            # Create a comparison table
+            print(f"\n{'Image':30} {'Top Prediction':20} {'Confidence':12}")
+            print("-" * 65)
+            
+            for img_path, result in zip(available_images, batch_results):
+                top_cat = max(result, key=result.get)
+                confidence = result[top_cat]
+                
+                # Shorten filename if too long
+                display_name = os.path.basename(img_path)
+                if len(display_name) > 28:
+                    display_name = display_name[:25] + "..."
+                
+                # Color code high confidence predictions
+                if confidence > 0.8:
+                    confidence_str = f"{confidence:.4f} (HIGH)"
+                elif confidence > 0.6:
+                    confidence_str = f"{confidence:.4f} (MED)"
+                else:
+                    confidence_str = f"{confidence:.4f} (LOW)"
+                
+                print(f"{display_name:30} {top_cat:20} {confidence_str:12}")
+            
+            print("\nBatch processing complete!")
+            
+        except Exception as e:
+            print(f"Error in batch prediction: {e}")
+    
+    # ========================================================================
+    # PART 3: Quick summary
+    # ========================================================================
+    print("\n" + "=" * 60)
+    print("DEMO SUMMARY")
+    print("-" * 40)
+    
+    print(f"✓ Model loaded successfully on: {classifier.device}")
+    print(f"✓ Analyzed {len(available_images)} example image(s)")
+    print(f"✓ Example categories: {classifier.CATEGORIES}")
+    
+    missing_images = [img for img in DEMO_IMAGES if img not in available_images]
+    if missing_images:
+        print(f"\nNote: Missing example images: {missing_images}")
+        print("To complete the demo, add these images to the directory.")
+    
+    print("\n" + "=" * 60)
+    print("Demo completed successfully!")
+    print("\nNext steps:")
+    print("1. Try your own images by modifying the DEMO_IMAGES list")
+    print("2. Use the classifier in your own code:")
+    print("   ```python")
+    print("   from example_usage import CoherenceClassifier")
+    print("   classifier = CoherenceClassifier('your_model.pth')")
+    print("   result = classifier.predict('your_image.jpg')")
+    print("   ```")
+
+
+if __name__ == "__main__":
+    main()

model_architecture.py

@@ -0,0 +1,235 @@
+"""
+Exact architecture for Coherence Detection Model.
+Uses key matching to validate the safetensors file.
+"""
+
+import sys
+import torch
+import torch.nn as nn
+
+# ============================================================================
+# AdaptiveConcatPool2d for FastAI model compatitbility
+# ============================================================================
+class AdaptiveConcatPool2d(nn.Module):
+    """FastAI-style adaptive concatenation pooling."""
+    def __init__(self, sz=None):
+        super().__init__()
+        self.ap = nn.AdaptiveAvgPool2d(sz or 1)
+        self.mp = nn.AdaptiveMaxPool2d(sz or 1)
+        
+    def forward(self, x):
+        return torch.cat([self.mp(x), self.ap(x)], 1)
+
+# Also register in __main__ if it exists
+if '__main__' in sys.modules:
+    main_module = sys.modules['__main__']
+    if not hasattr(main_module, 'AdaptiveConcatPool2d'):
+        setattr(main_module, 'AdaptiveConcatPool2d', AdaptiveConcatPool2d)
+        
+# Also register under common FastAI names
+if '__main__' in sys.modules:
+    main_module = sys.modules['__main__']
+    # Some FastAI models might expect this
+    if not hasattr(main_module, 'AdaptiveConcatPool'):
+        setattr(main_module, 'AdaptiveConcatPool', AdaptiveConcatPool2d)
+
+# ============================================================================
+# Utility function to check torchvision version
+# ============================================================================
+def _get_resnet_backbone():
+    """Helper to get ResNet backbone with version-appropriate API."""
+    from torchvision.models import resnet34
+    import torchvision
+    
+    # Parse version to determine API
+    version = torchvision.__version__.split('.')
+    major = int(version[0]) if version[0].isdigit() else 0
+    minor = int(version[1]) if len(version) > 1 and version[1].isdigit() else 0
+    
+    if major >= 0 and minor >= 13:
+        # Use new weights API
+        return resnet34(weights=None)
+    else:
+        # Use old pretrained API
+        return resnet34(pretrained=False)
+
+
+# ============================================================================
+# Clean model with version detection
+# ============================================================================
+class CoherenceDetectionModel(nn.Sequential):
+    """
+    Clean version that handles torchvision API changes properly.
+    """
+    def __init__(self, num_classes=3):
+        # Get backbone using version-appropriate API
+        backbone = _get_resnet_backbone()
+        backbone = nn.Sequential(*list(backbone.children())[:-2])
+        
+        # Classifier head
+        classifier = nn.Sequential(
+            AdaptiveConcatPool2d(),
+            nn.Flatten(start_dim=1, end_dim=-1),
+            nn.BatchNorm1d(1024),
+            nn.Dropout(p=0.25, inplace=False),
+            nn.Linear(1024, 512, bias=True),
+            nn.ReLU(inplace=True),
+            nn.BatchNorm1d(512),
+            nn.Dropout(p=0.5, inplace=False),
+            nn.Linear(512, num_classes, bias=True)
+        )
+        
+        super().__init__(backbone, classifier)
+
+
+# ============================================================================
+# Loading function
+# ============================================================================
+def load_coherence_model(safetensors_path, device='auto'):
+    """
+    Load safetensors weights with automatic key remapping.
+    
+    Args:
+        safetensors_path: Path to .safetensors file
+        device: 'auto', 'cuda', or 'cpu'
+    """
+    import safetensors.torch
+    
+    # Determine device
+    if device == 'auto':
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    else:
+        device = torch.device(device)
+    
+    # Create model instance
+    model = CoherenceDetectionModel(num_classes=3)
+    
+    # Load safetensors
+    try:
+        state_dict = safetensors.torch.load_file(safetensors_path, device='cpu')
+    except FileNotFoundError:
+        print(f"Error: File '{safetensors_path}' not found.")
+        print("Testing with sample model structure...")
+        state_dict = model.state_dict()
+    
+    # Load directly (keys should match exactly)
+    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+    
+    if missing_keys:
+        print(f"Warning: Missing keys: {missing_keys}")
+    
+    if unexpected_keys:
+        print(f"Warning: Unexpected keys: {unexpected_keys}")
+    
+    if not missing_keys and not unexpected_keys:
+        print(f"✓ CoherenceDetectionModel loaded successfully (exact match)")
+    else:
+        print(f"⚠ CoherenceDetectionModel loaded with key mismatches")
+    
+    model = model.to(device)
+    model.eval()
+    
+    print(f"  Device: {device}")
+    print(f"  Parameters: {sum(p.numel() for p in model.parameters()):,}")
+    
+    return model
+
+
+# ============================================================================
+# Test functions
+# ============================================================================
+def test_key_matching(safetensors_path="coherence_model.safetensors"):
+    """Test that keys match between model and safetensors."""
+    import safetensors.torch
+    
+    print("\nTesting key matching...")
+    
+    try:
+        state_dict = safetensors.torch.load_file(safetensors_path, device='cpu')
+    except FileNotFoundError:
+        print(f"  ⚠ File '{safetensors_path}' not found, using model weights")
+        model = CoherenceDetectionModel(num_classes=3)
+        state_dict = model.state_dict()
+    
+    print("\nTesting CoherenceDetectionModel:")
+    model_clean = CoherenceDetectionModel(num_classes=3)
+    missing, unexpected = model_clean.load_state_dict(state_dict, strict=False)
+    
+    if not missing and not unexpected:
+        print("  ✅ Load successful (exact key match)")
+    else:
+        print(f"  ⚠ Load completed with issues")
+        if missing:
+            print(f"    Missing keys: {len(missing)}")
+        if unexpected:
+            print(f"    Unexpected keys: {len(unexpected)}")
+    
+    return model_clean if not missing and not unexpected else None
+
+
+def print_key_samples(safetensors_path="coherence_model.safetensors"):
+    """Print sample keys for debugging."""
+    import safetensors.torch
+    
+    print("\nKey samples:")
+    
+    try:
+        state_dict = safetensors.torch.load_file(safetensors_path, device='cpu')
+        print("From safetensors file (first 5 keys):")
+        for i, key in enumerate(sorted(state_dict.keys())[:5]):
+            print(f"  {i}: {key}")
+    except FileNotFoundError:
+        print(f"Safetensors file '{safetensors_path}' not found")
+        print("Showing model structure keys instead:")
+        state_dict = None
+    
+    print("\nFrom CoherenceDetectionModel:")
+    model = CoherenceDetectionModel(num_classes=3)
+    for i, key in enumerate(sorted(model.state_dict().keys())[:5]):
+        print(f"  {i}: {key}")
+    
+    return state_dict
+
+
+# ============================================================================
+# Version compatibility info
+# ============================================================================
+def print_version_info():
+    """Print version information for debugging."""
+    import torch
+    import torchvision
+    
+    print("\n" + "=" * 60)
+    print("Version Information")
+    print("=" * 60)
+    print(f"Torch: {torch.__version__}")
+    print(f"Torchvision: {torchvision.__version__}")
+    print(f"CUDA Available: {torch.cuda.is_available()}")
+    
+    # Check API compatibility
+    version = torchvision.__version__.split('.')
+    major = int(version[0]) if version[0].isdigit() else 0
+    minor = int(version[1]) if len(version) > 1 and version[1].isdigit() else 0
+    
+    if major >= 0 and minor >= 13:
+        print("✓ Using modern torchvision API (weights parameter)")
+    else:
+        print("⚠ Using legacy torchvision API (pretrained parameter)")
+    print("=" * 60)
+
+
+if __name__ == "__main__":
+    print("=" * 60)
+    print("Coherence Detection Model Architecture")
+    print("=" * 60)
+    
+    print_version_info()
+    state_dict = print_key_samples()
+    print("\n" + "=" * 60)
+    model = test_key_matching()
+    
+    if model:
+        print("\nModel summary:")
+        print(f"  Backbone layers: {len(model[0])}")
+        print(f"  Classifier layers: {len(model[1])}")
+        print(f"  Total sequential blocks: {len(model)}")