Add model.py

model.py

@@ -0,0 +1,177 @@
+"""
+Autism Detector Model
+
+A feedforward neural network for ASD risk classification
+from structured clinical data.
+"""
+
+import torch
+import torch.nn as nn
+
+
+class AutismDetector(nn.Module):
+    """
+    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=8, hidden_sizes=None, num_classes=2, dropout=0.3):
+        super().__init__()
+
+        if hidden_sizes is None:
+            hidden_sizes = [64, 32]
+
+        layers = []
+        prev_size = input_size
+
+        for hidden_size in hidden_sizes:
+            layers.extend([
+                nn.Linear(prev_size, hidden_size),
+                nn.ReLU(),
+                nn.Dropout(dropout),
+            ])
+            prev_size = hidden_size
+
+        layers.append(nn.Linear(prev_size, num_classes))
+        self.classifier = nn.Sequential(*layers)
+
+        # Store config
+        self.input_size = input_size
+        self.hidden_sizes = hidden_sizes
+        self.num_classes = num_classes
+        self.dropout = dropout
+
+    def forward(self, x):
+        """
+        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():
+            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 TorchScript model."""
+    model = torch.jit.load(model_path, map_location=device)
+    model.eval()
+    return model
+
+
+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 = []
+
+    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)
+
+        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__':
+    # 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}")