1{
2  "op": "replace",
3  "end_line": 103,
4  "start_line": 16,
5  "filename": "qlib/custom_model.py",
6  "content": "# =====================================================================\n# EDITABLE: CustomModel — implement your stock prediction model here\n# =====================================================================\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\n\nclass TemporalDataset(Dataset):\n    def __init__(self, features, labels):\n        self.features = torch.tensor(features, dtype=torch.float)\n        self.labels = torch.tensor(labels, dtype=torch.float)\n    def __len__(self): return len(self.features)\n    def __getitem__(self, idx): return self.features[idx], self.labels[idx]\n\nclass Extractor(nn.Module):\n    def __init__(self, input_dim, hidden_dim, dropout=0.3):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(input_dim, hidden_dim),\n            nn.LayerNorm(hidden_dim),\n            nn.GELU(),\n            nn.Dropout(dropout),\n            nn.Linear(hidden_dim, hidden_dim),\n            nn.LayerNorm(hidden_dim),\n            nn.GELU(),\n            nn.Dropout(dropout),\n        )\n    def forward(self, x): return self.net(x)\n\nclass Predictor(nn.Module):\n    def __init__(self, hidden_dim):\n        super().__init__()\n        self.net = nn.Linear(hidden_dim, 1)\n    def forward(self, x): return self.net(x)\n\ndef coral_loss(h1, h2):\n    if h1.shape[0] < 2 or h2.shape[0] < 2:\n        return torch.tensor(0.0).to(h1.device)\n    h1_c = h1 - torch.mean(h1, dim=0, keepdim=True)\n    h2_c = h2 - torch.mean(h2, dim=0, keepdim=True)\n    # Adding epsilon to denominator\n    cov1 = (h1_c.t() @ h1_c) / (h1.shape[0] - 1.0 + 1e-6)\n    cov2 = (h2_c.t() @ h2_c) / (h2.shape[0] - 1.0 + 1e-6)\n    return torch.mean((cov1 - cov2)**2)\n\ndef pcc_loss(pred, target):\n    p_m = pred - pred.mean()\n    t_m = target - target.mean()\n    cov = torch.sum(p_m * t_m)\n    std_p = torch.sqrt(torch.sum(p_m**2) + 1e-6)\n    std_t = torch.sqrt(torch.sum(t_m**2) + 1e-6)\n    return 1.0 - (cov / (std_p * std_t))\n\nclass CustomModel(Model):\n    def __init__(self, input_dim=158, hidden_dim=256, total_steps=2000, \n                 batch_size=512, lr=1e-3, k_domains=4, beta=5.0, lambda_coral=1.0, num_domains=14):\n        super().__init__()\n        self.extractor = Extractor(input_dim, hidden_dim).to(DEVICE)\n        self.predictor = Predictor(hidden_dim).to(DEVICE)\n        self.total_steps = total_steps\n        self.batch_size = batch_size\n        self.lr = lr\n        self.k_domains = k_domains\n        self.beta = beta\n        self.lambda_coral = lambda_coral\n        self.num_domains = num_domains\n        self.fitted = False\n\n    def fit(self, dataset: DatasetH):\n        df_train = dataset.prepare(\"train\", col_set=[\"feature\", \"label\"], data_key=DataHandlerLP.DK_L)\n        features = df_train[\"feature\"].values\n        labels = df_train[\"label\"].values.ravel()\n        mask = ~(np.isnan(features).any(axis=1) | np.isnan(labels))\n        features = features[mask]\n        labels = labels[mask]\n        times = df_train.index.get_level_values('datetime')[mask]\n        \n        unique_times = np.sort(np.unique(times))\n        domain_splits = np.array_split(unique_times, self.num_domains)\n        \n        domain_dataloaders = []\n        for split_times in domain_splits:\n            idx = np.isin(times, split_times)\n            pts = features[idx]\n            lbls = labels[idx]\n            if len(pts) > 0:\n                ds = TemporalDataset(pts, lbls)\n                dl = DataLoader(ds, batch_size=self.batch_size, shuffle=True, drop_last=True)\n                if len(dl) > 0:\n                    domain_dataloaders.append(dl)\n                    \n        M = len(domain_dataloaders)\n        optimizer = optim.AdamW(list(self.extractor.parameters()) + list(self.predictor.parameters()), \n                                lr=self.lr, weight_decay=1e-4)\n        lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=self.total_steps)\n        \n        self.extractor.train()\n        self.predictor.train()\n        \n        iters = [iter(dl) for dl in domain_dataloaders]\n        \n        for step in range(self.total_steps):\n            sampled_indices = np.random.choice(M, min(self.k_domains, M), replace=False)\n            losses = []\n            features_list = []\n            \n            optimizer.zero_grad()\n            for idx in sampled_indices:\n                try:\n                    x, y = next(iters[idx])\n                except StopIteration:\n                    iters[idx] = iter(domain_dataloaders[idx])\n                    x, y = next(iters[idx])\n                    \n                x, y = x.to(DEVICE), y.to(DEVICE)\n                h = self.extractor(x)\n                p = self.predictor(h).squeeze()\n                \n                loss = F.mse_loss(p, y) + 0.5 * pcc_loss(p, y)\n                losses.append(loss)\n                features_list.append(h)\n                \n            losses_stack = torch.stack(losses)\n            risk_mean = losses_stack.mean()\n            risk_var = losses_stack.var() if len(losses) > 1 else 0.0\n            \n            coral_penalty = 0.0\n            if len(features_list) > 1 and self.lambda_coral > 0:\n                count = 0\n                for i in range(len(features_list)):\n                    for j in range(i+1, len(features_list)):\n                        coral_penalty += coral_loss(features_list[i], features_list[j])\n                        count += 1\n                coral_penalty /= count\n            \n            total_loss = risk_mean + self.beta * risk_var + self.lambda_coral * coral_penalty\n            total_loss.backward()\n            torch.nn.utils.clip_grad_norm_(self.extractor.parameters(), 1.0)\n            torch.nn.utils.clip_grad_norm_(self.predictor.parameters(), 1.0)\n            optimizer.step()\n            lr_scheduler.step()\n            \n        self.fitted = True\n\n    def predict(self, dataset: DatasetH, segment=\"test\"):\n        if not self.fitted:\n            raise ValueError(\"Model is not fitted yet!\")\n        \n        df_test = dataset.prepare(segment, col_set=[\"feature\", \"label\"], data_key=DataHandlerLP.DK_I)\n        features = df_test[\"feature\"]\n        index = features.index\n        \n        features_np = features.values\n        features_np = np.nan_to_num(features_np, nan=0.0)\n        \n        self.extractor.eval()\n        self.predictor.eval()\n        \n        preds = []\n        bs = 2048\n        with torch.no_grad():\n            for i in range(0, len(features_np), bs):\n                x = torch.tensor(features_np[i:i+bs], dtype=torch.float).to(DEVICE)\n                h = self.extractor(x)\n                p = self.predictor(h).squeeze(-1)\n                preds.append(p.cpu().numpy())\n                \n        preds = np.concatenate(preds)\n        if preds.ndim == 0:\n            preds = np.expand_dims(preds, 0)\n        return pd.Series(preds, index=index, name=\"score\")"
7}

1{
2  "op": "replace",
3  "filename": "qlib/custom_model.py",
4  "end_line": 103,
5  "start_line": 16,
6  "content": "import copy\n\nclass MSCMoE(nn.Module):\n    def __init__(self, in_d=158, hid_d=128, n_exp=4):\n        super().__init__()\n        self.enc = nn.Sequential(nn.Linear(in_d, hid_d), nn.LayerNorm(hid_d), nn.SELU(), \n                                 nn.Dropout(0.2), nn.Linear(hid_d, hid_d), nn.LayerNorm(hid_d), nn.SELU())\n        self.exps = nn.ModuleList([nn.Sequential(nn.Linear(hid_d, hid_d//2), nn.SELU(), \n                                                 nn.Linear(hid_d//2, 1)) for _ in range(n_exp)])\n        self.gate = nn.Sequential(nn.Linear(in_d, hid_d), nn.LayerNorm(hid_d), nn.SELU(), \n                                  nn.Linear(hid_d, n_exp), nn.Softmax(dim=1))\n    def forward(self, x, c):\n        z = self.enc(x)\n        preds = torch.cat([E(z) for E in self.exps], dim=1)\n        g = self.gate(c)\n        return torch.sum(preds * g, dim=1, keepdim=True), g\n\nclass CustomModel(Model):\n    def __init__(self):\n        super().__init__()\n        self.fitted = False\n        self.dev = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n        self.model = MSCMoE().to(self.dev)\n        self.opt = torch.optim.Adam(self.model.parameters(), lr=1e-3, weight_decay=1e-4)\n\n    def prep_data(self, df):\n        hl = \"label\" in df\n        if hl: df = df[~df[\"label\"].isna().any(axis=1)].copy()\n        ft = df[\"feature\"].fillna(0.0)\n        c = ft.groupby(\"datetime\").mean().loc[ft.index.get_level_values(\"datetime\")].values\n        t_ft = torch.tensor(ft.values, dtype=torch.float32)\n        t_c = torch.tensor(c, dtype=torch.float32)\n        if hl: \n            return t_ft, t_c, torch.tensor(df[\"label\"].values.ravel(), dtype=torch.float32)\n        return t_ft, t_c, None\n\n    def fit(self, dataset: DatasetH):\n        t_x, t_c, t_y = self.prep_data(dataset.prepare(\"train\", col_set=[\"feature\", \"label\"], data_key=DataHandlerLP.DK_L))\n        v_x, v_c, v_y = self.prep_data(dataset.prepare(\"valid\", col_set=[\"feature\", \"label\"], data_key=DataHandlerLP.DK_L))\n        ldr = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(t_x, t_c, t_y), batch_size=4096, shuffle=True)\n        bst_l, pat, tr = float('inf'), 5, 0\n        bst_w = copy.deepcopy(self.model.state_dict())\n        for _ in range(40):\n            self.model.train()\n            for x, c, y in ldr:\n                self.opt.zero_grad()\n                p, g = self.model(x.to(self.dev), c.to(self.dev))\n                mg = g.mean(dim=0)\n                loss = F.mse_loss(p.squeeze(), y.to(self.dev)) + F.mse_loss(mg, torch.ones_like(mg)/4.0)\n                loss.backward()\n                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 5.0)\n                self.opt.step()\n            self.model.eval()\n            with torch.no_grad():\n                pv, _ = self.model(v_x.to(self.dev), v_c.to(self.dev))\n                vl = F.mse_loss(pv.squeeze(), v_y.to(self.dev)).item()\n            if vl < bst_l: \n                bst_l, tr, bst_w = copy.deepcopy(vl), 0, copy.deepcopy(self.model.state_dict())\n            else: \n                tr += 1\n                if tr >= pat: break\n        self.model.load_state_dict(bst_w)\n        self.fitted = True\n\n    def predict(self, dataset: DatasetH, segment=\"test\"):\n        df = dataset.prepare(segment, col_set=[\"feature\", \"label\"], data_key=DataHandlerLP.DK_I)\n        x, c, _ = self.prep_data(df)\n        ldr = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(x, c), batch_size=4096)\n        self.model.eval()\n        ps = []\n        with torch.no_grad():\n            for bx, bc in ldr:\n                p, _ = self.model(bx.to(self.dev), bc.to(self.dev))\n                ps.append(p.cpu().numpy().ravel())\n        return pd.Series(np.concatenate(ps), index=df[\"feature\"].index, name=\"score\")"
7}

1{
2  "start_line": 41,
3  "filename": "qlib/custom_model.py",
4  "content": "    def prep_data(self, df, is_train=False):\n        hl = \"label\" in df\n        if hl and is_train: df = df[~df[\"label\"].isna().any(axis=1)].copy()\n        ft = df[\"feature\"].fillna(0.0)\n        c = ft.groupby(\"datetime\").mean().loc[ft.index.get_level_values(\"datetime\")].values\n        t_ft = torch.tensor(ft.values, dtype=torch.float32)\n        t_c = torch.tensor(c, dtype=torch.float32)\n        if hl and is_train: \n            return t_ft, t_c, torch.tensor(df[\"label\"].values.ravel(), dtype=torch.float32)\n        return t_ft, t_c, None\n\n    def fit(self, dataset: DatasetH):\n        t_x, t_c, t_y = self.prep_data(dataset.prepare(\"train\", col_set=[\"feature\", \"label\"], data_key=DataHandlerLP.DK_L), True)\n        v_x, v_c, v_y = self.prep_data(dataset.prepare(\"valid\", col_set=[\"feature\", \"label\"], data_key=DataHandlerLP.DK_L), True)",
5  "op": "replace",
6  "end_line": 54
7}

1{
2  "n": 2
3}

1{
2  "filename": "custom_model.py",
3  "start_line": 16,
4  "op": "replace",
5  "end_line": 103,
6  "after_line": 15,
7  "content": "# =====================================================================\n# EDITABLE: CustomModel — implement your stock prediction model here\n# =====================================================================\nfrom torch.utils.data import DataLoader, TensorDataset\n\nclass ResBlock(nn.Module):\n    def __init__(self, dim, dropout_rate=0.2):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(dim, dim),\n            nn.LayerNorm(dim),\n            nn.GELU(),\n            nn.Dropout(dropout_rate),\n            nn.Linear(dim, dim),\n            nn.LayerNorm(dim)\n        )\n        self.act = nn.GELU()\n    def forward(self, x):\n        return self.act(x + self.net(x))\n\nclass RobustMLP(nn.Module):\n    def __init__(self, input_dim=158, hidden_dim=256):\n        super().__init__()\n        self.proj = nn.Sequential(\n            nn.Linear(input_dim, hidden_dim),\n            nn.LayerNorm(hidden_dim),\n            nn.GELU()\n        )\n        self.res1 = ResBlock(hidden_dim)\n        self.res2 = ResBlock(hidden_dim)\n        self.out = nn.Linear(hidden_dim, 1)\n        \n    def forward(self, x):\n        x = self.proj(x)\n        x = self.res1(x)\n        x = self.res2(x)\n        return self.out(x).squeeze(-1)\n\nclass CustomModel(Model):\n    \"\"\"\n    Concept Drift Adaptation Model using Temporal Variance Risk Extrapolation (V-REx).\n    \n    Algorithmic Contribution:\n    To combat temporal concept drift in financial series, we formulate a novel objective\n    that acts structurally beyond minimizing average loss. By treating historical time periods (years)\n    as distinct temporal domains, we minimize a combination of the global average loss \n    AND the variance of the losses across these temporal domains.\n    This limits the model's capacity to overfit to transient correlations of specific volatile years\n    and forces it to rely ONLY on regime-invariant stable features, effectively mitigating drift.\n    \"\"\"\n    def __init__(self):\n        super().__init__()\n        self.model = RobustMLP().to(DEVICE)\n        self.beta = 10.0  # V-REx variance penalty coefficient\n        self.epochs = 20\n        self.batch_size = 8192\n        self.lr = 5e-4\n        self.fitted = False\n\n    def fit(self, dataset: DatasetH):\n        df_train = dataset.prepare(\n            \"train\", col_set=[\"feature\", \"label\"], data_key=DataHandlerLP.DK_L\n        )\n        features = df_train[\"feature\"].values\n        labels = df_train[\"label\"].values.ravel()\n        \n        times = df_train.index.get_level_values(\"datetime\")\n        domains = times.year.values.astype(int)\n        \n        mask = ~(np.isnan(features).any(axis=1) | np.isnan(labels))\n        features = features[mask]\n        labels = labels[mask]\n        domains = domains[mask]\n        \n        x_tensor = torch.tensor(features, dtype=torch.float32)\n        y_tensor = torch.tensor(labels, dtype=torch.float32)\n        d_tensor = torch.tensor(domains, dtype=torch.long)\n        \n        train_ds = TensorDataset(x_tensor, y_tensor, d_tensor)\n        train_loader = DataLoader(train_ds, batch_size=self.batch_size, shuffle=True)\n        \n        optimizer = torch.optim.AdamW(self.model.parameters(), lr=self.lr, weight_decay=1e-3)\n        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=self.epochs)\n        \n        self.model.train()\n        for epoch in range(self.epochs):\n            for x_batch, y_batch, d_batch in train_loader:\n                x_batch = x_batch.to(DEVICE)\n                y_batch = y_batch.to(DEVICE)\n                d_batch = d_batch.to(DEVICE)\n                \n                optimizer.zero_grad()\n                preds = self.model(x_batch)\n                \n                unique_domains = torch.unique(d_batch)\n                domain_losses = []\n                for d in unique_domains:\n                    idx = (d_batch == d)\n                    if idx.sum() > 10:\n                        loss_d = F.mse_loss(preds[idx], y_batch[idx])\n                        domain_losses.append(loss_d)\n                \n                if len(domain_losses) > 1:\n                    losses_stack = torch.stack(domain_losses)\n                    mean_loss = losses_stack.mean()\n                    var_loss = losses_stack.var()\n                    loss = mean_loss + self.beta * var_loss\n                else:\n                    loss = F.mse_loss(preds, y_batch)\n                    \n                loss.backward()\n                torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)\n                optimizer.step()\n                \n            scheduler.step()\n        self.fitted = True\n\n    def predict(self, dataset: DatasetH, segment=\"test\"):\n        if not self.fitted:\n            raise ValueError(\"Model is not fitted yet!\")\n            \n        df_test = dataset.prepare(\n            segment, col_set=\"feature\", data_key=DataHandlerLP.DK_I\n        )\n        if hasattr(df_test, \"columns\") and isinstance(df_test.columns, pd.MultiIndex) and \"feature\" in df_test.columns.levels[0]:\n            features = df_test[\"feature\"]\n        elif hasattr(df_test, \"columns\") and \"feature\" in df_test.columns:\n            features = df_test[\"feature\"]\n        else:\n            features = df_test\n            \n        index = features.index\n        features_np = np.nan_to_num(features.values, nan=0.0)\n        \n        self.model.eval()\n        with torch.no_grad():\n            x_tensor = torch.tensor(features_np, dtype=torch.float32)\n            test_ds = TensorDataset(x_tensor)\n            test_loader = DataLoader(test_ds, batch_size=8192, shuffle=False)\n            preds_list = []\n            for (x_batch,) in test_loader:\n                x_batch = x_batch.to(DEVICE)\n                preds = self.model(x_batch)\n                preds_list.append(preds.cpu().numpy())\n            \n        preds_all = np.concatenate(preds_list)\n        return pd.Series(preds_all, index=index, name=\"score\")"
8}

1{
2  "op": "replace",
3  "filename": "qlib/custom_model.py",
4  "start_line": 16,
5  "content": "# =====================================================================\n# EDITABLE: CustomModel — implement your stock prediction model here\n# =====================================================================\nfrom torch.utils.data import DataLoader, TensorDataset\n\nclass ResBlock(nn.Module):\n    def __init__(self, dim, dropout_rate=0.2):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(dim, dim),\n            nn.LayerNorm(dim),\n            nn.GELU(),\n            nn.Dropout(dropout_rate),\n            nn.Linear(dim, dim),\n            nn.LayerNorm(dim)\n        )\n        self.act = nn.GELU()\n    def forward(self, x):\n        return self.act(x + self.net(x))\n\nclass RobustMLP(nn.Module):\n    def __init__(self, input_dim=158, hidden_dim=256):\n        super().__init__()\n        self.proj = nn.Sequential(\n            nn.Linear(input_dim, hidden_dim),\n            nn.LayerNorm(hidden_dim),\n            nn.GELU()\n        )\n        self.res1 = ResBlock(hidden_dim)\n        self.res2 = ResBlock(hidden_dim)\n        self.out = nn.Linear(hidden_dim, 1)\n        \n    def forward(self, x):\n        x = self.proj(x)\n        x = self.res1(x)\n        x = self.res2(x)\n        return self.out(x).squeeze(-1)\n\nclass CustomModel(Model):\n    \"\"\"\n    Concept Drift Adaptation Model using Temporal Variance Risk Extrapolation (V-REx).\n    \n    Algorithmic Contribution:\n    To combat temporal concept drift in financial series, we formulate a novel objective\n    that acts structurally beyond minimizing average loss. By treating historical time periods (years)\n    as distinct temporal domains, we minimize a combination of the global average loss \n    AND the variance of the losses across these temporal domains.\n    This limits the model's capacity to overfit to transient correlations of specific volatile years\n    and forces it to rely ONLY on regime-invariant stable features, effectively mitigating drift.\n    \"\"\"\n    def __init__(self):\n        super().__init__()\n        self.model = RobustMLP().to(DEVICE)\n        self.beta = 10.0  # V-REx variance penalty coefficient\n        self.epochs = 20\n        self.batch_size = 8192\n        self.lr = 5e-4\n        self.fitted = False\n\n    def fit(self, dataset: DatasetH):\n        df_train = dataset.prepare(\n            \"train\", col_set=[\"feature\", \"label\"], data_key=DataHandlerLP.DK_L\n        )\n        features = df_train[\"feature\"].values\n        labels = df_train[\"label\"].values.ravel()\n        \n        times = df_train.index.get_level_values(\"datetime\")\n        domains = times.year.values.astype(int)\n        \n        mask = ~(np.isnan(features).any(axis=1) | np.isnan(labels))\n        features = features[mask]\n        labels = labels[mask]\n        domains = domains[mask]\n        \n        x_tensor = torch.tensor(features, dtype=torch.float32)\n        y_tensor = torch.tensor(labels, dtype=torch.float32)\n        d_tensor = torch.tensor(domains, dtype=torch.long)\n        \n        train_ds = TensorDataset(x_tensor, y_tensor, d_tensor)\n        train_loader = DataLoader(train_ds, batch_size=self.batch_size, shuffle=True)\n        \n        optimizer = torch.optim.AdamW(self.model.parameters(), lr=self.lr, weight_decay=1e-3)\n        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=self.epochs)\n        \n        self.model.train()\n        for epoch in range(self.epochs):\n            for x_batch, y_batch, d_batch in train_loader:\n                x_batch = x_batch.to(DEVICE)\n                y_batch = y_batch.to(DEVICE)\n                d_batch = d_batch.to(DEVICE)\n                \n                optimizer.zero_grad()\n                preds = self.model(x_batch)\n                \n                unique_domains = torch.unique(d_batch)\n                domain_losses = []\n                for d in unique_domains:\n                    idx = (d_batch == d)\n                    if idx.sum() > 10:\n                        loss_d = F.mse_loss(preds[idx], y_batch[idx])\n                        domain_losses.append(loss_d)\n                \n                if len(domain_losses) > 1:\n                    losses_stack = torch.stack(domain_losses)\n                    mean_loss = losses_stack.mean()\n                    var_loss = losses_stack.var()\n                    loss = mean_loss + self.beta * var_loss\n                else:\n                    loss = F.mse_loss(preds, y_batch)\n                    \n                loss.backward()\n                torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)\n                optimizer.step()\n                \n            scheduler.step()\n        self.fitted = True\n\n    def predict(self, dataset: DatasetH, segment=\"test\"):\n        if not self.fitted:\n            raise ValueError(\"Model is not fitted yet!\")\n            \n        df_test = dataset.prepare(\n            segment, col_set=\"feature\", data_key=DataHandlerLP.DK_I\n        )\n        if hasattr(df_test, \"columns\") and isinstance(df_test.columns, pd.MultiIndex) and \"feature\" in df_test.columns.levels[0]:\n            features = df_test[\"feature\"]\n        elif hasattr(df_test, \"columns\") and \"feature\" in df_test.columns:\n            features = df_test[\"feature\"]\n        else:\n            features = df_test\n            \n        index = features.index\n        features_np = np.nan_to_num(features.values, nan=0.0)\n        \n        self.model.eval()\n        with torch.no_grad():\n            x_tensor = torch.tensor(features_np, dtype=torch.float32)\n            test_ds = TensorDataset(x_tensor)\n            test_loader = DataLoader(test_ds, batch_size=8192, shuffle=False)\n            preds_list = []\n            for (x_batch,) in test_loader:\n                x_batch = x_batch.to(DEVICE)\n                preds = self.model(x_batch)\n                preds_list.append(preds.cpu().numpy())\n            \n        preds_all = np.concatenate(preds_list)\n        return pd.Series(preds_all, index=index, name=\"score\")",
6  "after_line": 15,
7  "end_line": 103
8}

1{
2  "start_line": 16,
3  "content": "# =====================================================================\n# EDITABLE: CustomModel — implement your stock prediction model here\n# =====================================================================\nimport numpy as np\nimport pandas as pd\nfrom qlib.model.base import Model\nfrom qlib.data.dataset.handler import DataHandlerLP\nfrom qlib.data.dataset import DatasetH\nimport lightgbm as lgb\n\nclass CustomModel(Model):\n    \"\"\"\n    Concept Drift Adaptation via Distribution-Matched Meta-Ensemble.\n    \n    Algorithmic Contribution:\n    We propose an unsupervised test-time temporal domain routing algorithm.\n    During training, we partition the data sequentially into distinct temporal regimes \n    and train an ensemble of expert models, storing their regime representation \n    (the centroid of cross-sectional distributions). We also train a universal global model.\n    During inference, without requiring labels, we dynamically assess the macroscopic \n    market condition of the current target period and route predictions by mapping \n    the target state back to the most similar historical regimes using softmax over \n    Euclidean distances. This combats temporal shift continuously.\n    \"\"\"\n    def __init__(self):\n        super().__init__()\n        self.chunk_models = []\n        self.global_model = None\n        self.regime_centroids = []\n        self.n_regimes = 6\n        self.fitted = False\n        self.gamma = 3.0\n        \n        self.params = {\n            \"objective\": \"mse\",\n            \"learning_rate\": 0.05,\n            \"max_depth\": 8,\n            \"colsample_bytree\": 0.8879,\n            \"subsample\": 0.8789,\n            \"lambda_l1\": 200.0,\n            \"lambda_l2\": 500.0,\n            \"num_leaves\": 210,\n            \"num_threads\": 20,\n            \"verbosity\": -1,\n        }\n\n    def fit(self, dataset: DatasetH):\n        df_train = dataset.prepare(\n            \"train\", col_set=[\"feature\", \"label\"], data_key=DataHandlerLP.DK_L\n        )\n        features = df_train[\"feature\"].values\n        labels = df_train[\"label\"].values.ravel()\n        \n        times = df_train.index.get_level_values(\"datetime\").unique().sort_values()\n        \n        # Train Global Model\n        mask = ~(np.isnan(features).any(axis=1) | np.isnan(labels))\n        lgb_global = lgb.Dataset(features[mask], labels[mask])\n        self.global_model = lgb.train(self.params, lgb_global, num_boost_round=300)\n        \n        # Split into consecutive regimes\n        chunk_size = len(times) // self.n_regimes\n        \n        for k in range(self.n_regimes):\n            start_t = times[k * chunk_size]\n            if k == self.n_regimes - 1:\n                end_t = times[-1]\n            else:\n                end_t = times[(k + 1) * chunk_size - 1]\n                \n            chunk_df = df_train.loc[pd.IndexSlice[start_t:end_t, :], :]\n            \n            X_chunk = chunk_df[\"feature\"].values\n            Y_chunk = chunk_df[\"label\"].values.ravel()\n            \n            mask_chunk = ~(np.isnan(X_chunk).any(axis=1) | np.isnan(Y_chunk))\n            X_chunk = X_chunk[mask_chunk]\n            Y_chunk = Y_chunk[mask_chunk]\n            \n            centroid = np.median(np.nan_to_num(X_chunk, nan=0.0), axis=0)\n            self.regime_centroids.append(centroid)\n            \n            lgb_chunk = lgb.Dataset(X_chunk, Y_chunk)\n            model = lgb.train(self.params, lgb_chunk, num_boost_round=150)\n            self.chunk_models.append(model)\n            \n        self.regime_centroids = np.array(self.regime_centroids)\n        self.fitted = True\n\n    def predict(self, dataset: DatasetH, segment=\"test\"):\n        if not self.fitted:\n            raise ValueError(\"Model is not fitted yet!\")\n            \n        df_test = dataset.prepare(\n            segment, col_set=\"feature\", data_key=DataHandlerLP.DK_I\n        )\n        if hasattr(df_test, \"columns\") and isinstance(df_test.columns, pd.MultiIndex) and \"feature\" in df_test.columns.levels[0]:\n            features = df_test[\"feature\"]\n        elif hasattr(df_test, \"columns\") and \"feature\" in df_test.columns:\n            features = df_test[\"feature\"]\n        else:\n            features = df_test\n            \n        predictions = pd.Series(0.0, index=features.index, name=\"score\")\n        times = features.index.get_level_values(\"datetime\")\n        year_months = times.year.values * 100 + times.month.values\n        \n        for ym in np.unique(year_months):\n            idx_mask = (year_months == ym)\n            X_val = np.nan_to_num(features.values[idx_mask], nan=0.0)\n            \n            if len(X_val) == 0:\n                continue\n                \n            current_centroid = np.median(X_val, axis=0)\n            dists = np.linalg.norm(self.regime_centroids - current_centroid, axis=1)\n            \n            d_min = np.min(dists)\n            d_std = np.std(dists) + 1e-8\n            dists_z = (dists - d_min) / d_std\n            weights = np.exp(-self.gamma * dists_z)\n            weights /= np.sum(weights)\n            \n            pred_ym_expert = np.zeros(len(X_val))\n            for k, model in enumerate(self.chunk_models):\n                pred_ym_expert += weights[k] * model.predict(X_val)\n                \n            pred_ym_global = self.global_model.predict(X_val)\n            \n            predictions.iloc[idx_mask] = 0.5 * pred_ym_global + 0.5 * pred_ym_expert\n            \n        return predictions",
4  "end_line": 162,
5  "filename": "qlib/custom_model.py",
6  "after_line": 15,
7  "op": "replace"
8}