{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.10.12"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":84493,"databundleVersionId":9871156,"sourceType":"competition"},{"sourceId":10376360,"sourceType":"datasetVersion","datasetId":6427477},{"sourceId":10376415,"sourceType":"datasetVersion","datasetId":6427521},{"sourceId":201379473,"sourceType":"kernelVersion"},{"sourceId":201788317,"sourceType":"kernelVersion"},{"sourceId":204479873,"sourceType":"kernelVersion"}],"dockerImageVersionId":30823,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":181.780351,"end_time":"2024-12-27T02:35:55.804678","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2024-12-27T02:32:54.024327","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"#!cp -rf /kaggle/input/test04/* .\n#!cp -rf /kaggle/input/test06/* .\n#!cp -rf /kaggle/input/checkpoint1/* .\n#!mv axial_positional_embedding-0.2.1 axial_positional_embedding-0.2.1.tar.gz\n!ls -l","metadata":{"papermill":{"duration":20.748837,"end_time":"2024-12-27T02:33:16.895821","exception":false,"start_time":"2024-12-27T02:32:56.146984","status":"completed"},"tags":[],"trusted":true,"execution":{"iopub.status.busy":"2025-01-09T02:15:41.260319Z","iopub.execute_input":"2025-01-09T02:15:41.260640Z","iopub.status.idle":"2025-01-09T02:15:41.378994Z","shell.execute_reply.started":"2025-01-09T02:15:41.260603Z","shell.execute_reply":"2025-01-09T02:15:41.378109Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\n\nclass TriangularCausalMask():\n    def __init__(self, B, L, device=\"cpu\"):\n        mask_shape = [B, 1, L, L]\n        with torch.no_grad():\n            self._mask = torch.triu(torch.ones(mask_shape, dtype=torch.bool), diagonal=1).to(device)\n\n    @property\n    def mask(self):\n        return self._mask\n\n\nclass ProbMask():\n    def __init__(self, B, H, L, index, scores, device=\"cpu\"):\n        _mask = torch.ones(L, scores.shape[-1], dtype=torch.bool).to(device).triu(1)\n        _mask_ex = _mask[None, None, :].expand(B, H, L, scores.shape[-1])\n        indicator = _mask_ex[torch.arange(B)[:, None, None],\n                    torch.arange(H)[None, :, None],\n                    index, :].to(device)\n        self._mask = indicator.view(scores.shape).to(device)\n\n    @property\n    def mask(self):\n        return self._mask\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-09T02:15:41.380330Z","iopub.execute_input":"2025-01-09T02:15:41.380633Z","iopub.status.idle":"2025-01-09T02:15:42.857240Z","shell.execute_reply.started":"2025-01-09T02:15:41.380607Z","shell.execute_reply":"2025-01-09T02:15:42.856590Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport math\n\n\nclass PositionalEmbedding(nn.Module):\n    def __init__(self, d_model, max_len=5000):\n        super(PositionalEmbedding, self).__init__()\n        # Compute the positional encodings once in log space.\n        pe = torch.zeros(max_len, d_model).float()\n        pe.require_grad = False\n\n        position = torch.arange(0, max_len).float().unsqueeze(1)\n        div_term = (torch.arange(0, d_model, 2).float()\n                    * -(math.log(10000.0) / d_model)).exp()\n\n        pe[:, 0::2] = torch.sin(position * div_term)\n        pe[:, 1::2] = torch.cos(position * div_term)\n\n        pe = pe.unsqueeze(0)\n        self.register_buffer('pe', pe)\n\n    def forward(self, x):\n        return self.pe[:, :x.size(1)]\n\n\nclass TokenEmbedding(nn.Module):\n    def __init__(self, c_in, d_model):\n        super(TokenEmbedding, self).__init__()\n        padding = 1 if torch.__version__ >= '1.5.0' else 2\n        self.tokenConv = nn.Conv1d(in_channels=c_in, out_channels=d_model,\n                                   kernel_size=3, padding=padding, padding_mode='circular', bias=False)\n        for m in self.modules():\n            if isinstance(m, nn.Conv1d):\n                nn.init.kaiming_normal_(\n                    m.weight, mode='fan_in', nonlinearity='leaky_relu')\n\n    def forward(self, x):\n        x = self.tokenConv(x.permute(0, 2, 1)).transpose(1, 2)\n        return x\n\n\nclass FixedEmbedding(nn.Module):\n    def __init__(self, c_in, d_model):\n        super(FixedEmbedding, self).__init__()\n\n        w = torch.zeros(c_in, d_model).float()\n        w.require_grad = False\n\n        position = torch.arange(0, c_in).float().unsqueeze(1)\n        div_term = (torch.arange(0, d_model, 2).float()\n                    * -(math.log(10000.0) / d_model)).exp()\n\n        w[:, 0::2] = torch.sin(position * div_term)\n        w[:, 1::2] = torch.cos(position * div_term)\n\n        self.emb = nn.Embedding(c_in, d_model)\n        self.emb.weight = nn.Parameter(w, requires_grad=False)\n\n    def forward(self, x):\n        return self.emb(x).detach()\n\n\nclass TemporalEmbedding(nn.Module):\n    def __init__(self, d_model, embed_type='fixed', freq='h'):\n        super(TemporalEmbedding, self).__init__()\n\n        minute_size = 4\n        hour_size = 24\n        weekday_size = 7\n        day_size = 32\n        month_size = 13\n\n        Embed = FixedEmbedding if embed_type == 'fixed' else nn.Embedding\n        if freq == 't':\n            self.minute_embed = Embed(minute_size, d_model)\n        self.hour_embed = Embed(hour_size, d_model)\n        self.weekday_embed = Embed(weekday_size, d_model)\n        self.day_embed = Embed(day_size, d_model)\n        self.month_embed = Embed(month_size, d_model)\n\n    def forward(self, x):\n        x = x.long()\n        minute_x = self.minute_embed(x[:, :, 4]) if hasattr(\n            self, 'minute_embed') else 0.\n        hour_x = self.hour_embed(x[:, :, 3])\n        weekday_x = self.weekday_embed(x[:, :, 2])\n        day_x = self.day_embed(x[:, :, 1])\n        month_x = self.month_embed(x[:, :, 0])\n\n        return hour_x + weekday_x + day_x + month_x + minute_x\n\n\nclass TimeFeatureEmbedding(nn.Module):\n    def __init__(self, d_model, embed_type='timeF', freq='h'):\n        super(TimeFeatureEmbedding, self).__init__()\n\n        freq_map = {'h': 4, 't': 5, 's': 6,\n                    'm': 1, 'a': 1, 'w': 2, 'd': 3, 'b': 3}\n        d_inp = freq_map[freq]\n        self.embed = nn.Linear(d_inp, d_model, bias=False)\n\n    def forward(self, x):\n        return self.embed(x)\n\n\nclass DataEmbedding(nn.Module):\n    def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):\n        super(DataEmbedding, self).__init__()\n\n        self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)\n        self.position_embedding = PositionalEmbedding(d_model=d_model)\n        self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type,\n                                                    freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding(\n            d_model=d_model, embed_type=embed_type, freq=freq)\n        self.dropout = nn.Dropout(p=dropout)\n\n    def forward(self, x, x_mark):\n        if x_mark is None:\n            x = self.value_embedding(x) + self.position_embedding(x)\n        else:\n            x = self.value_embedding(\n                x) + self.temporal_embedding(x_mark) + self.position_embedding(x)\n        return self.dropout(x)\n\n\nclass DataEmbedding_inverted(nn.Module):\n    def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):\n        super(DataEmbedding_inverted, self).__init__()\n        self.value_embedding = nn.Linear(c_in, d_model)\n        self.dropout = nn.Dropout(p=dropout)\n\n    def forward(self, x, x_mark):\n        x = x.permute(0, 2, 1)\n        # x: [Batch Variate Time]\n        if x_mark is None:\n            x = self.value_embedding(x)\n        else:\n            # the potential to take covariates (e.g. timestamps) as tokens\n            x = self.value_embedding(torch.cat([x, x_mark.permute(0, 2, 1)], 1)) \n        # x: [Batch Variate d_model]\n        return self.dropout(x)\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-09T02:15:42.858940Z","iopub.execute_input":"2025-01-09T02:15:42.859313Z","iopub.status.idle":"2025-01-09T02:15:42.877157Z","shell.execute_reply.started":"2025-01-09T02:15:42.859290Z","shell.execute_reply":"2025-01-09T02:15:42.876326Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport numpy as np\nfrom math import sqrt\n#from utils.masking import TriangularCausalMask, ProbMask\nfrom reformer_pytorch import LSHSelfAttention\nfrom einops import rearrange\n\n\n# Code implementation from https://github.com/thuml/Flowformer\nclass FlowAttention(nn.Module):\n    def __init__(self, attention_dropout=0.1):\n        super(FlowAttention, self).__init__()\n        self.dropout = nn.Dropout(attention_dropout)\n\n    def kernel_method(self, x):\n        return torch.sigmoid(x)\n\n    def forward(self, queries, keys, values, attn_mask, tau=None, delta=None):\n        queries = queries.transpose(1, 2)\n        keys = keys.transpose(1, 2)\n        values = values.transpose(1, 2)\n        # kernel\n        queries = self.kernel_method(queries)\n        keys = self.kernel_method(keys)\n        # incoming and outgoing\n        normalizer_row = 1.0 / (torch.einsum(\"nhld,nhd->nhl\", queries + 1e-6, keys.sum(dim=2) + 1e-6))\n        normalizer_col = 1.0 / (torch.einsum(\"nhsd,nhd->nhs\", keys + 1e-6, queries.sum(dim=2) + 1e-6))\n        # reweighting\n        normalizer_row_refine = (\n            torch.einsum(\"nhld,nhd->nhl\", queries + 1e-6, (keys * normalizer_col[:, :, :, None]).sum(dim=2) + 1e-6))\n        normalizer_col_refine = (\n            torch.einsum(\"nhsd,nhd->nhs\", keys + 1e-6, (queries * normalizer_row[:, :, :, None]).sum(dim=2) + 1e-6))\n        # competition and allocation\n        normalizer_row_refine = torch.sigmoid(\n            normalizer_row_refine * (float(queries.shape[2]) / float(keys.shape[2])))\n        normalizer_col_refine = torch.softmax(normalizer_col_refine, dim=-1) * keys.shape[2]  # B h L vis\n        # multiply\n        kv = keys.transpose(-2, -1) @ (values * normalizer_col_refine[:, :, :, None])\n        x = (((queries @ kv) * normalizer_row[:, :, :, None]) * normalizer_row_refine[:, :, :, None]).transpose(1,\n                                                                                                                2).contiguous()\n        return x, None\n\n\n# Code implementation from https://github.com/shreyansh26/FlashAttention-PyTorch\nclass FlashAttention(nn.Module):\n    def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):\n        super(FlashAttention, self).__init__()\n        self.scale = scale\n        self.mask_flag = mask_flag\n        self.output_attention = output_attention\n        self.dropout = nn.Dropout(attention_dropout)\n\n    def flash_attention_forward(self, Q, K, V, mask=None):\n        BLOCK_SIZE = 32\n        NEG_INF = -1e10  # -infinity\n        EPSILON = 1e-10\n        # mask = torch.randint(0, 2, (128, 8)).to(device='cuda')\n        O = torch.zeros_like(Q, requires_grad=True)\n        l = torch.zeros(Q.shape[:-1])[..., None]\n        m = torch.ones(Q.shape[:-1])[..., None] * NEG_INF\n\n        O = O.to(device='cuda')\n        l = l.to(device='cuda')\n        m = m.to(device='cuda')\n\n        Q_BLOCK_SIZE = min(BLOCK_SIZE, Q.shape[-1])\n        KV_BLOCK_SIZE = BLOCK_SIZE\n\n        Q_BLOCKS = torch.split(Q, Q_BLOCK_SIZE, dim=2)\n        K_BLOCKS = torch.split(K, KV_BLOCK_SIZE, dim=2)\n        V_BLOCKS = torch.split(V, KV_BLOCK_SIZE, dim=2)\n        if mask is not None:\n            mask_BLOCKS = list(torch.split(mask, KV_BLOCK_SIZE, dim=1))\n\n        Tr = len(Q_BLOCKS)\n        Tc = len(K_BLOCKS)\n\n        O_BLOCKS = list(torch.split(O, Q_BLOCK_SIZE, dim=2))\n        l_BLOCKS = list(torch.split(l, Q_BLOCK_SIZE, dim=2))\n        m_BLOCKS = list(torch.split(m, Q_BLOCK_SIZE, dim=2))\n\n        for j in range(Tc):\n            Kj = K_BLOCKS[j]\n            Vj = V_BLOCKS[j]\n            if mask is not None:\n                maskj = mask_BLOCKS[j]\n\n            for i in range(Tr):\n                Qi = Q_BLOCKS[i]\n                Oi = O_BLOCKS[i]\n                li = l_BLOCKS[i]\n                mi = m_BLOCKS[i]\n\n                scale = 1 / np.sqrt(Q.shape[-1])\n                Qi_scaled = Qi * scale\n\n                S_ij = torch.einsum('... i d, ... j d -> ... i j', Qi_scaled, Kj)\n                if mask is not None:\n                    # Masking\n                    maskj_temp = rearrange(maskj, 'b j -> b 1 1 j')\n                    S_ij = torch.where(maskj_temp > 0, S_ij, NEG_INF)\n\n                m_block_ij, _ = torch.max(S_ij, dim=-1, keepdims=True)\n                P_ij = torch.exp(S_ij - m_block_ij)\n                if mask is not None:\n                    # Masking\n                    P_ij = torch.where(maskj_temp > 0, P_ij, 0.)\n\n                l_block_ij = torch.sum(P_ij, dim=-1, keepdims=True) + EPSILON\n\n                P_ij_Vj = torch.einsum('... i j, ... j d -> ... i d', P_ij, Vj)\n\n                mi_new = torch.maximum(m_block_ij, mi)\n                li_new = torch.exp(mi - mi_new) * li + torch.exp(m_block_ij - mi_new) * l_block_ij\n\n                O_BLOCKS[i] = (li / li_new) * torch.exp(mi - mi_new) * Oi + (\n                        torch.exp(m_block_ij - mi_new) / li_new) * P_ij_Vj\n                l_BLOCKS[i] = li_new\n                m_BLOCKS[i] = mi_new\n\n        O = torch.cat(O_BLOCKS, dim=2)\n        l = torch.cat(l_BLOCKS, dim=2)\n        m = torch.cat(m_BLOCKS, dim=2)\n        return O, l, m\n\n    def forward(self, queries, keys, values, attn_mask, tau=None, delta=None):\n        res = \\\n        self.flash_attention_forward(queries.permute(0, 2, 1, 3), keys.permute(0, 2, 1, 3), values.permute(0, 2, 1, 3),\n                                     attn_mask)[0]\n        return res.permute(0, 2, 1, 3).contiguous(), None\n\n\nclass FullAttention(nn.Module):\n    def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):\n        super(FullAttention, self).__init__()\n        self.scale = scale\n        self.mask_flag = mask_flag\n        self.output_attention = output_attention\n        self.dropout = nn.Dropout(attention_dropout)\n\n    def forward(self, queries, keys, values, attn_mask, tau=None, delta=None):\n        B, L, H, E = queries.shape\n        _, S, _, D = values.shape\n        scale = self.scale or 1. / sqrt(E)\n\n        scores = torch.einsum(\"blhe,bshe->bhls\", queries, keys)\n\n        if self.mask_flag:\n            if attn_mask is None:\n                attn_mask = TriangularCausalMask(B, L, device=queries.device)\n\n            scores.masked_fill_(attn_mask.mask, -np.inf)\n\n        A = self.dropout(torch.softmax(scale * scores, dim=-1))\n        V = torch.einsum(\"bhls,bshd->blhd\", A, values)\n\n        if self.output_attention:\n            return (V.contiguous(), A)\n        else:\n            return (V.contiguous(), None)\n\n\n# Code implementation from https://github.com/zhouhaoyi/Informer2020\nclass ProbAttention(nn.Module):\n    def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):\n        super(ProbAttention, self).__init__()\n        self.factor = factor\n        self.scale = scale\n        self.mask_flag = mask_flag\n        self.output_attention = output_attention\n        self.dropout = nn.Dropout(attention_dropout)\n\n    def _prob_QK(self, Q, K, sample_k, n_top):  # n_top: c*ln(L_q)\n        # Q [B, H, L, D]\n        B, H, L_K, E = K.shape\n        _, _, L_Q, _ = Q.shape\n\n        # calculate the sampled Q_K\n        K_expand = K.unsqueeze(-3).expand(B, H, L_Q, L_K, E)\n        # real U = U_part(factor*ln(L_k))*L_q\n        index_sample = torch.randint(L_K, (L_Q, sample_k))\n        K_sample = K_expand[:, :, torch.arange(\n            L_Q).unsqueeze(1), index_sample, :]\n        Q_K_sample = torch.matmul(\n            Q.unsqueeze(-2), K_sample.transpose(-2, -1)).squeeze()\n\n        # find the Top_k query with sparisty measurement\n        M = Q_K_sample.max(-1)[0] - torch.div(Q_K_sample.sum(-1), L_K)\n        M_top = M.topk(n_top, sorted=False)[1]\n\n        # use the reduced Q to calculate Q_K\n        Q_reduce = Q[torch.arange(B)[:, None, None],\n                   torch.arange(H)[None, :, None],\n                   M_top, :]  # factor*ln(L_q)\n        Q_K = torch.matmul(Q_reduce, K.transpose(-2, -1))  # factor*ln(L_q)*L_k\n\n        return Q_K, M_top\n\n    def _get_initial_context(self, V, L_Q):\n        B, H, L_V, D = V.shape\n        if not self.mask_flag:\n            # V_sum = V.sum(dim=-2)\n            V_sum = V.mean(dim=-2)\n            contex = V_sum.unsqueeze(-2).expand(B, H,\n                                                L_Q, V_sum.shape[-1]).clone()\n        else:  # use mask\n            # requires that L_Q == L_V, i.e. for self-attention only\n            assert (L_Q == L_V)\n            contex = V.cumsum(dim=-2)\n        return contex\n\n    def _update_context(self, context_in, V, scores, index, L_Q, attn_mask):\n        B, H, L_V, D = V.shape\n\n        if self.mask_flag:\n            attn_mask = ProbMask(B, H, L_Q, index, scores, device=V.device)\n            scores.masked_fill_(attn_mask.mask, -np.inf)\n\n        attn = torch.softmax(scores, dim=-1)  # nn.Softmax(dim=-1)(scores)\n\n        context_in[torch.arange(B)[:, None, None],\n        torch.arange(H)[None, :, None],\n        index, :] = torch.matmul(attn, V).type_as(context_in)\n        if self.output_attention:\n            attns = (torch.ones([B, H, L_V, L_V]) /\n                     L_V).type_as(attn).to(attn.device)\n            attns[torch.arange(B)[:, None, None], torch.arange(H)[\n                                                  None, :, None], index, :] = attn\n            return (context_in, attns)\n        else:\n            return (context_in, None)\n\n    def forward(self, queries, keys, values, attn_mask, tau=None, delta=None):\n        B, L_Q, H, D = queries.shape\n        _, L_K, _, _ = keys.shape\n\n        queries = queries.transpose(2, 1)\n        keys = keys.transpose(2, 1)\n        values = values.transpose(2, 1)\n\n        U_part = self.factor * \\\n                 np.ceil(np.log(L_K)).astype('int').item()  # c*ln(L_k)\n        u = self.factor * \\\n            np.ceil(np.log(L_Q)).astype('int').item()  # c*ln(L_q)\n\n        U_part = U_part if U_part < L_K else L_K\n        u = u if u < L_Q else L_Q\n\n        scores_top, index = self._prob_QK(\n            queries, keys, sample_k=U_part, n_top=u)\n\n        # add scale factor\n        scale = self.scale or 1. / sqrt(D)\n        if scale is not None:\n            scores_top = scores_top * scale\n        # get the context\n        context = self._get_initial_context(values, L_Q)\n        # update the context with selected top_k queries\n        context, attn = self._update_context(\n            context, values, scores_top, index, L_Q, attn_mask)\n\n        return context.contiguous(), attn\n\n\nclass AttentionLayer(nn.Module):\n    def __init__(self, attention, d_model, n_heads, d_keys=None,\n                 d_values=None):\n        super(AttentionLayer, self).__init__()\n\n        d_keys = d_keys or (d_model // n_heads)\n        d_values = d_values or (d_model // n_heads)\n\n        self.inner_attention = attention\n        self.query_projection = nn.Linear(d_model, d_keys * n_heads)\n        self.key_projection = nn.Linear(d_model, d_keys * n_heads)\n        self.value_projection = nn.Linear(d_model, d_values * n_heads)\n        self.out_projection = nn.Linear(d_values * n_heads, d_model)\n        self.n_heads = n_heads\n\n    def forward(self, queries, keys, values, attn_mask, tau=None, delta=None):\n        B, L, _ = queries.shape\n        _, S, _ = keys.shape\n        H = self.n_heads\n\n        queries = self.query_projection(queries).view(B, L, H, -1)\n        keys = self.key_projection(keys).view(B, S, H, -1)\n        values = self.value_projection(values).view(B, S, H, -1)\n\n        out, attn = self.inner_attention(\n            queries,\n            keys,\n            values,\n            attn_mask,\n            tau=tau,\n            delta=delta\n        )\n        out = out.view(B, L, -1)\n\n        return self.out_projection(out), attn\n\n\nclass ReformerLayer(nn.Module):\n    def __init__(self, attention, d_model, n_heads, d_keys=None,\n                 d_values=None, causal=False, bucket_size=4, n_hashes=4):\n        super().__init__()\n        self.bucket_size = bucket_size\n        self.attn = LSHSelfAttention(\n            dim=d_model,\n            heads=n_heads,\n            bucket_size=bucket_size,\n            n_hashes=n_hashes,\n            causal=causal\n        )\n\n    def fit_length(self, queries):\n        # inside reformer: assert N % (bucket_size * 2) == 0\n        B, N, C = queries.shape\n        if N % (self.bucket_size * 2) == 0:\n            return queries\n        else:\n            # fill the time series\n            fill_len = (self.bucket_size * 2) - (N % (self.bucket_size * 2))\n            return torch.cat([queries, torch.zeros([B, fill_len, C]).to(queries.device)], dim=1)\n\n    def forward(self, queries, keys, values, attn_mask, tau, delta):\n        # in Reformer: defalut queries=keys\n        B, N, C = queries.shape\n        queries = self.attn(self.fit_length(queries))[:, :N, :]\n        return queries, None\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-09T02:15:42.878372Z","iopub.execute_input":"2025-01-09T02:15:42.878615Z","iopub.status.idle":"2025-01-09T02:15:43.566875Z","shell.execute_reply.started":"2025-01-09T02:15:42.878594Z","shell.execute_reply":"2025-01-09T02:15:43.565970Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch.nn as nn\nimport torch.nn.functional as F\n\n\nclass ConvLayer(nn.Module):\n    def __init__(self, c_in):\n        super(ConvLayer, self).__init__()\n        self.downConv = nn.Conv1d(in_channels=c_in,\n                                  out_channels=c_in,\n                                  kernel_size=3,\n                                  padding=2,\n                                  padding_mode='circular')\n        self.norm = nn.BatchNorm1d(c_in)\n        self.activation = nn.ELU()\n        self.maxPool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)\n\n    def forward(self, x):\n        x = self.downConv(x.permute(0, 2, 1))\n        x = self.norm(x)\n        x = self.activation(x)\n        x = self.maxPool(x)\n        x = x.transpose(1, 2)\n        return x\n\n\nclass EncoderLayer(nn.Module):\n    def __init__(self, attention, d_model, d_ff=None, dropout=0.1, activation=\"relu\"):\n        super(EncoderLayer, self).__init__()\n        d_ff = d_ff or 4 * d_model\n        self.attention = attention\n        self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)\n        self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)\n        self.norm1 = nn.LayerNorm(d_model)\n        self.norm2 = nn.LayerNorm(d_model)\n        self.dropout = nn.Dropout(dropout)\n        self.activation = F.relu if activation == \"relu\" else F.gelu\n\n    def forward(self, x, attn_mask=None, tau=None, delta=None):\n        new_x, attn = self.attention(\n            x, x, x,\n            attn_mask=attn_mask,\n            tau=tau, delta=delta\n        )\n        x = x + self.dropout(new_x)\n\n        y = x = self.norm1(x)\n        y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))\n        y = self.dropout(self.conv2(y).transpose(-1, 1))\n\n        return self.norm2(x + y), attn\n\n\nclass Encoder(nn.Module):\n    def __init__(self, attn_layers, conv_layers=None, norm_layer=None):\n        super(Encoder, self).__init__()\n        self.attn_layers = nn.ModuleList(attn_layers)\n        self.conv_layers = nn.ModuleList(conv_layers) if conv_layers is not None else None\n        self.norm = norm_layer\n\n    def forward(self, x, attn_mask=None, tau=None, delta=None):\n        # x [B, L, D]\n        attns = []\n        if self.conv_layers is not None:\n            for i, (attn_layer, conv_layer) in enumerate(zip(self.attn_layers, self.conv_layers)):\n                delta = delta if i == 0 else None\n                x, attn = attn_layer(x, attn_mask=attn_mask, tau=tau, delta=delta)\n                x = conv_layer(x)\n                attns.append(attn)\n            x, attn = self.attn_layers[-1](x, tau=tau, delta=None)\n            attns.append(attn)\n        else:\n            for attn_layer in self.attn_layers:\n                x, attn = attn_layer(x, attn_mask=attn_mask, tau=tau, delta=delta)\n                attns.append(attn)\n\n        if self.norm is not None:\n            x = self.norm(x)\n\n        return x, attns\n\n\nclass DecoderLayer(nn.Module):\n    def __init__(self, self_attention, cross_attention, d_model, d_ff=None,\n                 dropout=0.1, activation=\"relu\"):\n        super(DecoderLayer, self).__init__()\n        d_ff = d_ff or 4 * d_model\n        self.self_attention = self_attention\n        self.cross_attention = cross_attention\n        self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)\n        self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)\n        self.norm1 = nn.LayerNorm(d_model)\n        self.norm2 = nn.LayerNorm(d_model)\n        self.norm3 = nn.LayerNorm(d_model)\n        self.dropout = nn.Dropout(dropout)\n        self.activation = F.relu if activation == \"relu\" else F.gelu\n\n    def forward(self, x, cross, x_mask=None, cross_mask=None, tau=None, delta=None):\n        x = x + self.dropout(self.self_attention(\n            x, x, x,\n            attn_mask=x_mask,\n            tau=tau, delta=None\n        )[0])\n        x = self.norm1(x)\n\n        x = x + self.dropout(self.cross_attention(\n            x, cross, cross,\n            attn_mask=cross_mask,\n            tau=tau, delta=delta\n        )[0])\n\n        y = x = self.norm2(x)\n        y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))\n        y = self.dropout(self.conv2(y).transpose(-1, 1))\n\n        return self.norm3(x + y)\n\n\nclass Decoder(nn.Module):\n    def __init__(self, layers, norm_layer=None, projection=None):\n        super(Decoder, self).__init__()\n        self.layers = nn.ModuleList(layers)\n        self.norm = norm_layer\n        self.projection = projection\n\n    def forward(self, x, cross, x_mask=None, cross_mask=None, tau=None, delta=None):\n        for layer in self.layers:\n            x = layer(x, cross, x_mask=x_mask, cross_mask=cross_mask, tau=tau, delta=delta)\n\n        if self.norm is not None:\n            x = self.norm(x)\n\n        if self.projection is not None:\n            x = self.projection(x)\n        return x\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-09T02:15:43.567792Z","iopub.execute_input":"2025-01-09T02:15:43.568134Z","iopub.status.idle":"2025-01-09T02:15:43.584694Z","shell.execute_reply.started":"2025-01-09T02:15:43.568100Z","shell.execute_reply":"2025-01-09T02:15:43.583689Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"1. #### iTransformer","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n#from layers.Transformer_EncDec import Encoder, EncoderLayer\n#from layers.SelfAttention_Family import FullAttention, AttentionLayer\n#from layers.Embed import DataEmbedding_inverted\nimport numpy as np\n\n\nclass iTransformer(nn.Module):\n    \"\"\"\n    Paper link: https://arxiv.org/abs/2310.06625\n    \"\"\"\n\n    def __init__(self, configs):\n        super(iTransformer, self).__init__()\n        self.seq_len = configs.seq_len\n        self.pred_len = configs.pred_len\n        self.output_attention = configs.output_attention\n        self.use_norm = configs.use_norm\n        # Embedding\n        self.enc_embedding = DataEmbedding_inverted(configs.seq_len, configs.d_model, configs.embed, configs.freq,\n                                                    configs.dropout)\n        self.class_strategy = configs.class_strategy\n        # Encoder-only architecture\n        self.encoder = Encoder(\n            [\n                EncoderLayer(\n                    AttentionLayer(\n                        FullAttention(False, configs.factor, attention_dropout=configs.dropout,\n                                      output_attention=configs.output_attention), configs.d_model, configs.n_heads),\n                    configs.d_model,\n                    configs.d_ff,\n                    dropout=configs.dropout,\n                    activation=configs.activation\n                ) for l in range(configs.e_layers)\n            ],\n            norm_layer=torch.nn.LayerNorm(configs.d_model)\n        )\n        self.projector = nn.Linear(configs.d_model, configs.pred_len, bias=True)\n\n    def forecast(self, x_enc, x_mark_enc, x_dec, x_mark_dec):\n        if self.use_norm:\n            # Normalization from Non-stationary Transformer\n            means = x_enc.mean(1, keepdim=True).detach()\n            x_enc = x_enc - means\n            stdev = torch.sqrt(torch.var(x_enc, dim=1, keepdim=True, unbiased=False) + 1e-5)\n            x_enc /= stdev\n\n        _, _, N = x_enc.shape # B L N\n        # B: batch_size;    E: d_model; \n        # L: seq_len;       S: pred_len;\n        # N: number of variate (tokens), can also includes covariates\n\n        # Embedding\n        # B L N -> B N E                (B L N -> B L E in the vanilla Transformer)\n        enc_out = self.enc_embedding(x_enc, x_mark_enc) # covariates (e.g timestamp) can be also embedded as tokens\n        \n        # B N E -> B N E                (B L E -> B L E in the vanilla Transformer)\n        # the dimensions of embedded time series has been inverted, and then processed by native attn, layernorm and ffn modules\n        enc_out, attns = self.encoder(enc_out, attn_mask=None)\n\n        # B N E -> B N S -> B S N \n        dec_out = self.projector(enc_out).permute(0, 2, 1)[:, :, :N] # filter the covariates\n\n        if self.use_norm:\n            # De-Normalization from Non-stationary Transformer\n            dec_out = dec_out * (stdev[:, 0, :].unsqueeze(1).repeat(1, self.pred_len, 1))\n            dec_out = dec_out + (means[:, 0, :].unsqueeze(1).repeat(1, self.pred_len, 1))\n\n        return dec_out\n\n\n    def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec, mask=None):\n        dec_out = self.forecast(x_enc, x_mark_enc, x_dec, x_mark_dec)\n        return dec_out[:, -self.pred_len:, :]  # [B, L, D]\n#iTransformer = Model\n#import os\n#import torch\n#from model import iTransformer\n\n\nclass Exp_Basic(object):\n    def __init__(self, args):\n        self.args = args\n        self.model_dict = {\n            'iTransformer': iTransformer,\n        }\n        self.device = self._acquire_device()\n        # print(self.device)\n        self.model = self._build_model().to(self.device)\n\n    def _build_model(self):\n        raise NotImplementedError\n        return None\n\n    def _acquire_device(self):\n        if self.args.use_gpu:\n            os.environ[\"CUDA_VISIBLE_DEVICES\"] = str(\n                self.args.gpu) if not self.args.use_multi_gpu else self.args.devices\n            device = torch.device('cuda:{}'.format(self.args.gpu))\n            print('Use GPU: cuda:{}'.format(self.args.gpu))\n        else:\n            device = torch.device('cpu')\n            print('Use CPU')\n        return device\n\n    def _get_data(self):\n        pass\n\n    def vali(self):\n        pass\n\n    def train(self):\n        pass\n\n    def test(self):\n        pass\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-09T02:15:43.586257Z","iopub.execute_input":"2025-01-09T02:15:43.586600Z","iopub.status.idle":"2025-01-09T02:15:43.606342Z","shell.execute_reply.started":"2025-01-09T02:15:43.586568Z","shell.execute_reply":"2025-01-09T02:15:43.605559Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### prepare_prediction_data","metadata":{}},{"cell_type":"code","source":"#!cp -rf /kaggle/input/tes#t06/* .\n\nimport torch\nfrom torch.utils.data import DataLoader\nimport numpy as np\nimport polars as pl  \nimport pandas as pd  \n#from experiments.exp_basic import Exp_Basic\nimport pandas as pd\nimport numpy as np\nimport torch\nfrom torch.utils.data import Dataset\n\ndef prepare_prediction_data(df_raw, args, scaler=None, timeenc=0, inverse=False):\n\n    # Define borders for data slicing\n    border1 = len(df_raw) - args.seq_len\n    border2 = len(df_raw)\n\n    # Prepare input data based on feature type\n    if args.features in ['M', 'MS']:\n        cols_data = df_raw.columns[1:]  # Exclude date column\n        df_data = df_raw[cols_data]\n    elif args.features == 'S':\n        df_data = df_raw[[args.target]]\n    else:\n        raise ValueError(\"Features parameter should be 'M', 'MS', or 'S'.\")\n\n    data = df_data.values\n\n    # Prepare timestamps for predictions\n    tmp_stamp = df_raw[['date']][border1:border2]\n    tmp_stamp['date'] = pd.to_datetime(tmp_stamp.date)\n    pred_dates = pd.date_range(tmp_stamp.date.values[-1], periods=args.pred_len + 1, freq=args.freq)\n\n    df_stamp = pd.DataFrame(columns=['date'])\n    df_stamp.date = list(tmp_stamp.date.values) + list(pred_dates[1:])\n    \n    # Create time-related features based on encoding option\n    if timeenc == 0:\n        df_stamp['month'] = df_stamp.date.dt.month\n        df_stamp['day'] = df_stamp.date.dt.day\n        df_stamp['weekday'] = df_stamp.date.dt.weekday\n        df_stamp['hour'] = df_stamp.date.dt.hour\n        df_stamp['minute'] = df_stamp.date.dt.minute\n        df_stamp['minute'] = df_stamp.minute.map(lambda x: x // 15)\n        data_stamp = df_stamp.drop(['date'], axis=1).values\n    elif timeenc == 1:\n        data_stamp = time_features(pd.to_datetime(df_stamp['date'].values), freq=args.freq)\n        data_stamp = data_stamp.transpose(1, 0)\n    else:\n        raise ValueError(\"timeenc should be 0 or 1.\")\n\n    # Prepare model input data\n    data_x = data[border1:border2]\n    data_y = df_data.values[border1:border2] if not inverse else data[border1:border2]\n\n    def get_item(index):\n        s_begin = index\n        s_end = s_begin + args.seq_len\n        r_begin = s_end - args.label_len\n        r_end = r_begin + args.label_len + args.pred_len\n\n        seq_x = data_x[s_begin:s_end]\n        seq_y = data_x[r_begin:r_begin + args.label_len] if inverse else data_y[r_begin:r_begin + args.label_len]\n        seq_x_mark = data_stamp[s_begin:s_end]\n        seq_y_mark = data_stamp[r_begin:r_end]\n\n        return seq_x, seq_y, seq_x_mark, seq_y_mark\n\n    def len_function():\n        return len(data_x) - args.seq_len + 1\n\n    return data_x, data_y, data_stamp, get_item, len_function","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-09T02:15:43.607201Z","iopub.execute_input":"2025-01-09T02:15:43.607444Z","iopub.status.idle":"2025-01-09T02:15:44.093714Z","shell.execute_reply.started":"2025-01-09T02:15:43.607423Z","shell.execute_reply":"2025-01-09T02:15:44.092801Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class Exp_Long_Term_Forecast(Exp_Basic):\n    def __init__(self, args):\n        super(Exp_Long_Term_Forecast, self).__init__(args)\n\n    def _build_model(self):\n        model = self.model_dict[self.args.model](self.args).float()\n        if self.args.use_multi_gpu and self.args.use_gpu:\n            model = nn.DataParallel(model, device_ids=self.args.device_ids)\n        return model\n\n    def _get_data(self, flag):\n        data_set, data_loader = data_provider(self.args, flag)\n        return data_set, data_loader\n\n    def _select_optimizer(self):\n        model_optim = optim.Adam(self.model.parameters(), lr=self.args.learning_rate)\n        return model_optim\n\n    def _select_criterion(self):\n        criterion = nn.MSELoss()\n        return criterion\n\n    def _make_predictions(self, pred_loader, test):\n        \"\"\"执行模型预测的核心逻辑\"\"\"\n        # 加载模型\n        self.model.load_state_dict(torch.load('/kaggle/input/test04/checkpoint.pth', weights_only=True))\n        self.model.eval()\n        self.model.to(self.device)\n    \n        # 进行预测\n        preds = []\n        with torch.no_grad():\n            for i, (batch_x, batch_y, batch_x_mark, batch_y_mark) in enumerate(pred_loader):\n                batch_x = batch_x.float().to(self.device)\n                batch_y = batch_y.float().to(self.device)\n                batch_x_mark = batch_x_mark.float().to(self.device)\n                batch_y_mark = batch_y_mark.float().to(self.device)\n\n                # 准备解码器输入\n                dec_inp = torch.zeros_like(batch_y[:, -self.args.pred_len:, :]).float()\n                dec_inp = torch.cat([batch_y[:, :self.args.label_len, :], dec_inp], dim=1).float().to(self.device)\n\n                # 执行预测\n                if self.args.use_amp:\n                    with torch.cuda.amp.autocast():\n                        outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)\n                else:\n                    outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)\n\n                outputs = outputs.detach().cpu().numpy()\n                preds.append(outputs)\n\n        # 处理预测结果\n        preds = np.array(preds)\n        preds = preds.reshape(-1, preds.shape[-2], preds.shape[-1])[:, :, -78:-39]\n        \n        # 转换为Polars DataFrame\n        test = pl.from_pandas(test)\n        predictions = test.select(\"row_id\").with_columns(\n            pl.Series(\"responder_6\", preds.flatten())\n        )\n        \n        return predictions\n    \n    def predict(self, test: pd.DataFrame, lags: pd.DataFrame | None) -> pl.DataFrame | pd.DataFrame:\n        \"\"\"Make a prediction based on the provided test data and optional lags.\"\"\"\n        \n        # Store lags globally if provided\n        if lags is not None:\n            lags_ = lags\n        \n        # Initialize predictions DataFrame\n        predictions = test.select(\n            'row_id',\n            pl.lit(0.0).alias('responder_6'),\n        )\n        \n        # Get symbol IDs from the test data\n        symbol_ids = test.select('symbol_id').to_numpy()[:, 0]\n        \n        # Handle lags data\n        if lags is not None:\n            # Group and get the last record for each date and symbol\n            lags = lags.group_by([\"date_id\", \"symbol_id\"], maintain_order=True).last()\n            test = test.join(lags, on=[\"date_id\", \"symbol_id\"], how=\"left\")\n        else:\n            # Create placeholder columns for responder lags if no lags data is provided\n            test = test.with_columns(\n                (pl.lit(0.0).alias(f'responder_{idx}_lag_1') for idx in range(9))\n            )\n            \n        # Add a date column\n        test = test.with_columns(\n            (test['date_id'] * 1000 + test['time_id']).alias('date')\n        )    \n        test = test.to_pandas()\n    \n        # Select needed features and responder column\n        needed_features = [f'feature_{i:02d}' for i in range(79)]\n        needed_responder = 'responder_6_lag_1'\n        \n        # Prepare data for prediction\n        df_raw = test[['date'] + needed_features + [needed_responder]]\n\n        # 使用 pivot_table 重塑数据  \n        pivot_df = test.pivot_table(  \n            index=['date'],  \n            columns='symbol_id',  \n            values=needed_features + [needed_responder],  \n            aggfunc='first'  \n        )  \n\n        # 将多层列索引转为单层列名  \n        pivot_df.columns = [f\"{metric}_{col}\" for metric, col in pivot_df.columns]\n        pivot_df.reset_index(inplace=True)  \n        df_raw = pivot_df.fillna(value=0)\n\n        # 准备预测数据\n        data_x, data_y, data_stamp, get_item, len_function = prepare_prediction_data(\n            df_raw, self.args, scaler=None, timeenc=0, inverse=False\n        )\n        \n        # 创建数据集\n        dataset = CustomDataset(get_item, len_function)\n        pred_loader = DataLoader(dataset, batch_size=1024, shuffle=False, num_workers=4)\n\n        # 执行预测\n        predictions = self._make_predictions(pred_loader, test)\n        \n        return predictions\n\n# 辅助类定义\nclass CustomDataset(torch.utils.data.Dataset):\n    def __init__(self, get_item, len_function):\n        self.get_item = get_item\n        self.len = len_function\n\n    def __getitem__(self, index):\n        return self.get_item(index)\n\n    def __len__(self):\n        return self.len()","metadata":{"papermill":{"duration":3.506553,"end_time":"2024-12-27T02:35:29.193341","exception":false,"start_time":"2024-12-27T02:35:25.686788","status":"completed"},"tags":[],"trusted":true,"execution":{"iopub.status.busy":"2025-01-09T02:15:44.096234Z","iopub.execute_input":"2025-01-09T02:15:44.096686Z","iopub.status.idle":"2025-01-09T02:15:44.112277Z","shell.execute_reply.started":"2025-01-09T02:15:44.096659Z","shell.execute_reply":"2025-01-09T02:15:44.111249Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import argparse\nimport torch\nimport random\nimport numpy as np\nimport types  \nimport polars as pl  \nimport os  \n\nargs = torch.load('/kaggle/input/test04/parse_args.pth', weights_only=True) \nargs = types.SimpleNamespace(**args)  \nprint(args)  # 可以查看 args 的内容  \n\nExp = Exp_Long_Term_Forecast\nexp = Exp(args)\n\n#################\ntest_data_path = '/kaggle/input/jane-street-real-time-market-data-forecasting/test.parquet/date_id=0/part-0.parquet'  \ntest = pl.read_parquet(test_data_path)  \ntest = test.fill_null(value=0)\n# 第二步：加载滞后数据（如果使用滞后数据）  \nlags_data_path = '/kaggle/input/jane-street-real-time-market-data-forecasting/lags.parquet/date_id=0/part-0.parquet'  \nlags = pl.read_parquet(lags_data_path)  \nlags = lags.fill_null(value=0)\nok = exp.predict(test, lags) \nok","metadata":{"papermill":{"duration":18.355382,"end_time":"2024-12-27T02:35:47.552878","exception":false,"start_time":"2024-12-27T02:35:29.197496","status":"completed"},"tags":[],"trusted":true,"execution":{"iopub.status.busy":"2025-01-09T02:15:44.113521Z","iopub.execute_input":"2025-01-09T02:15:44.113928Z","iopub.status.idle":"2025-01-09T02:15:44.816221Z","shell.execute_reply.started":"2025-01-09T02:15:44.113883Z","shell.execute_reply":"2025-01-09T02:15:44.815240Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## stop","metadata":{}},{"cell_type":"code","source":"#import pandas as pd\n#import os\n#import re\n#from glob import glob\n#from tqdm import tqdm  \n## 获取所有 date_id 的文件路径\n#test_files = sorted(glob(os.path.join(\"/kaggle/input/janestreet-updated-simulator-for-time-series-api/debug/test.parquet/\", \"date_id=121\", \"part-0.parquet\")))\n#lags_files = sorted(glob(os.path.join(\"/kaggle/input/janestreet-updated-simulator-for-time-series-api/debug/lags.parquet/\", \"date_id=121\", \"part-0.parquet\")))\n#test_files = sorted(test_files, key=lambda x: int(re.search(r\"date_id=(\\d+)\", x).group(1)))\n#lags_files = sorted(lags_files, key=lambda x: int(re.search(r\"date_id=(\\d+)\", x).group(1)))\n#print(lags_files)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-09T02:15:44.817421Z","iopub.execute_input":"2025-01-09T02:15:44.817794Z","iopub.status.idle":"2025-01-09T02:15:44.821436Z","shell.execute_reply.started":"2025-01-09T02:15:44.817757Z","shell.execute_reply":"2025-01-09T02:15:44.820546Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#print(test_files)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-09T02:15:44.822382Z","iopub.execute_input":"2025-01-09T02:15:44.822762Z","iopub.status.idle":"2025-01-09T02:15:44.837073Z","shell.execute_reply.started":"2025-01-09T02:15:44.822731Z","shell.execute_reply":"2025-01-09T02:15:44.836275Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#date_id_pattern = re.compile(r\"date_id=(\\d+)\")\n## list for saving the prediction\n#results = []\n#prev_date_id = None  # Record for last date_id\n#for test_file in tqdm(test_files, desc=\"Processing Test Files\"):\n#    print(test_file)\n#    # get current_date_id\n#    match = date_id_pattern.search(test_file)\n#    if not match:\n#        continue\n#    current_date_id = int(match.group(1))    \n#    test_data = pl.read_parquet(test_file)\n#    print(current_date_id)\n#    print(test_data)\n#    # the predict function will be called on different time_id(time_id=0,1,2...)\n#    for time_id in tqdm(test_data[\"time_id\"].unique(), desc=f\"Processing date_id={current_date_id}\", leave=False):\n#        time_batch = test_data.filter(pl.col(\"time_id\") == time_id)\n#\n#        if current_date_id != prev_date_id and time_id == 0:\n#            # ①If we go to the next date_id, we can get the full lags data when time_id==0.\n#            lags_file = next((f for f in lags_files if f\"date_id={current_date_id}\" in f), None)\n#            lags_data = pl.read_parquet(lags_file) if lags_file else None\n##            print(f\"Processing new date_id: {current_date_id}, time_id: {time_id} with lags\")\n#        else:\n#            # ②And we get None for the rest of the time_id(1,2,3...)\n#            lags_data = None\n##            print(f\"Processing date_id: {current_date_id}, time_id: {time_id} without lags\")\n#\n#        # call predict func for every time_id batch\n#        prediction = exp.predict(time_batch, lags_data)\n#        results.append(prediction)\n#\n#    # Update prev_date_id\n#    prev_date_id = current_date_id\n#\n#results_df = pl.concat(results)\n#print(\"Prediction summary:\")\n#print(results_df.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-09T02:15:44.837937Z","iopub.execute_input":"2025-01-09T02:15:44.838454Z","iopub.status.idle":"2025-01-09T02:15:44.853637Z","shell.execute_reply.started":"2025-01-09T02:15:44.838422Z","shell.execute_reply":"2025-01-09T02:15:44.852803Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#### 主程序 ####  \n#import kaggle_evaluation.jane_street_inference_server  \n#inference_server = kaggle_evaluation.jane_street_inference_server.JSInferenceServer(exp.predict)  \n#if os.getenv('KAGGLE_IS_COMPETITION_RERUN'):  \n#    inference_server.serve()  \n#else:  \n#    inference_server.run_local_gateway(  \n#        (  \n#            '/kaggle/input/jane-street-real-time-market-data-forecasting/test.parquet',  \n#            '/kaggle/input/jane-street-real-time-market-data-forecasting/lags.parquet',  \n#        )  \n#    )\n#    print(\"YES\")\n##!ls | grep -v \"submission.parquet\" | xargs rm -rf\n!ls\n","metadata":{"papermill":{"duration":5.947126,"end_time":"2024-12-27T02:35:53.504447","exception":false,"start_time":"2024-12-27T02:35:47.557321","status":"completed"},"tags":[],"trusted":true,"execution":{"iopub.status.busy":"2025-01-09T02:15:44.854435Z","iopub.execute_input":"2025-01-09T02:15:44.854726Z","iopub.status.idle":"2025-01-09T02:15:45.002944Z","shell.execute_reply.started":"2025-01-09T02:15:44.854706Z","shell.execute_reply":"2025-01-09T02:15:45.001682Z"}},"outputs":[],"execution_count":null}]}