ins-pricing 0.3.2__py3-none-any.whl → 0.3.4__py3-none-any.whl

ins_pricing/cli/BayesOpt_entry.py

@@ -6,6 +6,8 @@ The main implementation lives in bayesopt_entry_runner.py.
 from __future__ import annotations
 
 from pathlib import Path
+import json
+import os
 import sys
 
 if __package__ in {None, ""}:
@@ -13,6 +15,36 @@ if __package__ in {None, ""}:
     if str(repo_root) not in sys.path:
         sys.path.insert(0, str(repo_root))
 
+def _apply_env_from_config(argv: list[str]) -> None:
+    if "--config-json" not in argv:
+        return
+    idx = argv.index("--config-json")
+    if idx + 1 >= len(argv):
+        return
+    raw_path = argv[idx + 1]
+    try:
+        cfg_path = Path(raw_path).expanduser()
+        if not cfg_path.is_absolute():
+            cfg_path = cfg_path.resolve()
+        if not cfg_path.exists():
+            script_dir = Path(__file__).resolve().parents[1]
+            candidate = (script_dir / raw_path).resolve()
+            if candidate.exists():
+                cfg_path = candidate
+        if not cfg_path.exists():
+            return
+        cfg = json.loads(cfg_path.read_text(encoding="utf-8", errors="replace"))
+        env = cfg.get("env", {})
+        if isinstance(env, dict):
+            for key, value in env.items():
+                if key is None:
+                    continue
+                os.environ.setdefault(str(key), str(value))
+    except Exception:
+        return
+
+_apply_env_from_config(sys.argv)
+
 try:
     from .bayesopt_entry_runner import main
 except Exception:  # pragma: no cover

ins_pricing/cli/utils/import_resolver.py

@@ -13,6 +13,7 @@ Usage:
 from __future__ import annotations
 
 import importlib
+import os
 import sys
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -70,14 +71,39 @@ class ResolvedImports:
     plot_loss_curve: Optional[Callable] = None
 
 
+def _debug_imports_enabled() -> bool:
+    value = os.environ.get("BAYESOPT_DEBUG_IMPORTS")
+    if value is None:
+        return False
+    return str(value).strip().lower() in {"1", "true", "yes", "y", "on"}
+
+
 def _try_import(module_path: str, attr_name: Optional[str] = None) -> Optional[Any]:
     """Attempt to import a module or attribute, returning None on failure."""
     try:
         module = importlib.import_module(module_path)
         if attr_name:
-            return getattr(module, attr_name, None)
-        return module
-    except Exception:
+            result = getattr(module, attr_name, None)
+        else:
+            result = module
+        if _debug_imports_enabled():
+            origin = getattr(module, "__file__", None)
+            origin = origin or getattr(module, "__path__", None)
+            print(
+                f"[BAYESOPT_DEBUG_IMPORTS] imported {module_path}"
+                f"{'::' + attr_name if attr_name else ''} from {origin}",
+                file=sys.stderr,
+                flush=True,
+            )
+        return result
+    except Exception as exc:
+        if _debug_imports_enabled():
+            print(
+                f"[BAYESOPT_DEBUG_IMPORTS] failed import {module_path}"
+                f"{'::' + attr_name if attr_name else ''}: {exc.__class__.__name__}: {exc}",
+                file=sys.stderr,
+                flush=True,
+            )
         return None
 
 

ins_pricing/cli/utils/notebook_utils.py

@@ -8,9 +8,9 @@ from pathlib import Path
 from typing import Iterable, List, Optional, Sequence, cast
 
 try:
-    from .cli_config import add_config_json_arg  # type: ignore
+    from .cli_config import add_config_json_arg, set_env  # type: ignore
 except Exception:  # pragma: no cover
-    from cli_config import add_config_json_arg  # type: ignore
+    from cli_config import add_config_json_arg, set_env  # type: ignore
 
 
 def _find_ins_pricing_dir(cwd: Optional[Path] = None) -> Path:
@@ -261,6 +261,7 @@ def run_from_config(config_json: str | Path) -> subprocess.CompletedProcess:
     if not config_path.is_absolute():
         config_path = (pkg_dir / config_path).resolve() if (pkg_dir / config_path).exists() else config_path.resolve()
     raw = json.loads(config_path.read_text(encoding="utf-8", errors="replace"))
+    set_env(raw.get("env", {}))
     runner = cast(dict, raw.get("runner") or {})
 
     mode = str(runner.get("mode") or "entry").strip().lower()

ins_pricing/docs/modelling/BayesOpt_USAGE.md

@@ -75,13 +75,13 @@ Under `ins_pricing/modelling/core/bayesopt/`:
 
 1) **Tools and utilities**
 
-- `IOUtils / TrainingUtils / PlotUtils`: I/O, training utilities (batch size, tweedie loss, free_cuda), plotting helpers
+- `IOUtils / TrainingUtils / PlotUtils`: I/O, training utilities (batch size, loss functions, free_cuda), plotting helpers
 - `DistributedUtils`: DDP init, rank/world_size helpers
 
 2) **TorchTrainerMixin (common components for torch tabular training)**
 
 - DataLoader: `_build_dataloader()` / `_build_val_dataloader()` (prints batch/accum/workers)
-- Loss: `_compute_losses()` / `_compute_weighted_loss()` (regression uses tweedie; classification uses BCEWithLogits)
+- Loss: `_compute_losses()` / `_compute_weighted_loss()` (regression supports tweedie/poisson/gamma/mse/mae; classification uses BCEWithLogits)
 - Early stop: `_early_stop_update()`
 
 3) **Sklearn-style model classes (core training objects)**
@@ -292,7 +292,7 @@ FT role is controlled by `ft_role` (from config or CLI `--ft-role`):
 ### 4.1 Supervised models (GLM/XGB/ResNet/FT-as-model)
 
 - `TrainerBase.tune()` calls each trainer's `cross_val()` and minimizes validation metric (default direction `minimize`)
-- Regression typically uses Tweedie deviance or related loss; classification uses logloss
+- Regression loss is configurable (tweedie/poisson/gamma/mse/mae); classification uses logloss
 
 ### 4.2 FT self-supervised (`unsupervised_embedding`)
 

ins_pricing/modelling/core/bayesopt/__init__.py

@@ -4,6 +4,10 @@ from __future__ import annotations
 
 import torch
 
+from ins_pricing.utils.torch_compat import disable_torch_dynamo_if_requested
+
+disable_torch_dynamo_if_requested()
+
 from .config_preprocess import (
     BayesOptConfig,
     DatasetPreprocessor,

ins_pricing/modelling/core/bayesopt/config_preprocess.py

@@ -12,6 +12,7 @@ import pandas as pd
 from sklearn.preprocessing import StandardScaler
 
 from .utils import IOUtils
+from .utils.losses import normalize_loss_name
 from ....exceptions import ConfigurationError, DataValidationError
 
 # NOTE: Some CSV exports may contain invisible BOM characters or leading/trailing
@@ -81,6 +82,7 @@ class BayesOptConfig:
         task_type: Either 'regression' or 'classification'
         binary_resp_nme: Column name for binary response (optional)
         cate_list: List of categorical feature column names
+        loss_name: Regression loss ('auto', 'tweedie', 'poisson', 'gamma', 'mse', 'mae')
         prop_test: Proportion of data for validation (0.0-1.0)
         rand_seed: Random seed for reproducibility
         epochs: Number of training epochs
@@ -117,6 +119,7 @@ class BayesOptConfig:
     task_type: str = 'regression'
     binary_resp_nme: Optional[str] = None
     cate_list: Optional[List[str]] = None
+    loss_name: str = "auto"
 
     # Training configuration
     prop_test: float = 0.25
@@ -207,6 +210,15 @@ class BayesOptConfig:
             errors.append(
                 f"task_type must be one of {valid_task_types}, got '{self.task_type}'"
             )
+        # Validate loss_name
+        try:
+            normalized_loss = normalize_loss_name(self.loss_name, self.task_type)
+            if self.task_type == "classification" and normalized_loss not in {"auto", "logloss", "bce"}:
+                errors.append(
+                    "loss_name must be 'auto', 'logloss', or 'bce' for classification tasks."
+                )
+        except ValueError as exc:
+            errors.append(str(exc))
 
         # Validate prop_test
         if not 0.0 < self.prop_test < 1.0:

ins_pricing/modelling/core/bayesopt/core.py

@@ -17,6 +17,12 @@ from .model_plotting_mixin import BayesOptPlottingMixin
 from .models import GraphNeuralNetSklearn
 from .trainers import FTTrainer, GLMTrainer, GNNTrainer, ResNetTrainer, XGBTrainer
 from .utils import EPS, infer_factor_and_cate_list, set_global_seed
+from .utils.losses import (
+    infer_loss_name_from_model_name,
+    normalize_loss_name,
+    resolve_tweedie_power,
+    resolve_xgb_objective,
+)
 
 
 class _CVSplitter:
@@ -293,6 +299,14 @@ class BayesOptModel(BayesOptPlottingMixin, BayesOptExplainMixin):
         self.config = cfg
         self.model_nme = cfg.model_nme
         self.task_type = cfg.task_type
+        normalized_loss = normalize_loss_name(getattr(cfg, "loss_name", None), self.task_type)
+        if self.task_type == "classification":
+            self.loss_name = "logloss" if normalized_loss == "auto" else normalized_loss
+        else:
+            if normalized_loss == "auto":
+                self.loss_name = infer_loss_name_from_model_name(self.model_nme)
+            else:
+                self.loss_name = normalized_loss
         self.resp_nme = cfg.resp_nme
         self.weight_nme = cfg.weight_nme
         self.factor_nmes = cfg.factor_nmes
@@ -339,14 +353,7 @@ class BayesOptModel(BayesOptPlottingMixin, BayesOptExplainMixin):
         if self.task_type == 'classification':
             self.obj = 'binary:logistic'
         else:  # regression task
-            if 'f' in self.model_nme:
-                self.obj = 'count:poisson'
-            elif 's' in self.model_nme:
-                self.obj = 'reg:gamma'
-            elif 'bc' in self.model_nme:
-                self.obj = 'reg:tweedie'
-            else:
-                self.obj = 'reg:tweedie'
+            self.obj = resolve_xgb_objective(self.loss_name)
         self.fit_params = {
             'sample_weight': self.train_data[self.weight_nme].values
         }
@@ -426,6 +433,11 @@ class BayesOptModel(BayesOptPlottingMixin, BayesOptExplainMixin):
     def default_tweedie_power(self, obj: Optional[str] = None) -> Optional[float]:
         if self.task_type == 'classification':
             return None
+        loss_name = getattr(self, "loss_name", None)
+        if loss_name:
+            resolved = resolve_tweedie_power(str(loss_name), default=1.5)
+            if resolved is not None:
+                return resolved
         objective = obj or getattr(self, "obj", None)
         if objective == 'count:poisson':
             return 1.0
@@ -503,6 +515,7 @@ class BayesOptModel(BayesOptPlottingMixin, BayesOptExplainMixin):
                 patience=5,
                 task_type=self.task_type,
                 tweedie_power=tw_power,
+                loss_name=self.loss_name,
                 use_data_parallel=False,
                 use_ddp=False,
                 use_approx_knn=self.config.gnn_use_approx_knn,

ins_pricing/modelling/core/bayesopt/models/model_ft_components.py

@@ -106,31 +106,58 @@ class ScaledTransformerEncoderLayer(nn.Module):
         self.res_scale_attn = residual_scale_attn
         self.res_scale_ffn = residual_scale_ffn
 
-    def forward(self, src, src_mask=None, src_key_padding_mask=None):
+    def forward(self, src, src_mask=None, src_key_padding_mask=None, is_causal: Optional[bool] = None, **_kwargs):
         # Input tensor shape: (batch, seq_len, d_model).
         x = src
 
         if self.norm_first:
             # Pre-norm before attention.
-            x = x + self._sa_block(self.norm1(x), src_mask,
-                                   src_key_padding_mask)
+            x = x + self._sa_block(
+                self.norm1(x),
+                src_mask,
+                src_key_padding_mask,
+                is_causal=is_causal,
+            )
             x = x + self._ff_block(self.norm2(x))
         else:
             # Post-norm (usually disabled).
             x = self.norm1(
-                x + self._sa_block(x, src_mask, src_key_padding_mask))
+                x + self._sa_block(
+                    x,
+                    src_mask,
+                    src_key_padding_mask,
+                    is_causal=is_causal,
+                )
+            )
             x = self.norm2(x + self._ff_block(x))
 
         return x
 
-    def _sa_block(self, x, attn_mask, key_padding_mask):
+    def _sa_block(self, x, attn_mask, key_padding_mask, *, is_causal: Optional[bool] = None):
         # Self-attention with residual scaling.
-        attn_out, _ = self.self_attn(
-            x, x, x,
-            attn_mask=attn_mask,
-            key_padding_mask=key_padding_mask,
-            need_weights=False
-        )
+        if is_causal is None:
+            attn_out, _ = self.self_attn(
+                x, x, x,
+                attn_mask=attn_mask,
+                key_padding_mask=key_padding_mask,
+                need_weights=False,
+            )
+        else:
+            try:
+                attn_out, _ = self.self_attn(
+                    x, x, x,
+                    attn_mask=attn_mask,
+                    key_padding_mask=key_padding_mask,
+                    need_weights=False,
+                    is_causal=is_causal,
+                )
+            except TypeError:
+                attn_out, _ = self.self_attn(
+                    x, x, x,
+                    attn_mask=attn_mask,
+                    key_padding_mask=key_padding_mask,
+                    need_weights=False,
+                )
         return self.res_scale_attn * self.dropout1(attn_out)
 
     def _ff_block(self, x):
@@ -313,4 +340,3 @@ class MaskedTabularDataset(Dataset):
             None if self.X_cat_true is None else self.X_cat_true[idx],
             None if self.cat_mask is None else self.cat_mask[idx],
         )
-

ins_pricing/modelling/core/bayesopt/models/model_ft_trainer.py

@@ -16,6 +16,11 @@ from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.nn.utils import clip_grad_norm_
 
 from ..utils import DistributedUtils, EPS, TorchTrainerMixin
+from ..utils.losses import (
+    infer_loss_name_from_model_name,
+    normalize_loss_name,
+    resolve_tweedie_power,
+)
 from .model_ft_components import FTTransformerCore, MaskedTabularDataset, TabularDataset
 
 
@@ -159,7 +164,8 @@ class FTTransformerSklearn(TorchTrainerMixin, nn.Module):
                  weight_decay: float = 0.0,
                  use_data_parallel: bool = True,
                  use_ddp: bool = False,
-                 num_numeric_tokens: Optional[int] = None
+                 num_numeric_tokens: Optional[int] = None,
+                 loss_name: Optional[str] = None
                  ):
         super().__init__()
 
@@ -187,14 +193,18 @@ class FTTransformerSklearn(TorchTrainerMixin, nn.Module):
         self.weight_decay = weight_decay
         self.task_type = task_type
         self.patience = patience
+        resolved_loss = normalize_loss_name(loss_name, self.task_type)
         if self.task_type == 'classification':
+            self.loss_name = "logloss"
             self.tw_power = None  # No Tweedie power for classification.
-        elif 'f' in self.model_nme:
-            self.tw_power = 1.0
-        elif 's' in self.model_nme:
-            self.tw_power = 2.0
         else:
-            self.tw_power = tweedie_power
+            if resolved_loss == "auto":
+                resolved_loss = infer_loss_name_from_model_name(self.model_nme)
+            self.loss_name = resolved_loss
+            if self.loss_name == "tweedie":
+                self.tw_power = float(tweedie_power) if tweedie_power is not None else 1.5
+            else:
+                self.tw_power = resolve_tweedie_power(self.loss_name, default=1.5)
 
         if self.is_ddp_enabled:
             self.device = torch.device(f"cuda:{self.local_rank}")

ins_pricing/modelling/core/bayesopt/models/model_gnn.py

@@ -17,6 +17,11 @@ from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.nn.utils import clip_grad_norm_
 
 from ..utils import DistributedUtils, EPS, IOUtils, TorchTrainerMixin
+from ..utils.losses import (
+    infer_loss_name_from_model_name,
+    normalize_loss_name,
+    resolve_tweedie_power,
+)
 
 try:
     from torch_geometric.nn import knn_graph
@@ -109,7 +114,8 @@ class GraphNeuralNetSklearn(TorchTrainerMixin, nn.Module):
                  max_gpu_knn_nodes: Optional[int] = None,
                  knn_gpu_mem_ratio: float = 0.9,
                  knn_gpu_mem_overhead: float = 2.0,
-                 knn_cpu_jobs: Optional[int] = -1) -> None:
+                 knn_cpu_jobs: Optional[int] = -1,
+                 loss_name: Optional[str] = None) -> None:
         super().__init__()
         self.model_nme = model_nme
         self.input_dim = input_dim
@@ -139,14 +145,18 @@ class GraphNeuralNetSklearn(TorchTrainerMixin, nn.Module):
         self._adj_cache_key: Optional[Tuple[Any, ...]] = None
         self._adj_cache_tensor: Optional[torch.Tensor] = None
 
+        resolved_loss = normalize_loss_name(loss_name, self.task_type)
         if self.task_type == 'classification':
+            self.loss_name = "logloss"
             self.tw_power = None
-        elif 'f' in self.model_nme:
-            self.tw_power = 1.0
-        elif 's' in self.model_nme:
-            self.tw_power = 2.0
         else:
-            self.tw_power = tweedie_power
+            if resolved_loss == "auto":
+                resolved_loss = infer_loss_name_from_model_name(self.model_nme)
+            self.loss_name = resolved_loss
+            if self.loss_name == "tweedie":
+                self.tw_power = float(tweedie_power) if tweedie_power is not None else 1.5
+            else:
+                self.tw_power = resolve_tweedie_power(self.loss_name, default=1.5)
 
         self.ddp_enabled = False
         self.local_rank = int(os.environ.get("LOCAL_RANK", 0))

ins_pricing/modelling/core/bayesopt/models/model_resn.py

@@ -12,6 +12,11 @@ from torch.nn.utils import clip_grad_norm_
 from torch.utils.data import TensorDataset
 
 from ..utils import DistributedUtils, EPS, TorchTrainerMixin
+from ..utils.losses import (
+    infer_loss_name_from_model_name,
+    normalize_loss_name,
+    resolve_tweedie_power,
+)
 
 
 # =============================================================================
@@ -140,7 +145,8 @@ class ResNetSklearn(TorchTrainerMixin, nn.Module):
                  stochastic_depth: float = 0.0,
                  weight_decay: float = 1e-4,
                  use_data_parallel: bool = True,
-                 use_ddp: bool = False):
+                 use_ddp: bool = False,
+                 loss_name: Optional[str] = None):
         super(ResNetSklearn, self).__init__()
 
         self.use_ddp = use_ddp
@@ -179,15 +185,18 @@ class ResNetSklearn(TorchTrainerMixin, nn.Module):
         else:
             self.device = torch.device('cpu')
 
-        # Tweedie power (unused for classification)
+        resolved_loss = normalize_loss_name(loss_name, self.task_type)
         if self.task_type == 'classification':
+            self.loss_name = "logloss"
             self.tw_power = None
-        elif 'f' in self.model_nme:
-            self.tw_power = 1
-        elif 's' in self.model_nme:
-            self.tw_power = 2
         else:
-            self.tw_power = tweedie_power
+            if resolved_loss == "auto":
+                resolved_loss = infer_loss_name_from_model_name(self.model_nme)
+            self.loss_name = resolved_loss
+            if self.loss_name == "tweedie":
+                self.tw_power = float(tweedie_power) if tweedie_power is not None else 1.5
+            else:
+                self.tw_power = resolve_tweedie_power(self.loss_name, default=1.5)
 
         # Build network (construct on CPU first)
         core = ResNetSequential(

ins_pricing/modelling/core/bayesopt/trainers/trainer_base.py

@@ -578,6 +578,7 @@ class TrainerBase:
                     "n_layers": getattr(self.model, "n_layers", 4),
                     "dropout": getattr(self.model, "dropout", 0.1),
                     "task_type": getattr(self.model, "task_type", "regression"),
+                    "loss_name": getattr(self.model, "loss_name", None),
                     "tw_power": getattr(self.model, "tw_power", 1.5),
                     "num_geo": getattr(self.model, "num_geo", 0),
                     "num_numeric_tokens": getattr(self.model, "num_numeric_tokens", None),
@@ -638,6 +639,7 @@ class TrainerBase:
                             n_layers=model_config.get("n_layers", 4),
                             dropout=model_config.get("dropout", 0.1),
                             task_type=model_config.get("task_type", "regression"),
+                            loss_name=model_config.get("loss_name", None),
                             tweedie_power=model_config.get("tw_power", 1.5),
                             num_numeric_tokens=model_config.get("num_numeric_tokens"),
                             use_data_parallel=False,

ins_pricing/modelling/core/bayesopt/trainers/trainer_ft.py

@@ -5,11 +5,12 @@ from typing import Any, Dict, List, Optional, Tuple
 import numpy as np
 import optuna
 import pandas as pd
-from sklearn.metrics import log_loss, mean_tweedie_deviance
+from sklearn.metrics import log_loss
 from sklearn.model_selection import GroupKFold, TimeSeriesSplit
 
 from .trainer_base import TrainerBase
 from ..models import FTTransformerSklearn
+from ..utils.losses import regression_loss
 
 class FTTrainer(TrainerBase):
     def __init__(self, context: "BayesOptModel") -> None:
@@ -67,6 +68,7 @@ class FTTrainer(TrainerBase):
 
     def cross_val_unsupervised(self, trial: Optional[optuna.trial.Trial]) -> float:
         """Optuna objective A: minimize validation loss for masked reconstruction."""
+        loss_name = getattr(self.ctx, "loss_name", "tweedie")
         param_space: Dict[str, Callable[[optuna.trial.Trial], Any]] = {
             "learning_rate": lambda t: t.suggest_float('learning_rate', 1e-5, 5e-3, log=True),
             "d_model": lambda t: t.suggest_int('d_model', 16, 128, step=16),
@@ -159,6 +161,7 @@ class FTTrainer(TrainerBase):
             use_data_parallel=self.ctx.config.use_ft_data_parallel,
             use_ddp=self.ctx.config.use_ft_ddp,
             num_numeric_tokens=num_numeric_tokens,
+            loss_name=loss_name,
         )
         model.set_params(model_params)
         try:
@@ -191,7 +194,8 @@ class FTTrainer(TrainerBase):
             "dropout": lambda t: t.suggest_float('dropout', 0.0, 0.2),
             "weight_decay": lambda t: t.suggest_float('weight_decay', 1e-6, 1e-2, log=True),
         }
-        if self.ctx.task_type == 'regression' and self.ctx.obj == 'reg:tweedie':
+        loss_name = getattr(self.ctx, "loss_name", "tweedie")
+        if self.ctx.task_type == 'regression' and loss_name == 'tweedie':
             param_space["tw_power"] = lambda t: t.suggest_float(
                 'tw_power', 1.0, 2.0)
         geo_enabled = bool(
@@ -231,10 +235,12 @@ class FTTrainer(TrainerBase):
             tw_power = params.get("tw_power")
             if self.ctx.task_type == 'regression':
                 base_tw = self.ctx.default_tweedie_power()
-                if self.ctx.obj in ('count:poisson', 'reg:gamma'):
-                    tw_power = base_tw
-                elif tw_power is None:
+                if loss_name == "tweedie":
+                    tw_power = base_tw if tw_power is None else tw_power
+                elif loss_name in ("poisson", "gamma"):
                     tw_power = base_tw
+                else:
+                    tw_power = None
             metric_ctx["tw_power"] = tw_power
 
             adaptive_heads, _ = self._resolve_adaptive_heads(
@@ -259,6 +265,7 @@ class FTTrainer(TrainerBase):
                 use_data_parallel=self.ctx.config.use_ft_data_parallel,
                 use_ddp=self.ctx.config.use_ft_ddp,
                 num_numeric_tokens=num_numeric_tokens,
+                loss_name=loss_name,
             ).set_params({"_geo_params": geo_params_local} if geo_enabled else {})
 
         def fit_predict(model, X_train, y_train, w_train, X_val, y_val, w_val, trial_obj):
@@ -286,11 +293,12 @@ class FTTrainer(TrainerBase):
 
         def metric_fn(y_true, y_pred, weight):
             if self.ctx.task_type == 'regression':
-                return mean_tweedie_deviance(
+                return regression_loss(
                     y_true,
                     y_pred,
-                    sample_weight=weight,
-                    power=metric_ctx.get("tw_power", 1.5)
+                    weight,
+                    loss_name=loss_name,
+                    tweedie_power=metric_ctx.get("tw_power", 1.5),
                 )
             return log_loss(y_true, y_pred, sample_weight=weight)
 
@@ -313,6 +321,7 @@ class FTTrainer(TrainerBase):
     def train(self) -> None:
         if not self.best_params:
             raise RuntimeError("Run tune() first to obtain best FT-Transformer parameters.")
+        loss_name = getattr(self.ctx, "loss_name", "tweedie")
         resolved_params = dict(self.best_params)
         d_model_value = resolved_params.get("d_model", 64)
         adaptive_heads, heads_adjusted = self._resolve_adaptive_heads(
@@ -342,6 +351,7 @@ class FTTrainer(TrainerBase):
                 use_ddp=self.ctx.config.use_ft_ddp,
                 num_numeric_tokens=self._resolve_numeric_tokens(),
                 weight_decay=float(resolved_params.get("weight_decay", 0.0)),
+                loss_name=loss_name,
             )
             tmp_model.set_params(resolved_params)
             geo_train_full = self.ctx.train_geo_tokens
@@ -375,6 +385,7 @@ class FTTrainer(TrainerBase):
             use_ddp=self.ctx.config.use_ft_ddp,
             num_numeric_tokens=self._resolve_numeric_tokens(),
             weight_decay=float(resolved_params.get("weight_decay", 0.0)),
+            loss_name=loss_name,
         )
         if refit_epochs is not None:
             self.model.epochs = int(refit_epochs)
@@ -408,6 +419,7 @@ class FTTrainer(TrainerBase):
     def ensemble_predict(self, k: int) -> None:
         if not self.best_params:
             raise RuntimeError("Run tune() first to obtain best FT-Transformer parameters.")
+        loss_name = getattr(self.ctx, "loss_name", "tweedie")
         k = max(2, int(k))
         X_all = self.ctx.train_data[self.ctx.factor_nmes]
         y_all = self.ctx.train_data[self.ctx.resp_nme]
@@ -446,6 +458,7 @@ class FTTrainer(TrainerBase):
                 use_ddp=self.ctx.config.use_ft_ddp,
                 num_numeric_tokens=self._resolve_numeric_tokens(),
                 weight_decay=float(resolved_params.get("weight_decay", 0.0)),
+                loss_name=loss_name,
             )
             model.set_params(resolved_params)
 
@@ -541,6 +554,7 @@ class FTTrainer(TrainerBase):
         return splitter, None, oof_folds
 
     def _build_ft_feature_model(self, resolved_params: Dict[str, Any]) -> FTTransformerSklearn:
+        loss_name = getattr(self.ctx, "loss_name", "tweedie")
         model = FTTransformerSklearn(
             model_nme=self.ctx.model_nme,
             num_cols=self.ctx.num_features,
@@ -549,6 +563,7 @@ class FTTrainer(TrainerBase):
             use_data_parallel=self.ctx.config.use_ft_data_parallel,
             use_ddp=self.ctx.config.use_ft_ddp,
             num_numeric_tokens=self._resolve_numeric_tokens(),
+            loss_name=loss_name,
         )
         adaptive_heads, heads_adjusted = self._resolve_adaptive_heads(
             d_model=resolved_params.get("d_model", model.d_model),
@@ -702,6 +717,7 @@ class FTTrainer(TrainerBase):
                                          num_loss_weight: float = 1.0,
                                          cat_loss_weight: float = 1.0) -> None:
         """Self-supervised pretraining (masked reconstruction) and cache embeddings."""
+        loss_name = getattr(self.ctx, "loss_name", "tweedie")
         self.model = FTTransformerSklearn(
             model_nme=self.ctx.model_nme,
             num_cols=self.ctx.num_features,
@@ -710,6 +726,7 @@ class FTTrainer(TrainerBase):
             use_data_parallel=self.ctx.config.use_ft_data_parallel,
             use_ddp=self.ctx.config.use_ft_ddp,
             num_numeric_tokens=self._resolve_numeric_tokens(),
+            loss_name=loss_name,
         )
         resolved_params = dict(params or {})
         # Reuse supervised tuning structure params unless explicitly overridden.