Add model.py

model.py

@@ -1,98 +1,37 @@
 """
-Simplified Autism Spectrum Disorder (ASD) Detector Model
+Autism Detector Model
 
-8-feature model capturing 84% of predictive power.
+A feedforward neural network for ASD risk classification
+from structured clinical data.
 """
 
 import torch
 import torch.nn as nn
-import pandas as pd
-import numpy as np
-from sklearn.preprocessing import StandardScaler, LabelEncoder
 
 
-# Original column names used in training
-SELECTED_FEATURES = [
-    'Developmental milestones- global delay (G), motor delay (M), cognitive delay (C)',
-    'IQ/DQ',
-    'ICD',
-    'Language disorder Y= present, N=absent',
-    'Language development: delay, normal=N, absent=A',
-    'Dysmorphysm y=present, no=absent',
-    'Behaviour disorder- agressivity, agitation, irascibility',
-    'Neurological Examination; N=normal, text = abnormal; free cell = examination not performed ???'
-]
-
-
-class SimplePreprocessor:
-    """Preprocessor for the 8 selected features."""
-
-    def __init__(self):
-        self.label_encoders = {}
-        self.scaler = StandardScaler()
-        self.numeric_cols = ['IQ/DQ']
-        self.categorical_cols = [f for f in SELECTED_FEATURES if f != 'IQ/DQ']
-
-    def fit(self, X):
-        X = X.copy()
-        X['IQ/DQ'] = pd.to_numeric(X['IQ/DQ'], errors='coerce')
-
-        for col in self.categorical_cols:
-            X[col] = X[col].fillna('_missing_').astype(str)
-            all_values = list(X[col].unique()) + ['_missing_', '_unknown_']
-            self.label_encoders[col] = LabelEncoder()
-            self.label_encoders[col].fit(all_values)
-
-        X_encoded = self._encode(X)
-        self.scaler.fit(X_encoded)
-        return self
-
-    def _encode(self, X):
-        X = X.copy()
-        X['IQ/DQ'] = pd.to_numeric(X['IQ/DQ'], errors='coerce').fillna(70)
-
-        for col in self.categorical_cols:
-            X[col] = X[col].fillna('_missing_').astype(str)
-            known_classes = set(self.label_encoders[col].classes_)
-            X[col] = X[col].apply(lambda x: x if x in known_classes else '_unknown_')
-            X[col] = self.label_encoders[col].transform(X[col])
-
-        return X[SELECTED_FEATURES].values
-
-    def transform(self, X):
-        X_encoded = self._encode(X)
-        return self.scaler.transform(X_encoded)
-
-    def fit_transform(self, X):
-        self.fit(X)
-        return self.transform(X)
-
-
-class SimplifiedASDDetector(nn.Module):
+class AutismDetector(nn.Module):
     """
-    Simplified neural network for ASD detection using 8 key features.
-
-    Features:
-        1. developmental_milestones - N/G/M/C
-        2. iq_dq - numeric (0-150)
-        3. intellectual_disability - N/F70.0/F71/F72
-        4. language_disorder - N/Y
-        5. language_development - N/delay/A
-        6. dysmorphism - NO/Y
-        7. behaviour_disorder - N/Y
-        8. neurological_exam - N/abnormal text
-
-    Args:
-        input_size (int): Number of input features (8 after encoding)
-        hidden_sizes (list): Hidden layer sizes. Default: [32, 16]
-        dropout_rate (float): Dropout probability. Default: 0.3
+    Binary classifier for autism spectrum disorder screening.
+
+    Input: 8 preprocessed clinical features
+    Output: 2 logits (Healthy, ASD)
+
+    Features (in order):
+        1. developmental_milestones - N/G/M/C (encoded 0-3)
+        2. iq_dq - numeric, normalized 0-1
+        3. intellectual_disability - N/F70.0/F71/F72 (encoded 0-3)
+        4. language_disorder - N/Y (encoded 0-1)
+        5. language_development - N/delay/A (encoded 0-2)
+        6. dysmorphism - NO/Y (encoded 0-1)
+        7. behaviour_disorder - N/Y (encoded 0-1)
+        8. neurological_exam - N/abnormal (encoded 0-1)
     """
 
-    def __init__(self, input_size, hidden_sizes=None, dropout_rate=0.3):
-        super(SimplifiedASDDetector, self).__init__()
+    def __init__(self, input_size=8, hidden_sizes=None, num_classes=2, dropout=0.3):
+        super().__init__()
 
         if hidden_sizes is None:
-            hidden_sizes = [32, 16]
+            hidden_sizes = [64, 32]
 
         layers = []
         prev_size = input_size
@@ -100,94 +39,139 @@ class SimplifiedASDDetector(nn.Module):
         for hidden_size in hidden_sizes:
             layers.extend([
                 nn.Linear(prev_size, hidden_size),
-                nn.BatchNorm1d(hidden_size),
                 nn.ReLU(),
-                nn.Dropout(dropout_rate)
+                nn.Dropout(dropout),
             ])
             prev_size = hidden_size
 
-        layers.append(nn.Linear(prev_size, 1))
-        layers.append(nn.Sigmoid())
+        layers.append(nn.Linear(prev_size, num_classes))
+        self.classifier = nn.Sequential(*layers)
 
-        self.network = nn.Sequential(*layers)
+        # Store config
         self.input_size = input_size
         self.hidden_sizes = hidden_sizes
-        self.dropout_rate = dropout_rate
+        self.num_classes = num_classes
+        self.dropout = dropout
 
     def forward(self, x):
-        """Forward pass returning probability of ASD."""
-        return self.network(x)
-
-    def predict(self, x, threshold=0.5):
-        """Binary prediction (0=Healthy, 1=ASD)."""
+        """
+        Forward pass.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape (batch_size, 8)
+
+        Returns
+        -------
+        torch.Tensor
+            Output logits of shape (batch_size, num_classes)
+        """
+        return self.classifier(x)
+
+    def predict(self, x):
+        """
+        Make predictions with probabilities.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape (batch_size, 8)
+
+        Returns
+        -------
+        dict with 'prediction', 'probability', 'logits'
+        """
         self.eval()
         with torch.no_grad():
-            probs = self.forward(x)
-            return (probs > threshold).int()
+            logits = self.forward(x)
+            probs = torch.softmax(logits, dim=-1)
+            pred_class = torch.argmax(probs, dim=-1)
+
+        return {
+            'prediction': pred_class,
+            'probabilities': probs,
+            'logits': logits
+        }
 
 
 def load_model(model_path, device='cpu'):
-    """Load trained model from .pth file."""
-    checkpoint = torch.load(model_path, map_location=device, weights_only=False)
+    """Load TorchScript model."""
+    model = torch.jit.load(model_path, map_location=device)
+    model.eval()
+    return model
 
-    model = SimplifiedASDDetector(
-        input_size=checkpoint['input_size'],
-        hidden_sizes=checkpoint['hidden_sizes'],
-        dropout_rate=checkpoint['dropout_rate']
-    )
 
-    model.load_state_dict(checkpoint['model_state_dict'])
-    model.to(device)
-    model.eval()
+def preprocess(data, config):
+    """
+    Preprocess input data using JSON config.
+
+    Parameters
+    ----------
+    data : dict
+        Input features as dictionary
+    config : dict
+        Preprocessor configuration from preprocessor_config.json
+
+    Returns
+    -------
+    torch.Tensor
+        Preprocessed features tensor of shape (1, 8)
+    """
+    features = []
 
-    return model
+    for feature_name in config["feature_order"]:
+        if feature_name in config["categorical_features"]:
+            feat_config = config["categorical_features"][feature_name]
 
+            if feat_config["type"] == "text_binary":
+                # For neurological_exam: N -> 0, anything else -> 1
+                raw_value = str(data[feature_name]).strip().upper()
+                value = 0 if raw_value == feat_config["normal_value"] else 1
+            else:
+                # Standard categorical/binary mapping
+                raw_value = data[feature_name]
+                value = feat_config["mapping"].get(raw_value, 0)
 
-# Feature information for reference
-FEATURES = {
-    'developmental_milestones': {
-        'description': 'Developmental milestones status',
-        'values': {'N': 'Normal', 'G': 'Global delay', 'M': 'Motor delay', 'C': 'Cognitive delay'}
-    },
-    'iq_dq': {
-        'description': 'IQ or Developmental Quotient',
-        'values': 'numeric (typically 20-150, average ~100)'
-    },
-    'intellectual_disability': {
-        'description': 'ICD code for intellectual disability',
-        'values': {'N': 'None', 'F70.0': 'Mild (IQ 50-69)', 'F71': 'Moderate (IQ 35-49)', 'F72': 'Severe (IQ 20-34)'}
-    },
-    'language_disorder': {
-        'description': 'Presence of language disorder',
-        'values': {'N': 'No', 'Y': 'Yes'}
-    },
-    'language_development': {
-        'description': 'Language development status',
-        'values': {'N': 'Normal', 'delay': 'Delayed', 'A': 'Absent'}
-    },
-    'dysmorphism': {
-        'description': 'Physical dysmorphic features',
-        'values': {'NO': 'Absent', 'Y': 'Present'}
-    },
-    'behaviour_disorder': {
-        'description': 'Behavioral issues (aggression, agitation)',
-        'values': {'N': 'No', 'Y': 'Yes'}
-    },
-    'neurological_exam': {
-        'description': 'Neurological examination result',
-        'values': {'N': 'Normal', 'other': 'Abnormal (free text description)'}
-    }
-}
+        elif feature_name in config["numeric_features"]:
+            feat_config = config["numeric_features"][feature_name]
+            raw = float(data[feature_name])
+            # Min-max normalization
+            value = (raw - feat_config["min"]) / (feat_config["max"] - feat_config["min"])
+            value = max(0, min(1, value))  # Clamp to [0, 1]
+
+        features.append(value)
+
+    return torch.tensor([features], dtype=torch.float32)
+
+
+def get_risk_level(probability):
+    """
+    Get risk level from ASD probability.
+
+    Returns
+    -------
+    str: 'low', 'medium', or 'high'
+    """
+    if probability < 0.4:
+        return "low"
+    elif probability < 0.7:
+        return "medium"
+    else:
+        return "high"
 
 
 if __name__ == '__main__':
-    print("Simplified ASD Detector - 8 Features")
-    print("=" * 50)
-    print("\nRequired inputs:")
-    for i, (name, info) in enumerate(FEATURES.items(), 1):
-        print(f"\n{i}. {name}")
-        print(f"   Description: {info['description']}")
-        if isinstance(info['values'], dict):
-            print(f"   Values: {', '.join(f'{k}={v}' for k, v in info['values'].items())}")
-        else:
-            print(f"   Values: {info['values']}")
+    # Test model creation
+    model = AutismDetector()
+    print(f"Model architecture:\n{model}")
+
+    # Test forward pass
+    x = torch.randn(2, 8)
+    output = model(x)
+    print(f"\nInput shape: {x.shape}")
+    print(f"Output shape: {output.shape}")
+    print(f"Output (logits): {output}")
+
+    probs = torch.softmax(output, dim=-1)
+    print(f"Probabilities: {probs}")