{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":84493,"databundleVersionId":11305158,"sourceType":"competition"}],"dockerImageVersionId":30822,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"##########################################\n# This is a LightGBM regression model.\n# (gradient boosting decision tree model)\n##########################################\nimport os\nimport polars as pl\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom sklearn.cluster import KMeans","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-23T22:16:15.818689Z","iopub.execute_input":"2025-10-23T22:16:15.819046Z","iopub.status.idle":"2025-10-23T22:16:22.995370Z","shell.execute_reply.started":"2025-10-23T22:16:15.819011Z","shell.execute_reply":"2025-10-23T22:16:22.994242Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 1) Configuration\nDATA_PATH = '/kaggle/input/jane-street-real-time-market-data-forecasting/train.parquet'\nBATCH_SIZE = 512\nEPOCHS = 50\nLR = 1e-3\nTEMPERATURE = 0.5\nEMBED_DIM = 128\nPROJ_DIM = 64\nNUM_CLUSTERS = 2  # change as needed\nDEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-19T04:47:19.259205Z","iopub.execute_input":"2025-07-19T04:47:19.259561Z","iopub.status.idle":"2025-07-19T04:47:19.274776Z","shell.execute_reply.started":"2025-07-19T04:47:19.259505Z","shell.execute_reply":"2025-07-19T04:47:19.273642Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 2) Data loading & sampling\n# Load only feature columns, drop target\nFEATURES = [f'feature_{i:02d}' for i in range(79)]\ndf = pl.read_parquet(DATA_PATH, columns=FEATURES).sample(n=200_000, seed=42).to_pandas()\nX = df.values.astype(np.float32)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-19T04:47:23.272452Z","iopub.execute_input":"2025-07-19T04:47:23.272852Z","iopub.status.idle":"2025-07-19T04:48:13.198204Z","shell.execute_reply.started":"2025-07-19T04:47:23.272823Z","shell.execute_reply":"2025-07-19T04:48:13.197053Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 3) Define augmentations\nclass TabularAugment:\n    def __init__(self, noise_scale=0.01, dropout_prob=0.1):\n        self.noise_scale = noise_scale\n        self.dropout_prob = dropout_prob\n    def __call__(self, x):\n        # Gaussian noise\n        x = x + np.random.normal(scale=self.noise_scale, size=x.shape)\n        # Feature dropout\n        mask = np.random.rand(*x.shape) > self.dropout_prob\n        return x * mask\n\n# 4) Dataset that returns two augmented views\nclass ContrastiveDataset(Dataset):\n    def __init__(self, X, transform):\n        self.X = X\n        self.transform = transform\n    def __len__(self):\n        return len(self.X)\n    def __getitem__(self, idx):\n        x = self.X[idx]\n        return self.transform(x), self.transform(x)\n\n# 5) Encoder and projection head\nclass Encoder(nn.Module):\n    def __init__(self, input_dim, embed_dim=EMBED_DIM):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(input_dim, 512), nn.ReLU(),\n            nn.Linear(512, 256), nn.ReLU(),\n            nn.Linear(256, embed_dim)\n        )\n    def forward(self, x):\n        return self.net(x)\n\nclass ProjectionHead(nn.Module):\n    def __init__(self, embed_dim=EMBED_DIM, proj_dim=PROJ_DIM):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(embed_dim, embed_dim), nn.ReLU(),\n            nn.Linear(embed_dim, proj_dim)\n        )\n    def forward(self, z):\n        return self.net(z)\n\n# 6) Contrastive loss (NT-Xent)\ndef nt_xent_loss(z1, z2, temperature=TEMPERATURE):\n    z = torch.cat([z1, z2], dim=0)  # 2N x D\n    sim = F.cosine_similarity(z.unsqueeze(1), z.unsqueeze(0), dim=2)  # 2N x 2N\n\n    N = z1.size(0)\n    mask = torch.eye(2*N, dtype=torch.bool).to(DEVICE)\n    sim = sim / temperature\n\n    # numerator: sim between positive pairs\n    positive = torch.cat([torch.diag(sim, N), torch.diag(sim, -N)], dim=0)\n    # denominator: all except self\n    sim_masked = sim.masked_fill(mask, -9e15)\n    logits = sim_masked.exp().sum(dim=1)\n\n    loss = -torch.log(positive.exp() / logits)\n    return loss.mean()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-19T04:48:19.141030Z","iopub.execute_input":"2025-07-19T04:48:19.141355Z","iopub.status.idle":"2025-07-19T04:48:19.153164Z","shell.execute_reply.started":"2025-07-19T04:48:19.141329Z","shell.execute_reply":"2025-07-19T04:48:19.151761Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 7) Prepare dataloader\ntransform = TabularAugment()\ndataset = ContrastiveDataset(X, transform)\ndataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, drop_last=True)\n\n# 8) Model instantiation\nencoder = Encoder(input_dim=X.shape[1]).to(DEVICE)\nproj_head = ProjectionHead().to(DEVICE)\noptimizer = torch.optim.AdamW(list(encoder.parameters()) + list(proj_head.parameters()), lr=LR)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-19T04:48:22.603036Z","iopub.execute_input":"2025-07-19T04:48:22.603385Z","iopub.status.idle":"2025-07-19T04:48:24.291502Z","shell.execute_reply.started":"2025-07-19T04:48:22.603354Z","shell.execute_reply":"2025-07-19T04:48:24.290379Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 9) Training loop\nfor epoch in range(1, EPOCHS+1):\n    encoder.train(); proj_head.train()\n    total_loss = 0\n    for x1, x2 in dataloader:\n        x1 = x1.to(DEVICE).float()\n        x2 = x2.to(DEVICE).float()\n        z1 = proj_head(encoder(x1))\n        z2 = proj_head(encoder(x2))\n        loss = nt_xent_loss(z1, z2)\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n        total_loss += loss.item()\n    print(f\"Epoch {epoch:02d}: Loss = {total_loss/len(dataloader):.4f}\")\n\n# 10) Extract embeddings for clustering\nencoder.eval()\nwith torch.no_grad():\n    X_tensor = torch.tensor(X).to(DEVICE)\n    embeddings = encoder(X_tensor.float()).cpu().numpy()\n\n# 11) K-Means clustering on embeddings\nkmeans = KMeans(n_clusters=NUM_CLUSTERS, random_state=42).fit(embeddings)\nclusters = kmeans.labels_\nprint(\"Cluster distribution:\", np.bincount(clusters))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-19T04:49:16.465049Z","iopub.execute_input":"2025-07-19T04:49:16.465387Z","iopub.status.idle":"2025-07-19T06:23:12.366613Z","shell.execute_reply.started":"2025-07-19T04:49:16.465360Z","shell.execute_reply":"2025-07-19T06:23:12.364138Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 12) Save model and clusters if desired\nos.makedirs('artifacts', exist_ok=True)\ntorch.save(encoder.state_dict(), 'artifacts/encoder.pth')\nnp.save('artifacts/embeddings.npy', embeddings)\nnp.save('artifacts/clusters.npy', clusters)\nprint(\"Training complete. Artifacts saved in ./artifacts.\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import polars as pl\nimport torch\nfrom torch.utils.data import IterableDataset, DataLoader\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport random\n\n# -----------------------\n# Configuration\n# -----------------------\nDATA_PATH = \"./data/train.parquet\"\nFEATURES = [f\"feature_{i}\" for i in range(79)]\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# -----------------------\n# Tabular Augmentations\n# -----------------------\nclass TabularAugment:\n    def __init__(self, noise_std=0.05, dropout_prob=0.1):\n        self.noise_std = noise_std\n        self.dropout_prob = dropout_prob\n\n    def __call__(self, x):\n        x = x.copy()\n        # Gaussian noise\n        noise = self.noise_std * torch.randn_like(torch.tensor(x))\n        x += noise.numpy()\n        # Feature dropout\n        mask = torch.rand(len(x)) > self.dropout_prob\n        x = x * mask.numpy()\n        return x\n\n# -----------------------\n# Iterable Dataset\n# -----------------------\nclass ParquetIterableDataset(IterableDataset):\n    def __init__(self, path, features, chunk_size, transform):\n        self.path = path\n        self.features = features\n        self.chunk_size = chunk_size\n        self.transform = transform\n\n    def __iter__(self):\n        scan = pl.scan_parquet(self.path).select(self.features)\n        for batch in scan.iter_chunks(self.chunk_size):  # Use Polars' streaming iterator if available\n            arr = batch.to_numpy().astype('float32')\n            for x in arr:\n                x1 = self.transform(x)\n                x2 = self.transform(x)\n                yield torch.from_numpy(x1), torch.from_numpy(x2)\n\n\n# -----------------------\n# Model (placeholder)\n# -----------------------\nclass Encoder(nn.Module):\n    def __init__(self, input_dim=79, hidden_dim=128):\n        super().__init__()\n        self.fc1 = nn.Linear(input_dim, hidden_dim)\n        self.fc2 = nn.Linear(hidden_dim, hidden_dim)\n\n    def forward(self, x):\n        x = F.relu(self.fc1(x))\n        return self.fc2(x)\n\nclass ProjectionHead(nn.Module):\n    def __init__(self, dim=128):\n        super().__init__()\n        self.proj = nn.Sequential(\n            nn.Linear(dim, dim),\n            nn.ReLU(),\n            nn.Linear(dim, dim)\n        )\n\n    def forward(self, x):\n        return self.proj(x)\n\n# -----------------------\n# Loss (NT-Xent)\n# -----------------------\ndef nt_xent_loss(z1, z2, temperature=0.5):\n    z1 = F.normalize(z1, dim=1)\n    z2 = F.normalize(z2, dim=1)\n    representations = torch.cat([z1, z2], dim=0)\n    similarity = torch.matmul(representations, representations.T)\n    sim_exp = torch.exp(similarity / temperature)\n    mask = ~torch.eye(len(similarity), dtype=bool).to(DEVICE)\n    sim_exp = sim_exp.masked_select(mask).view(len(similarity), -1)\n    positives = torch.exp(torch.sum(z1 * z2, dim=-1) / temperature)\n    positives = torch.cat([positives, positives], dim=0)\n    loss = -torch.log(positives / sim_exp.sum(dim=1))\n    return loss.mean()\n\n# -----------------------\n# Train Loop\n# -----------------------\ntransform = TabularAugment()\ndataset = ParquetIterableDataset(DATA_PATH, FEATURES, 50_000, transform)\ndataloader = DataLoader(dataset, batch_size=512, num_workers=4, pin_memory=True)\n\nencoder = Encoder().to(DEVICE)\nproj_head = ProjectionHead().to(DEVICE)\noptimizer = torch.optim.AdamW(list(encoder.parameters()) + list(proj_head.parameters()), lr=1e-3)\nscaler = torch.amp.GradScaler(device_type='cuda')\n\nfor epoch in range(10):\n    for x1, x2 in dataloader:\n        x1, x2 = x1.to(DEVICE), x2.to(DEVICE)\n        optimizer.zero_grad()\n        with torch.cuda.amp.autocast():\n            z1 = proj_head(encoder(x1))\n            z2 = proj_head(encoder(x2))\n            loss = nt_xent_loss(z1, z2)\n        scaler.scale(loss).backward()\n        scaler.step(optimizer)\n        scaler.update()\n    print(f\"Epoch {epoch+1}: Loss = {loss.item():.4f}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-19T18:59:58.970252Z","iopub.execute_input":"2025-07-19T18:59:58.970732Z","iopub.status.idle":"2025-07-19T18:59:59.132589Z","shell.execute_reply.started":"2025-07-19T18:59:58.970660Z","shell.execute_reply":"2025-07-19T18:59:59.130402Z"}},"outputs":[],"execution_count":null}]}