phenotype_inference.py

import re
import math
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn


DEFAULT_SCHEMA = [
    "Gram Stain",
    "Shape",
    "Catalase",
    "Oxidase",
    "Colony Morphology",
    "Haemolysis",
    "Haemolysis Type",
    "Indole",
    "Growth Temperature",
    "Media Grown On",
    "Motility",
    "Motility Type",
    "Capsule",
    "Spore Formation",
    "Oxygen Requirement",
    "Methyl Red",
    "VP",
    "Citrate",
    "Urease",
    "H2S",
    "Lactose Fermentation",
    "Glucose Fermentation",
    "Sucrose Fermentation",
    "Nitrate Reduction",
    "Lysine Decarboxylase",
    "Ornithine Decarboxylase",
    "Arginine dihydrolase",
    "Gelatin Hydrolysis",
    "Esculin Hydrolysis",
    "DNase",
    "ONPG",
    "NaCl Tolerant (>=6%)",
    "Lipase Test",
    "Xylose Fermentation",
    "Rhamnose Fermentation",
    "Mannitol Fermentation",
    "Sorbitol Fermentation",
    "Maltose Fermentation",
    "Arabinose Fermentation",
    "Raffinose Fermentation",
    "Inositol Fermentation",
    "Trehalose Fermentation",
    "Coagulase",
    "TSI Pattern",
    "Gas Production",
]


class PhenotypeTinyTransformer(nn.Module):
    def __init__(
        self,
        vocab_size,
        num_classes,
        aux_dim,
        max_tokens,
        d_model=512,
        n_heads=8,
        n_layers=8,
        ff_dim=1536,
        dropout=0.12,
    ):
        super().__init__()

        self.token_embedding = nn.Embedding(vocab_size, d_model, padding_idx=0)
        self.position_embedding = nn.Embedding(max_tokens, d_model)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=n_heads,
            dim_feedforward=ff_dim,
            dropout=dropout,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )

        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)

        self.aux_mlp = nn.Sequential(
            nn.Linear(aux_dim, d_model),
            nn.GELU(),
            nn.LayerNorm(d_model),
            nn.Dropout(dropout),
        )

        self.classifier = nn.Sequential(
            nn.LayerNorm(d_model * 2),
            nn.Dropout(dropout),
            nn.Linear(d_model * 2, d_model),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_model, num_classes),
        )

    def forward(self, input_ids, attention_mask, aux):
        batch_size, seq_len = input_ids.shape

        positions = torch.arange(seq_len, device=input_ids.device)
        positions = positions.unsqueeze(0).expand(batch_size, seq_len)

        x = self.token_embedding(input_ids) + self.position_embedding(positions)

        key_padding_mask = attention_mask == 0

        encoded = self.encoder(x, src_key_padding_mask=key_padding_mask)

        cls_vec = encoded[:, 0, :]
        aux_vec = self.aux_mlp(aux)

        combined = torch.cat([cls_vec, aux_vec], dim=-1)

        return self.classifier(combined)


class PhenotypeClassifier:
    def __init__(self, model_path, device=None):
        self.model_path = Path(model_path)

        if not self.model_path.exists():
            raise FileNotFoundError(f"Model file not found: {self.model_path}")

        self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))

        try:
            checkpoint = torch.load(self.model_path, map_location=self.device, weights_only=False)
        except TypeError:
            checkpoint = torch.load(self.model_path, map_location=self.device)

        self.config = checkpoint["config"]
        self.vocab = checkpoint["vocab"]
        self.classes = checkpoint["classes"]

        temperature_info = checkpoint.get("temperature_scaling", {})
        self.temperature = float(temperature_info.get("temperature", 1.0))

        self.schema = self.config.get("schema", DEFAULT_SCHEMA)

        self.semicolon_fields = set(self.config.get("semicolon_fields", [
            "Shape",
            "Colony Morphology",
            "Media Grown On",
            "Motility Type",
            "Oxygen Requirement",
        ]))

        self.core_fields = self.config.get("core_fields", [
            "Gram Stain",
            "Shape",
            "Catalase",
            "Oxidase",
            "Oxygen Requirement",
            "Spore Formation",
            "Motility",
            "Indole",
            "Urease",
            "Citrate",
            "H2S",
            "Nitrate Reduction",
        ])

        self.growth_temp_thresholds = self.config.get(
            "growth_temperature_thresholds",
            [4, 10, 20, 25, 30, 37, 42, 45, 50, 55],
        )

        self.max_tokens = int(self.config.get("max_tokens", 192))
        self.aux_dim = int(self.config.get("aux_dim", 10))
        self.d_model = int(self.config.get("d_model", 512))
        self.n_heads = int(self.config.get("n_heads", 8))
        self.n_layers = int(self.config.get("n_layers", 8))
        self.ff_dim = int(self.config.get("ff_dim", 1536))
        self.dropout = float(self.config.get("dropout", 0.12))

        self.pad_token = "<PAD>"
        self.unk_token = "<UNK>"
        self.cls_token = "<CLS>"

        self.field_lookup = {field.lower(): field for field in self.schema}

        self.use_amp = self.device.type == "cuda"
        self.use_bf16 = self.device.type == "cuda" and torch.cuda.is_bf16_supported()
        self.amp_dtype = torch.bfloat16 if self.use_bf16 else torch.float16

        self.model = PhenotypeTinyTransformer(
            vocab_size=len(self.vocab),
            num_classes=len(self.classes),
            aux_dim=self.aux_dim,
            max_tokens=self.max_tokens,
            d_model=self.d_model,
            n_heads=self.n_heads,
            n_layers=self.n_layers,
            ff_dim=self.ff_dim,
            dropout=self.dropout,
        ).to(self.device)

        self.model.load_state_dict(checkpoint["model_state_dict"])
        self.model.eval()

    @staticmethod
    def safe_str(x):
        if x is None:
            return None
        if isinstance(x, float) and np.isnan(x):
            return None
        x = str(x).strip()
        return x if x else None

    @classmethod
    def split_value(cls, value):
        value = cls.safe_str(value)

        if value is None:
            return []

        if ";" in value:
            return [p.strip() for p in value.split(";") if p.strip()]

        return [value]

    def canonical_field_name(self, field):
        field = self.safe_str(field)

        if field is None:
            return None

        lowered = field.lower()

        if lowered in self.field_lookup:
            return self.field_lookup[lowered]

        return field

    @staticmethod
    def normalize_basic_value(value):
        value = PhenotypeClassifier.safe_str(value)

        if value is None:
            return None

        lowered = value.lower()

        value_map = {
            "pos": "Positive",
            "positive": "Positive",
            "+": "Positive",
            "neg": "Negative",
            "negative": "Negative",
            "-": "Negative",
            "var": "Variable",
            "variable": "Variable",
            "none": "None",
            "n/a": "None",
            "na": "None",
            "aerobe": "Aerobic",
            "aerobic": "Aerobic",
            "anaerobe": "Anaerobic",
            "anaerobic": "Anaerobic",
            "facultative anaerobe": "Facultative Anaerobic",
            "facultative anaerobic": "Facultative Anaerobic",
            "microaerophile": "Microaerophilic",
            "microaerophilic": "Microaerophilic",
            "rod": "Rods",
            "rods": "Rods",
            "short rods": "Short Rods",
            "short rod": "Short Rods",
            "cocci": "Cocci",
            "coccus": "Cocci",
            "yeast": "Yeast",
            "spiral": "Spiral",
        }

        return value_map.get(lowered, value)

    def normalize_features(self, input_features):
        if not isinstance(input_features, dict):
            raise TypeError("input_features must be a dictionary.")

        normalized = {}
        unknown_fields = []

        for raw_field, raw_value in input_features.items():
            field = self.canonical_field_name(raw_field)

            if field not in self.schema:
                unknown_fields.append(raw_field)
                continue

            if isinstance(raw_value, list):
                parts = [self.normalize_basic_value(v) for v in raw_value]
                parts = [p for p in parts if p is not None]
                value = "; ".join(parts)
            else:
                parts = self.split_value(raw_value)
                parts = [self.normalize_basic_value(p) for p in parts]
                parts = [p for p in parts if p is not None]
                value = "; ".join(parts)

            if value:
                normalized[field] = value

        return normalized, unknown_fields

    def parse_growth_temperature(self, value):
        raw = self.safe_str(value)

        if raw is None:
            return None, None, "missing"

        text = raw.strip()

        if text.lower() in {"unknown", "none", "variable", "var", "not known", "n/a", "na"}:
            return None, None, "non_numeric"

        if "//" in text:
            left, right = text.split("//", 1)
            left_nums = re.findall(r"-?\d+(?:\.\d+)?", left)
            right_nums = re.findall(r"-?\d+(?:\.\d+)?", right)

            if left_nums and right_nums:
                return float(left_nums[0]), float(right_nums[0]), "parsed_double_slash"

        range_patterns = [
            r"(-?\d+(?:\.\d+)?)\s*[-–—]\s*(-?\d+(?:\.\d+)?)",
            r"(-?\d+(?:\.\d+)?)\s*/\s*(-?\d+(?:\.\d+)?)",
            r"(-?\d+(?:\.\d+)?)\s+to\s+(-?\d+(?:\.\d+)?)",
        ]

        for pattern in range_patterns:
            match = re.search(pattern, text, flags=re.IGNORECASE)

            if match:
                low = float(match.group(1))
                high = float(match.group(2))
                return low, high, "parsed_range"

        nums = re.findall(r"-?\d+(?:\.\d+)?", text)

        if len(nums) >= 2:
            vals = [float(x) for x in nums]
            return min(vals), max(vals), "parsed_multiple_numbers_minmax"

        if len(nums) == 1:
            val = float(nums[0])
            return val, val, "single_number"

        return None, None, "parse_failed"

    def build_tokens_and_aux(self, features):
        tokens = []
        observed_count = 0
        core_observed_count = 0

        growth_low = 0.0
        growth_high = 0.0
        growth_range = 0.0
        growth_mid = 0.0
        growth_numeric_present = 0.0
        grows_at_37 = 0.0
        grows_45_plus = 0.0
        grows_10_or_below = 0.0

        for field in self.schema:
            value = self.safe_str(features.get(field))

            if value is None:
                continue

            value_tokens = self.split_value(value)

            if not value_tokens:
                continue

            observed_count += 1

            if field in self.core_fields:
                core_observed_count += 1

            tokens.append(f"FIELD::{field}")

            for token in value_tokens:
                tokens.append(f"{field}={token}")

            if field in self.semicolon_fields:
                tokens.append(f"FIELD_TOKEN_COUNT::{field}={len(value_tokens)}")

            if field == "Growth Temperature":
                low, high, status = self.parse_growth_temperature(value)
                tokens.append(f"GROWTH_TEMP_PARSE_STATUS={status}")

                if low is not None and high is not None:
                    if high < low:
                        low, high = high, low

                    growth_numeric_present = 1.0
                    growth_low = float(low)
                    growth_high = float(high)
                    growth_range = float(high - low)
                    growth_mid = float((low + high) / 2.0)

                    low_bin = int(math.floor(growth_low / 5.0) * 5)
                    high_bin = int(math.floor(growth_high / 5.0) * 5)
                    mid_bin = int(math.floor(growth_mid / 5.0) * 5)

                    tokens.append(f"GROWTH_TEMP_LOW_BIN={low_bin}")
                    tokens.append(f"GROWTH_TEMP_HIGH_BIN={high_bin}")
                    tokens.append(f"GROWTH_TEMP_MIDPOINT_BIN={mid_bin}")

                    for threshold in self.growth_temp_thresholds:
                        if growth_low <= threshold <= growth_high:
                            tokens.append(f"GROWS_AT_{threshold}C=YES")
                        else:
                            tokens.append(f"GROWS_AT_{threshold}C=NO")

                    grows_at_37 = 1.0 if growth_low <= 37 <= growth_high else 0.0
                    grows_45_plus = 1.0 if growth_high >= 45 else 0.0
                    grows_10_or_below = 1.0 if growth_low <= 10 else 0.0

        aux = np.array([
            observed_count / len(self.schema),
            core_observed_count / len(self.core_fields),
            growth_numeric_present,
            growth_low / 80.0,
            growth_high / 80.0,
            growth_range / 80.0,
            growth_mid / 80.0,
            grows_at_37,
            grows_45_plus,
            grows_10_or_below,
        ], dtype=np.float32)

        return tokens, aux

    def encode_tokens(self, tokens):
        ids = [self.vocab[self.cls_token]]
        unknown_tokens = []

        for token in tokens:
            token_id = self.vocab.get(token)

            if token_id is None:
                token_id = self.vocab[self.unk_token]
                unknown_tokens.append(token)

            ids.append(token_id)

        ids = ids[:self.max_tokens]
        attention_mask = [1] * len(ids)

        while len(ids) < self.max_tokens:
            ids.append(self.vocab[self.pad_token])
            attention_mask.append(0)

        return (
            np.array(ids, dtype=np.int64),
            np.array(attention_mask, dtype=np.int64),
            unknown_tokens,
        )

    @staticmethod
    def confidence_bucket(probability, margin):
        if probability >= 0.95 and margin >= 0.25:
            return "High"
        if probability >= 0.80 and margin >= 0.15:
            return "Medium-high"
        if probability >= 0.50 and margin >= 0.08:
            return "Medium"
        if probability >= 0.30:
            return "Low-medium"
        return "Low"

    @staticmethod
    def distinctness_bucket(margin):
        if margin >= 0.30:
            return "Very distinct"
        if margin >= 0.15:
            return "Fairly distinct"
        if margin >= 0.07:
            return "Ambiguous"
        return "Very ambiguous"

    @torch.no_grad()
    def predict(self, input_features, top_k=10):
        normalized_features, unknown_fields = self.normalize_features(input_features)
        tokens, aux = self.build_tokens_and_aux(normalized_features)

        if len(tokens) == 0:
            raise ValueError("No valid phenotype fields were provided.")

        input_ids_np, attention_mask_np, unknown_tokens = self.encode_tokens(tokens)

        input_ids = torch.tensor(input_ids_np, dtype=torch.long).unsqueeze(0).to(self.device)
        attention_mask = torch.tensor(attention_mask_np, dtype=torch.long).unsqueeze(0).to(self.device)
        aux_tensor = torch.tensor(aux, dtype=torch.float32).unsqueeze(0).to(self.device)

        self.model.eval()

        with torch.autocast(
            device_type=self.device.type,
            dtype=self.amp_dtype,
            enabled=self.use_amp,
        ):
            logits = self.model(input_ids, attention_mask, aux_tensor)

        logits = logits.float()
        calibrated_logits = logits / self.temperature
        probs = torch.softmax(calibrated_logits, dim=-1).squeeze(0).cpu().numpy()

        order = np.argsort(probs)[::-1]
        top_indices = order[:top_k]

        ranked = []

        for rank, idx in enumerate(top_indices, start=1):
            ranked.append({
                "rank": rank,
                "genus": self.classes[int(idx)],
                "probability": float(probs[int(idx)]),
            })

        top1_prob = float(probs[order[0]])
        top2_prob = float(probs[order[1]]) if len(order) > 1 else 0.0
        margin = top1_prob - top2_prob

        provided_fields = list(normalized_features.keys())
        missing_fields = [field for field in self.schema if field not in normalized_features]

        return {
            "top_genus": self.classes[int(order[0])],
            "top_probability": top1_prob,
            "top2_probability": top2_prob,
            "margin": margin,
            "confidence": self.confidence_bucket(top1_prob, margin),
            "distinctness": self.distinctness_bucket(margin),
            "temperature": self.temperature,
            "ranked_genera": ranked,
            "provided_fields": provided_fields,
            "missing_fields": missing_fields,
            "unknown_input_fields_ignored": unknown_fields,
            "unknown_model_tokens": unknown_tokens,
            "num_unknown_model_tokens": len(unknown_tokens),
            "num_provided_fields": len(provided_fields),
            "num_model_tokens": len(tokens),
            "note": "This is a phenotype-based genus prediction, not a confirmed laboratory identification.",
        }


def print_prediction(result):
    print()
    print("=" * 60)
    print("PHENOTYPECLASSIFIER PREDICTION")
    print("=" * 60)
    print(f"Top genus:      {result['top_genus']}")
    print(f"Probability:    {result['top_probability']:.4f}")
    print(f"Margin:         {result['margin']:.4f}")
    print(f"Confidence:     {result['confidence']}")
    print(f"Distinctness:   {result['distinctness']}")
    print(f"Fields used:    {result['num_provided_fields']}")
    print(f"Model tokens:   {result['num_model_tokens']}")
    print(f"Unknown tokens: {result['num_unknown_model_tokens']}")
    print("-" * 60)
    print("Top ranked genera:")

    for item in result["ranked_genera"]:
        print(f"{item['rank']:>2}. {item['genus']:<25} {item['probability']:.4f}")

    print("-" * 60)
    print(result["note"])
    print("=" * 60)