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

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/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_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.

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

@@ -6,10 +6,11 @@ import numpy as np
 import optuna
 import pandas as pd
 import statsmodels.api as sm
-from sklearn.metrics import log_loss, mean_tweedie_deviance
+from sklearn.metrics import log_loss
 
 from .trainer_base import TrainerBase
 from ..utils import EPS
+from ..utils.losses import regression_loss
 
 class GLMTrainer(TrainerBase):
     def __init__(self, context: "BayesOptModel") -> None:
@@ -19,10 +20,13 @@ class GLMTrainer(TrainerBase):
     def _select_family(self, tweedie_power: Optional[float] = None):
         if self.ctx.task_type == 'classification':
             return sm.families.Binomial()
-        if self.ctx.obj == 'count:poisson':
+        loss_name = getattr(self.ctx, "loss_name", "tweedie")
+        if loss_name == "poisson":
             return sm.families.Poisson()
-        if self.ctx.obj == 'reg:gamma':
+        if loss_name == "gamma":
             return sm.families.Gamma()
+        if loss_name in {"mse", "mae"}:
+            return sm.families.Gaussian()
         power = tweedie_power if tweedie_power is not None else 1.5
         return sm.families.Tweedie(var_power=power, link=sm.families.links.log())
 
@@ -45,7 +49,8 @@ class GLMTrainer(TrainerBase):
             "alpha": lambda t: t.suggest_float('alpha', 1e-6, 1e2, log=True),
             "l1_ratio": lambda t: t.suggest_float('l1_ratio', 0.0, 1.0)
         }
-        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["tweedie_power"] = lambda t: t.suggest_float(
                 'tweedie_power', 1.0, 2.0)
 
@@ -87,13 +92,12 @@ class GLMTrainer(TrainerBase):
             if self.ctx.task_type == 'classification':
                 y_pred_clipped = np.clip(y_pred, EPS, 1 - EPS)
                 return log_loss(y_true, y_pred_clipped, sample_weight=weight)
-            y_pred_safe = np.maximum(y_pred, EPS)
-            return mean_tweedie_deviance(
+            return regression_loss(
                 y_true,
-                y_pred_safe,
-                sample_weight=weight,
-                power=self._metric_power(
-                    metric_ctx.get("family"), metric_ctx.get("tweedie_power"))
+                y_pred,
+                weight,
+                loss_name=loss_name,
+                tweedie_power=metric_ctx.get("tweedie_power"),
             )
 
         return self.cross_val_generic(
@@ -192,4 +196,3 @@ class GLMTrainer(TrainerBase):
         preds_test = preds_test_sum / float(split_count)
         self._cache_predictions("glm", preds_train, preds_test)
 
-

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

@@ -6,11 +6,12 @@ from typing import Any, Dict, List, Optional, Tuple
 import numpy as np
 import optuna
 import torch
-from sklearn.metrics import log_loss, mean_tweedie_deviance
+from sklearn.metrics import log_loss
 
 from .trainer_base import TrainerBase
 from ..models import GraphNeuralNetSklearn
 from ..utils import EPS
+from ..utils.losses import regression_loss
 from ins_pricing.utils import get_logger
 from ins_pricing.utils.torch_compat import torch_load
 
@@ -25,6 +26,15 @@ class GNNTrainer(TrainerBase):
     def _build_model(self, params: Optional[Dict[str, Any]] = None) -> GraphNeuralNetSklearn:
         params = params or {}
         base_tw_power = self.ctx.default_tweedie_power()
+        loss_name = getattr(self.ctx, "loss_name", "tweedie")
+        tw_power = params.get("tw_power")
+        if self.ctx.task_type == "regression":
+            if loss_name == "tweedie":
+                tw_power = base_tw_power if tw_power is None else float(tw_power)
+            elif loss_name in ("poisson", "gamma"):
+                tw_power = base_tw_power
+            else:
+                tw_power = None
         model = GraphNeuralNetSklearn(
             model_nme=f"{self.ctx.model_nme}_gnn",
             input_dim=len(self.ctx.var_nmes),
@@ -36,7 +46,7 @@ class GNNTrainer(TrainerBase):
             epochs=int(params.get("epochs", self.ctx.epochs)),
             patience=int(params.get("patience", 5)),
             task_type=self.ctx.task_type,
-            tweedie_power=float(params.get("tw_power", base_tw_power or 1.5)),
+            tweedie_power=tw_power,
             weight_decay=float(params.get("weight_decay", 0.0)),
             use_data_parallel=bool(self.ctx.config.use_gnn_data_parallel),
             use_ddp=bool(self.ctx.config.use_gnn_ddp),
@@ -47,11 +57,13 @@ class GNNTrainer(TrainerBase):
             knn_gpu_mem_ratio=float(self.ctx.config.gnn_knn_gpu_mem_ratio),
             knn_gpu_mem_overhead=float(
                 self.ctx.config.gnn_knn_gpu_mem_overhead),
+            loss_name=loss_name,
         )
         return model
 
     def cross_val(self, trial: optuna.trial.Trial) -> float:
         base_tw_power = self.ctx.default_tweedie_power()
+        loss_name = getattr(self.ctx, "loss_name", "tweedie")
         metric_ctx: Dict[str, Any] = {}
 
         def data_provider():
@@ -60,8 +72,16 @@ class GNNTrainer(TrainerBase):
             return data[self.ctx.var_nmes], data[self.ctx.resp_nme], data[self.ctx.weight_nme]
 
         def model_builder(params: Dict[str, Any]):
-            tw_power = params.get("tw_power", base_tw_power)
+            if loss_name == "tweedie":
+                tw_power = params.get("tw_power", base_tw_power)
+            elif loss_name in ("poisson", "gamma"):
+                tw_power = base_tw_power
+            else:
+                tw_power = None
             metric_ctx["tw_power"] = tw_power
+            if tw_power is None:
+                params = dict(params)
+                params.pop("tw_power", None)
             return self._build_model(params)
 
         def preprocess_fn(X_train, X_val):
@@ -85,13 +105,12 @@ class GNNTrainer(TrainerBase):
             if self.ctx.task_type == 'classification':
                 y_pred_clipped = np.clip(y_pred, EPS, 1 - EPS)
                 return log_loss(y_true, y_pred_clipped, sample_weight=weight)
-            y_pred_safe = np.maximum(y_pred, EPS)
-            power = metric_ctx.get("tw_power", base_tw_power or 1.5)
-            return mean_tweedie_deviance(
+            return regression_loss(
                 y_true,
-                y_pred_safe,
-                sample_weight=weight,
-                power=power,
+                y_pred,
+                weight,
+                loss_name=loss_name,
+                tweedie_power=metric_ctx.get("tw_power", base_tw_power),
             )
 
         # Keep GNN BO lightweight: sample during CV, use full data for final training.
@@ -106,7 +125,7 @@ class GNNTrainer(TrainerBase):
             "dropout": lambda t: t.suggest_float('dropout', 0.0, 0.3),
             "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':
+        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)