src/syntax_pred/model.py

# src/syntax_pred/model.py — removed sigma & yulie, robust auto-loading, backward-compatible **kwargs
from __future__ import annotations
from typing import Any, Dict, List, Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models.video as tvmv
import lightning.pytorch as pl

class SyntaxLightningModule(pl.LightningModule):
    """
    R3D-18 backbone с настраиваемой "головой" для последовательностей.
    На выходе тензор (B, 2): [логит классификации, регрессия = log(1 + score)].
    """
    def __init__(
        self,
        num_classes: int,
        lr: float,
        variant: str,
        weight_decay: float = 0.0,
        max_epochs: Optional[int] = None,
        weight_path: Optional[str] = None,     # загрузка только бэкбона (опционально)
        pl_weight_path: Optional[str] = None,  # загрузка целого модуля (ckpt/pt)
        rnn_hidden_div: int = 4,
        rnn_dropout: float = 0.2,
        bert_nhead: int = 4,
        bert_layers: int = 1,
        bert_ff_div: int = 4,
        bert_dropout: float = 0.2,
        precision: str = "bf16",
        **_ignore,  # игнорирование лишних/неиспользуемых аргументов в вызове
    ) -> None:
        super().__init__()
        self.save_hyperparameters()
        self.num_classes = num_classes
        self.variant = variant
        self.lr = lr
        self.weight_decay = weight_decay
        self.max_epochs = max_epochs

        # --- Backbone (R3D-18)
        self.model = tvmv.r3d_18(weights=tvmv.R3D_18_Weights.DEFAULT)
        in_features = self.model.fc.in_features
        self.model.fc = nn.Linear(in_features, 2, bias=True)

        if weight_path is not None:
            self._load_backbone_weight(weight_path, self.model)

        # Для последовательных голов используем признаки до FC
        if self.variant != "mean_out":
            self.model.fc = nn.Identity()

        # --- Heads
        if self.variant == "mean_out":
            pass
        elif self.variant in ("gru_mean", "gru_last"):
            self.rnn = nn.GRU(in_features, in_features // rnn_hidden_div, batch_first=True)
            self.dropout = nn.Dropout(rnn_dropout)
            self.fc = nn.Linear(in_features // rnn_hidden_div, num_classes)
        elif self.variant in ("lstm_mean", "lstm_last"):
            self.lstm = nn.LSTM(
                input_size=in_features,
                hidden_size=in_features // rnn_hidden_div,
                proj_size=num_classes,
                batch_first=True,
            )
        elif self.variant == "mean":
            self.fc = nn.Linear(in_features, num_classes)
        elif self.variant in ("bert_mean", "bert_cls", "bert_cls2"):
            enc_layer = nn.TransformerEncoderLayer(
                d_model=in_features,
                nhead=bert_nhead,
                batch_first=True,
                dim_feedforward=in_features // bert_ff_div,
                dropout=bert_dropout,
            )
            self.encoder = nn.TransformerEncoder(enc_layer, num_layers=bert_layers)
            self.dropout = nn.Dropout(bert_dropout)
            self.fc = nn.Linear(in_features, num_classes)
            if self.variant == "bert_cls2":
                self.cls = nn.Parameter(torch.randn(1, 1, in_features))
        else:
            raise ValueError(f"Unknown variant: {self.variant}")

        # Загрузка полного state_dict (поддерживает форматы ckpt/pt)
        if pl_weight_path is not None:
            self._load_full_module(pl_weight_path)

        # Лоссы
        self.loss_clf = nn.BCEWithLogitsLoss(reduction="none")
        self.loss_reg = nn.MSELoss(reduction="none")

        # Кэши для валидационных метрик
        self.y_val: List[int] = []
        self.p_val: List[float] = []
        self.r_val: List[int] = []
        self.ty_val: List[float] = []
        self.tp_val: List[float] = []

    # ---------------- weights -----------------
    @staticmethod
    def _strip_prefix(key: str) -> str:
        for pref in ("model.", "backbone.", "module.", "net."):
            if key.startswith(pref):
                return key[len(pref):]
        return key

    def _load_backbone_weight(self, weight_path: str, model: nn.Module) -> None:
        ckpt = torch.load(weight_path, map_location="cpu")
        if isinstance(ckpt, dict) and "state_dict" in ckpt:
            ckpt = ckpt["state_dict"]
        new_sd = {
            self._strip_prefix(k).replace("fc.", ""): v
            for k, v in ckpt.items()
            if k.startswith(("model.", "backbone.", "module."))
        }
        model.load_state_dict(new_sd, strict=False)

    def _load_block(self, sd: Dict[str, torch.Tensor], module: nn.Module, prefix: str) -> None:
        filtered = {k.replace(f"{prefix}.", ""): v for k, v in sd.items() if k.startswith(prefix)}
        missing, unexpected = module.load_state_dict(filtered, strict=False)
        if missing:
            print(f"[load] Missing for {prefix}: {missing}")
        if unexpected:
            print(f"[load] Unexpected for {prefix}: {unexpected}")

    def _load_full_module(self, path: str) -> None:
        raw = torch.load(path, map_location="cpu")
        sd: Dict[str, torch.Tensor] = raw["state_dict"] if (isinstance(raw, dict) and "state_dict" in raw) else raw
        self._load_block(sd, self.model, "model")
        if self.variant == "mean_out":
            pass
        elif self.variant in ("gru_mean", "gru_last"):
            self._load_block(sd, self.rnn, "rnn")
            self._load_block(sd, self.fc, "fc")
        elif self.variant in ("lstm_mean", "lstm_last"):
            self._load_block(sd, self.lstm, "lstm")
        elif self.variant == "mean":
            self._load_block(sd, self.fc, "fc")
        elif self.variant in ("bert_mean", "bert_cls", "bert_cls2"):
            self._load_block(sd, self.encoder, "encoder")
            self._load_block(sd, self.fc, "fc")
            if self.variant == "bert_cls2" and "cls" in sd:
                with torch.no_grad():
                    self.cls.copy_(sd["cls"])

    # ---------------- forward -----------------
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Подача батча последовательностей формы (B,S,C,T,H,W),
        где S — число клипов внутри исследования для одной артерии.
        """
        b, s = x.shape[:2]
        x = x.flatten(0, 1)   # (B*S, C, T, H, W)
        x = self.model(x)     # (B*S, 2) или (B*S, F) при fc=Identity
        x = x.unflatten(0, (b, s))  # (B, S, ...)

        if self.variant == "mean_out":
            x = x.mean(dim=1)
        elif self.variant in ("gru_mean", "gru_last"):
            all_outs, last = self.rnn(x)
            x = all_outs.mean(dim=1) if self.variant == "gru_mean" else last[0]
            x = self.dropout(x)
            x = self.fc(x)
        elif self.variant in ("lstm_mean", "lstm_last"):
            all_outs, (last_out, _) = self.lstm(x)
            x = all_outs.mean(dim=1) if self.variant == "lstm_mean" else last_out
        elif self.variant == "mean":
            x = x.mean(dim=1)
            x = self.fc(x)
        elif self.variant in ("bert_mean", "bert_cls", "bert_cls2"):
            if self.variant == "bert_cls":
                x = F.pad(x, (0, 0, 1, 0), value=0.0)  # добавление CLS токена
            elif self.variant == "bert_cls2":
                bs = x.size(0)
                x = torch.cat([self.cls.expand(bs, -1, -1), x], dim=1)
            x = self.encoder(x)
            x = x.mean(dim=1) if self.variant == "bert_mean" else x[:, 0, :]
            x = self.dropout(x)
            x = self.fc(x)
        else:
            raise ValueError(self.variant)
        return x  # (B, 2)

    # --------------- training (kept for completeness) ---------------
    def training_step(self, batch, batch_idx):
        x, y, target, _ = batch
        y_hat = self(x)
        yp_clf = y_hat[:, 0:1]
        yp_reg = y_hat[:, 1:]

        weights_clf = torch.where(y > 0, 1.0, 0.2)
        clf_loss = (self.loss_clf(yp_clf, y) * weights_clf).mean()

        reg_loss = self.loss_reg(yp_reg, target).mean()
        return clf_loss + 0.5 * reg_loss

    def configure_optimizers(self):
        params = [p for p in self.parameters() if p.requires_grad]
        opt = torch.optim.Adam(params, lr=self.lr, weight_decay=self.weight_decay)
        if self.max_epochs:
            sch = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=self.lr, total_steps=self.max_epochs)
            return [opt], [sch]
        return opt