Upload 8 files

LICENSE

@@ -0,0 +1,12 @@
+MIT License
+
+Copyright (c) 2025
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND.

README.md

@@ -1,3 +1,20 @@
----
-license: apache-2.0
----
+# Hugging Face Model Repo
+
+This repository contains the exported model from `awe.ipynb`.
+
+## Contents
+- `model.py` : extracted model class definitions
+- `inference.py` : load and run the model with PyTorch
+- `demo.ipynb` : quick notebook demo for inference
+- `requirements.txt` : dependencies
+- `configs/config.json` : hyperparameters/config template
+- `examples/` : sample usage scripts
+- `scripts/upload_to_hf.py` : helper to push model to the Hub
+
+## Usage
+```bash
+pip install -r requirements.txt
+python inference.py
+```
+
+Replace `YourModelClass` in `inference.py` with the actual class defined in `model.py`.

configs/config.json

@@ -0,0 +1,7 @@
+{
+  "model_name": "YourModelClass",
+  "input_shape": [1, 3, 224, 224],
+  "learning_rate": 0.001,
+  "batch_size": 32,
+  "epochs": 10
+}

demo.ipynb

@@ -0,0 +1,42 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Demo Notebook for Hugging Face Spaces\n",
+    "\n",
+    "This notebook demonstrates how to load and run inference with the extracted model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from model import *  # Import your model classes\n",
+    "from inference import load_model\n",
+    "\n",
+    "model = load_model(\"model.pt\")\n",
+    "x = torch.randn(1, 3, 224, 224)\n",
+    "y = model(x)\n",
+    "print('Output:', y.shape)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

examples/example_usage.py

@@ -0,0 +1,9 @@
+import torch
+from model import *
+from inference import load_model
+
+if __name__ == "__main__":
+    model = load_model("model.pt")
+    x = torch.randn(1, 3, 224, 224)
+    y = model(x)
+    print("Example output:", y.shape)

inference.py

@@ -0,0 +1,20 @@
+import torch
+from model import *  # import extracted model classes
+
+def load_model(weights_path="model.pt", device=None):
+    if device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    # Replace 'YourModelClass' with the actual class name from model.py
+    model = YourModelClass()  
+    state = torch.load(weights_path, map_location=device)
+    model.load_state_dict(state)
+    model.to(device)
+    model.eval()
+    return model
+
+if __name__ == "__main__":
+    model = load_model()
+    # Example dummy input - adjust to your model's expected input
+    x = torch.randn(1, 3, 224, 224)
+    y = model(x)
+    print("Output shape:", y.shape)

model.py

@@ -0,0 +1,858 @@
+# Synthesized wrapper model file (inspect and adapt before use)
+import torch
+import torch.nn as nn
+
+# --- extracted class 1 ---
+class LossWeights:
+    lambda_task: float = 1.0
+    lambda_res: float = 0.5
+    lambda_ent: float = 0.2
+
+
+# --- extracted class 2 ---
+class RRF_Ultra_CNN(nn.Module):
+    def __init__(self, input_dim=1, output_dim=1):
+        super(RRF_Ultra_CNN, self).__init__()
+        self.conv1 = nn.Conv1d(input_dim, 64, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv1d(64, 128, kernel_size=3, padding=1)
+        self.fc1 = nn.Linear(128*160, 256)
+        self.fc2 = nn.Linear(256, output_dim)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        x = F.relu(self.conv2(x))
+        x = torch.flatten(x, 1)
+        x = F.relu(self.fc1(x))
+        return torch.sigmoid(self.fc2(x))
+
+
+# --- extracted class 3 ---
+class SavantRRF_Gauge(nn.Module):
+    def __init__(self, input_dim, hidden_dim, output_dim):
+        super(SavantRRF_Gauge, self).__init__()
+        self.conv1 = nn.Conv1d(input_dim, 64, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv1d(64, 128, kernel_size=3, padding=1)
+        self.conv3 = nn.Conv1d(128, 256, kernel_size=3, padding=1)
+        self.dropout = nn.Dropout(0.25)
+        # The input size to fc1 is based on the output size of conv3.
+        # Assuming input sequence length is 160, after 3 conv layers with kernel_size 3 and padding 1,
+        # the sequence length remains 160. 256 channels * 160 length = 40960.
+        self.fc1 = nn.Linear(256*160, 512) # Corrected input size based on sequence_length=160
+        self.fc2 = nn.Linear(512, 256)
+        self.fc3 = nn.Linear(256, output_dim)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        x = F.relu(self.conv2(x))
+        x = F.relu(self.conv3(x))
+        x = torch.flatten(x, 1)
+        x = self.dropout(x)
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        return torch.sigmoid(self.fc3(x))
+
+
+# --- extracted class 4 ---
+class DiracGraphConv(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, alpha: float = 1.0, bias: bool = True):
+        super().__init__()
+        self.lin = nn.Linear(in_dim, out_dim, bias=bias)
+        self.alpha = nn.Parameter(torch.tensor(alpha, dtype=torch.float32))
+        self.bias_edge = nn.Parameter(torch.tensor(0.0, dtype=torch.float32))
+
+    @staticmethod
+    def cosine_corr(z_i: torch.Tensor, z_j: torch.Tensor, eps: float = 1e-9) -> torch.Tensor:
+        num = (z_i * z_j).sum(dim=-1)
+        den = torch.clamp(z_i.norm(dim=-1) * z_j.norm(dim=-1), min=eps)
+        return num / den
+
+    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
+        N = x.size(0)
+        row, col = edge_index
+        corr = self.cosine_corr(z[row], z[col])
+        logits = self.alpha * corr + self.bias_edge
+        device = x.device
+        E = row.size(0)
+        ones = torch.ones(E, device=device)
+        max_per_row = torch.full((N,), -1e9, device=device)
+        max_per_row = max_per_row.index_put((row,), logits, accumulate=False).scatter_reduce_(0, row, logits, reduce="amax")
+        logits_centered = logits - max_per_row[row]
+        exp_logits = torch.exp(logits_centered)
+        denom = torch.zeros(N, device=device).index_add_(0, row, exp_logits)
+        attn = exp_logits / (denom[row] + 1e-9)
+        deg = torch.zeros(N, device=device).index_add_(0, row, ones)
+        norm = 1.0 / torch.clamp(deg[row], min=1.0)
+        msgs = norm.unsqueeze(-1) * attn.unsqueeze(-1) * x[col]
+        out = torch.zeros_like(x).index_add_(0, row, msgs)
+        return self.lin(out)
+
+
+# --- extracted class 5 ---
+class GNNDiracRRF(nn.Module):
+    def __init__(self, in_dim: int, hidden_dim: int, out_dim: int, num_layers: int, z_dim: int,
+                 alpha_attn: float = 1.0, dropout: float = 0.1):
+        super().__init__()
+        self.z_dim = z_dim
+        self.layers = nn.ModuleList()
+        self.layers.append(DiracGraphConv(in_dim, hidden_dim, alpha=alpha_attn))
+        for _ in range(num_layers - 2):
+            self.layers.append(DiracGraphConv(hidden_dim, hidden_dim, alpha=alpha_attn))
+        self.layers.append(DiracGraphConv(hidden_dim, out_dim, alpha=alpha_attn))
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
+        h = x
+        for i, layer in enumerate(self.layers):
+            h = layer(h, edge_index, z)
+            if i < len(self.layers) - 1:
+                h = F.gelu(h)
+                h = self.dropout(h)
+        return h
+
+
+# --- extracted class 6 ---
+class LossWeights:
+    lambda_task: float = 1.0
+    lambda_res: float = 0.5
+    lambda_ent: float = 0.2
+
+
+# --- extracted class 7 ---
+class IcosahedralRRF(nn.Module):
+    def __init__(self, input_dim, hidden_dim, output_dim, gnn_num_layers=2, gnn_z_dim=16, gnn_alpha_attn=1.0, gnn_dropout=0.1):
+        super(IcosahedralRRF, self).__init__()
+        # 12 nodos gauge
+        self.nodes = nn.ModuleList([
+            SavantRRF_Gauge(input_dim, hidden_dim, output_dim) for _ in range(12)
+        ])
+        # Núcleo ético
+        # The input to ethical_core is the concatenation of the outputs of the 12 gauge nodes.
+        # Each gauge node outputs a tensor of shape [batch_size, output_dim].
+        # Concatenating these along dim=1 results in a shape [batch_size, 12 * output_dim].
+        self.ethical_core = nn.Linear(12 * output_dim, output_dim)
+
+        # Subconsciente (dodecaedro) using GNNDiracRRF
+        # The input dimension (in_dim) for the GNN should match the feature dimension of its input nodes.
+        # There's ambiguity in the original code about what the GNN's nodes and features are.
+        # Interpretation 1 (based on original code passing 'regulated'): GNN operates on 'batch_size' nodes, with 'output_dim' features. in_dim = output_dim.
+        # Interpretation 2 (more conventional for graph on icosahedron/dodecahedron): GNN operates on 12 or 20 nodes, with features derived from gauge outputs.
+        # Let's assume interpretation 2, where the GNN operates on the 12 gauge nodes.
+        # The features for each of these 12 nodes would be the output of the corresponding gauge node, shape [batch_size, output_dim].
+        # For a GNN layer expecting [num_nodes, in_channels], the input should be [12, output_dim] per batch item.
+        # This means the GNN's in_dim should be output_dim. This matches the current GNN init below.
+        # The GNN's out_dim should match the desired output feature dimension per node (e.g., output_dim).
+        # The number of nodes for the GNN is 12 (for icosahedral).
+
+        # Let's define the memory_map GNN assuming it operates on the 12 gauge nodes.
+        # The input features to the GNN will be the outputs of the 12 gauge nodes.
+        # Each gauge node outputs a tensor of shape [batch_size, output_dim].
+        # We will treat output_dim as the feature dimension for the GNN nodes (the 12 gauge nodes).
+        # So, in_dim for GNN = output_dim.
+        # The GNN will output features for each of the 12 nodes. Let's assume out_dim for GNN is also output_dim.
+        self.memory_map = GNNDiracRRF(in_dim=output_dim, # Feature dimension for GNN nodes (output_dim of gauge nodes)
+                                      hidden_dim=hidden_dim,
+                                      out_dim=output_dim, # Output feature dimension per GNN node
+                                      num_layers=gnn_num_layers,
+                                      z_dim=gnn_z_dim,
+                                      alpha_attn=gnn_alpha_attn,
+                                      dropout=gnn_dropout)
+
+
+    def forward(self, x, edge_index=None, z=None):
+        # x is the input to the gauge nodes, shape [batch_size, input_dim, sequence_length]
+        outputs = [node(x) for node in self.nodes]
+        # outputs is a list of 12 tensors, each [batch_size, output_dim]
+
+        # Concatenate outputs for the ethical core
+        concat = torch.cat(outputs, dim=1) # [batch_size, 12 * output_dim]
+        regulated = torch.sigmoid(self.ethical_core(concat)) # [batch_size, output_dim]
+
+        # GNN operation on the 12 gauge nodes
+        if edge_index is not None and z is not None:
+            # Prepare input for the GNN: Features for the 12 nodes (the gauge node outputs).
+            # Stack the outputs to get [batch_size, 12, output_dim]
+            stacked_outputs = torch.stack(outputs, dim=1) # [batch_size, 12, output_dim]
+
+            # Reshape for GNN input: [num_nodes, in_channels] = [12, output_dim] per batch item.
+            # Need to process batch items. Simplest is to iterate.
+            # A more efficient way is to use torch_geometric.data.Batch
+
+            gnn_outputs_list = []
+            for i in range(stacked_outputs.size(0)):
+                # GNN input features for this batch item: [12, output_dim]
+                gnn_input_features_i = stacked_outputs[i]
+
+                # Ensure edge_index and z are on the correct device
+                edge_index_i = edge_index.to(x.device)
+                z_i = z.to(x.device)
+
+                # GNN forward pass for one batch item
+                gnn_output_i = self.memory_map(gnn_input_features_i, edge_index_i, z_i) # [12, output_dim]
+                gnn_outputs_list.append(gnn_output_i)
+
+            # Stack GNN outputs back into a batch tensor: [batch_size, 12, output_dim]
+            gnn_outputs_stacked = torch.stack(gnn_outputs_list, dim=0)
+
+            # Now, how to combine the GNN output [batch_size, 12, output_dim] with the 'regulated' output [batch_size, output_dim]?
+            # The original model returned just 'regulated'.
+            # A simple approach is to maybe combine them, e.g., add, concatenate, or use the GNN output as a modulation.
+            # Let's stick to returning the aggregated GNN output as the final output when GNN is used.
+            # This changes the model's behavior compared to the original.
+
+            # Alternative: The GNN output modulates the 'regulated' output.
+            # E.g., regulated * sigmoid(aggregated_gnn_output) or similar.
+            # Let's stick to returning the aggregated GNN output when edge_index and z are provided,
+            # and the original 'regulated' output otherwise. This seems the most direct path based on the conditional in the original forward.
+
+            # Aggregate the 12 nodes' outputs from the GNN
+            aggregated_gnn_output = gnn_outputs_stacked.mean(dim=1) # [batch_size, output_dim]
+
+            return aggregated_gnn_output # [batch_size, output_dim]
+
+        else:
+             # If edge_index and z are not provided, return the output of the ethical core as before.
+             return regulated
+
+
+# --- extracted class 8 ---
+class RRF_Ultra_CNN(nn.Module):
+    def __init__(self, input_dim=1, output_dim=1):
+        super(RRF_Ultra_CNN, self).__init__()
+        self.conv1 = nn.Conv1d(input_dim, 64, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv1d(64, 128, kernel_size=3, padding=1)
+        self.fc1 = nn.Linear(128*160, 256)
+        self.fc2 = nn.Linear(256, output_dim)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        x = F.relu(self.conv2(x))
+        x = torch.flatten(x, 1)
+        x = F.relu(self.fc1(x))
+        return torch.sigmoid(self.fc2(x))
+
+
+# --- extracted class 9 ---
+class SavantRRF_Gauge(nn.Module):
+    def __init__(self, input_dim, hidden_dim, output_dim):
+        super(SavantRRF_Gauge, self).__init__()
+        self.conv1 = nn.Conv1d(input_dim, 64, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv1d(64, 128, kernel_size=3, padding=1)
+        self.conv3 = nn.Conv1d(128, 256, kernel_size=3, padding=1)
+        self.dropout = nn.Dropout(0.25)
+        # The input size to fc1 is based on the output size of conv3.
+        # Assuming input sequence length is 160, after 3 conv layers with kernel_size 3 and padding 1,
+        # the sequence length remains 160. 256 channels * 160 length = 40960.
+        self.fc1 = nn.Linear(256*160, 512) # Corrected input size based on sequence_length=160
+        self.fc2 = nn.Linear(512, 256)
+        self.fc3 = nn.Linear(256, output_dim)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        x = F.relu(self.conv2(x))
+        x = F.relu(self.conv3(x))
+        x = torch.flatten(x, 1)
+        x = self.dropout(x)
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        return torch.sigmoid(self.fc3(x))
+
+
+# --- extracted class 10 ---
+class DiracGraphConv(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, alpha: float = 1.0, bias: bool = True):
+        super().__init__()
+        self.lin = nn.Linear(in_dim, out_dim, bias=bias)
+        self.alpha = nn.Parameter(torch.tensor(alpha, dtype=torch.float32))
+        self.bias_edge = nn.Parameter(torch.tensor(0.0, dtype=torch.float32))
+
+    @staticmethod
+    def cosine_corr(z_i: torch.Tensor, z_j: torch.Tensor, eps: float = 1e-9) -> torch.Tensor:
+        num = (z_i * z_j).sum(dim=-1)
+        den = torch.clamp(z_i.norm(dim=-1) * z_j.norm(dim=-1), min=eps)
+        return num / den
+
+    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
+        N = x.size(0)
+        row, col = edge_index
+        corr = self.cosine_corr(z[row], z[col])
+        logits = self.alpha * corr + self.bias_edge
+        device = x.device
+        E = row.size(0)
+        ones = torch.ones(E, device=device)
+        max_per_row = torch.full((N,), -1e9, device=device)
+        max_per_row = max_per_row.index_put((row,), logits, accumulate=False).scatter_reduce_(0, row, logits, reduce="amax")
+        logits_centered = logits - max_per_row[row]
+        exp_logits = torch.exp(logits_centered)
+        denom = torch.zeros(N, device=device).index_add_(0, row, exp_logits)
+        attn = exp_logits / (denom[row] + 1e-9)
+        deg = torch.zeros(N, device=device).index_add_(0, row, ones)
+        norm = 1.0 / torch.clamp(deg[row], min=1.0)
+        msgs = norm.unsqueeze(-1) * attn.unsqueeze(-1) * x[col]
+        out = torch.zeros_like(x).index_add_(0, row, msgs)
+        return self.lin(out)
+
+
+# --- extracted class 11 ---
+class GNNDiracRRF(nn.Module):
+    def __init__(self, in_dim: int, hidden_dim: int, out_dim: int, num_layers: int, z_dim: int,
+                 alpha_attn: float = 1.0, dropout: float = 0.1):
+        super().__init__()
+        self.z_dim = z_dim
+        self.layers = nn.ModuleList()
+        self.layers.append(DiracGraphConv(in_dim, hidden_dim, alpha=alpha_attn))
+        for _ in range(num_layers - 2):
+            self.layers.append(DiracGraphConv(hidden_dim, hidden_dim, alpha=alpha_attn))
+        self.layers.append(DiracGraphConv(hidden_dim, out_dim, alpha=alpha_attn))
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
+        h = x
+        for i, layer in enumerate(self.layers):
+            h = layer(h, edge_index, z)
+            if i < len(self.layers) - 1:
+                h = F.gelu(h)
+                h = self.dropout(h)
+        return h
+
+
+# --- extracted class 12 ---
+class LossWeights:
+    lambda_task: float = 1.0
+    lambda_res: float = 0.5
+    lambda_ent: float = 0.2
+
+
+# --- extracted class 13 ---
+class IcosahedralRRF(nn.Module):
+    def __init__(self, input_dim, hidden_dim, output_dim, gnn_num_layers=2, gnn_z_dim=16, gnn_alpha_attn=1.0, gnn_dropout=0.1):
+        super(IcosahedralRRF, self).__init__()
+        # 12 nodos gauge
+        self.nodes = nn.ModuleList([
+            SavantRRF_Gauge(input_dim, hidden_dim, output_dim) for _ in range(12)
+        ])
+        # Núcleo ético
+        # The input to ethical_core is the concatenation of the outputs of the 12 gauge nodes.
+        # Each gauge node outputs a tensor of shape [batch_size, output_dim].
+        # Concatenating these along dim=1 results in a shape [batch_size, 12 * output_dim].
+        self.ethical_core = nn.Linear(12 * output_dim, output_dim)
+
+        # Subconsciente (dodecaedro) using GNNDiracRRF
+        # The input dimension (in_dim) for the GNN should match the feature dimension of its input nodes.
+        # There's ambiguity in the original code about what the GNN's nodes and features are.
+        # Interpretation 1 (based on original code passing 'regulated'): GNN operates on 'batch_size' nodes, with 'output_dim' features. in_dim = output_dim.
+        # Interpretation 2 (more conventional for graph on icosahedron/dodecahedron): GNN operates on 12 or 20 nodes, with features derived from gauge outputs.
+        # Let's assume interpretation 2, where the GNN operates on the 12 gauge nodes.
+        # The features for each of these 12 nodes would be the output of the corresponding gauge node, shape [batch_size, output_dim].
+        # For a GNN layer expecting [num_nodes, in_channels], the input should be [12, output_dim] per batch item.
+        # This means the GNN's in_dim should be output_dim. This matches the current GNN init below.
+        # The GNN's out_dim should match the desired output feature dimension per node (e.g., output_dim).
+        # The number of nodes for the GNN is 12 (for icosahedral).
+
+        # Let's define the memory_map GNN assuming it operates on the 12 gauge nodes.
+        # The input features to the GNN will be the outputs of the 12 gauge nodes.
+        # Each gauge node outputs a tensor of shape [batch_size, output_dim].
+        # We will treat output_dim as the feature dimension for the GNN nodes (the 12 gauge nodes).
+        # So, in_dim for GNN = output_dim.
+        # The GNN will output features for each of the 12 nodes. Let's assume out_dim for GNN is also output_dim.
+        self.memory_map = GNNDiracRRF(in_dim=output_dim, # Feature dimension for GNN nodes (output_dim of gauge nodes)
+                                      hidden_dim=hidden_dim,
+                                      out_dim=output_dim, # Output feature dimension per GNN node
+                                      num_layers=gnn_num_layers,
+                                      z_dim=gnn_z_dim,
+                                      alpha_attn=gnn_alpha_attn,
+                                      dropout=gnn_dropout)
+
+
+    def forward(self, x, edge_index=None, z=None):
+        # x is the input to the gauge nodes, shape [batch_size, input_dim, sequence_length]
+        outputs = [node(x) for node in self.nodes]
+        # outputs is a list of 12 tensors, each [batch_size, output_dim]
+
+        # Concatenate outputs for the ethical core
+        concat = torch.cat(outputs, dim=1) # [batch_size, 12 * output_dim]
+        regulated = torch.sigmoid(self.ethical_core(concat)) # [batch_size, output_dim]
+
+        # GNN operation on the 12 gauge nodes
+        if edge_index is not None and z is not None:
+            # Prepare input for the GNN: Features for the 12 nodes (the gauge node outputs).
+            # Stack the outputs to get [batch_size, 12, output_dim]
+            stacked_outputs = torch.stack(outputs, dim=1) # [batch_size, 12, output_dim]
+
+            # Reshape for GNN input: [num_nodes, in_channels] = [12, output_dim] per batch item.
+            # Need to process batch items. Simplest is to iterate.
+            # A more efficient way is to use torch_geometric.data.Batch
+
+            gnn_outputs_list = []
+            for i in range(stacked_outputs.size(0)):
+                # GNN input features for this batch item: [12, output_dim]
+                gnn_input_features_i = stacked_outputs[i]
+
+                # Ensure edge_index and z are on the correct device
+                edge_index_i = edge_index.to(x.device)
+                z_i = z.to(x.device)
+
+                # GNN forward pass for one batch item
+                gnn_output_i = self.memory_map(gnn_input_features_i, edge_index_i, z_i) # [12, output_dim]
+                gnn_outputs_list.append(gnn_output_i)
+
+            # Stack GNN outputs back into a batch tensor: [batch_size, 12, output_dim]
+            gnn_outputs_stacked = torch.stack(gnn_outputs_list, dim=0)
+
+            # Now, how to combine the GNN output [batch_size, 12, output_dim] with the 'regulated' output [batch_size, output_dim]?
+            # The original model returned just 'regulated'.
+            # A simple approach is to maybe combine them, e.g., add, concatenate, or use the GNN output as a modulation.
+            # Let's stick to returning the aggregated GNN output as the final output when GNN is used.
+            # This changes the model's behavior compared to the original.
+
+            # Alternative: The GNN output modulates the 'regulated' output.
+            # E.g., regulated * sigmoid(aggregated_gnn_output) or similar.
+            # Let's stick to returning the aggregated GNN output when edge_index and z are provided,
+            # and the original 'regulated' output otherwise. This seems the most direct path based on the conditional in the original forward.
+
+            # Aggregate the 12 nodes' outputs from the GNN
+            aggregated_gnn_output = gnn_outputs_stacked.mean(dim=1) # [batch_size, output_dim]
+
+            return aggregated_gnn_output # [batch_size, output_dim]
+
+        else:
+             # If edge_index and z are not provided, return the output of the ethical core as before.
+             return regulated
+
+
+# --- extracted class 14 ---
+class LossWeights:
+    lambda_task: float = 1.0
+    lambda_res: float = 0.5
+    lambda_ent: float = 0.2
+
+
+# --- extracted class 15 ---
+class IcosahedralRRFDataset(InMemoryDataset):
+    def __init__(self, num_graphs: int = 64, k_modes: int = 16, feat_dim: int = 8,
+                 task_type: str = 'classification', split: str = 'train', transform=None, pre_transform=None):
+        super().__init__('.', transform, pre_transform)
+        self.task_type = task_type
+        self.num_graphs = num_graphs
+        self.k_modes = k_modes
+        self.feat_dim = feat_dim
+
+        # Generate graphs and process them
+        data_list = []
+        rng = np.random.default_rng(42 if split == 'train' else (43 if split == 'val' else 44))
+
+        for i in range(num_graphs):
+            G = nx.icosahedral_graph()
+            n_nodes = G.number_of_nodes()
+
+            # Build Dirac operator and compute spectral modes
+            D = build_dirac_operator(G, normalize=True)
+            # Use the modified dirac_eigendecomp that uses np.linalg.eigh
+            vals, vecs = dirac_eigendecomp(D, k=k_modes)
+            Z = node_spectral_coords_from_dirac(vecs, n_nodes)  # N x k
+
+            # Get edge index
+            edge_list = list(G.edges())
+            edge_index = torch.tensor(edge_list, dtype=torch.long).t().contiguous()
+            # Add reverse edges for undirected graph
+            row, col = edge_index
+            edge_index = torch.cat([edge_index, torch.stack([col, row], dim=0)], dim=1)
+
+            # Generate synthetic node features (x) and labels (y)
+            # Features: [n_nodes, feat_dim]
+            x = torch.randn(n_nodes, feat_dim, dtype=torch.float32)
+
+            # Labels: based on task_type
+            if task_type == 'classification':
+                # Example: Binary classification based on a simple rule, e.g., sum of features > threshold
+                threshold = 0.0 # Example threshold
+                y = (x.sum(dim=-1) > threshold).long() # [n_nodes]
+            elif task_type == 'regression':
+                 # Example: Regression target based on sum of features
+                 y = x.sum(dim=-1) # [n_nodes]
+            else:
+                raise ValueError("task_type must be 'classification' or 'regression'")
+
+
+            # Create Data object
+            # Note: The IcosahedralRRF model expects input 'x' as [batch_size, input_dim, sequence_length],
+            # edge_index [2, num_edges], and z [num_nodes, z_dim].
+            # The IcosahedralRRFDataset provides batch.x [num_nodes, feat_dim], batch.edge_index [2, num_edges], and batch.U [num_nodes, k_modes].
+            # There is a mismatch in the expected input format for the IcosahedralRRF model's forward pass when using the DataLoader.
+            # The IcosahedralRRF expects a single batch tensor `x` for the gauge nodes, and graph data (edge_index, z) for the GNN part which operates on gauge outputs.
+            # The IcosahedralRRFDataset provides node features `batch.x` that are intended as features *for the graph nodes themselves*, not as input to the gauge nodes.
+            # The current IcosahedralRRF forward pass processes a single input `x` [batch_size, input_dim, sequence_length] through all gauge nodes.
+            # The GNN then operates on the *outputs* of these gauge nodes, using the provided edge_index and z.
+
+            # To use the IcosahedralRRFDataset with the current IcosahedralRRF model structure,
+            # we need to map the dataset's structure to the model's expectations.
+            # The dataset provides graphs, each with nodes (typically 12 for icosahedral), node features (batch.x), edge_index, and spectral coords (batch.U).
+            # The IcosahedralRRF model has 12 gauge nodes, each designed to process a sequence [input_dim, sequence_length].
+            # It seems there is a conceptual mismatch in how the IcosahedralRRFDataset is structured (graph-centric with node features)
+            # and how the IcosahedralRRF model processes input (sequence-centric through gauge nodes first).
+
+            # Alternative Interpretation: The IcosahedralRRFDataset is meant to provide data where each *graph* is a sample in the batch.
+            # batch.x would be the concatenated node features for all graphs in the batch: [total_num_nodes_in_batch, feat_dim].
+            # batch.edge_index would be the block-diagonal edge indices for all graphs: [2, total_num_edges_in_batch].
+            # batch.U would be the concatenated spectral coordinates for all nodes: [total_num_nodes_in_batch, k_modes].
+            # In this case, the input to the IcosahedralRRF model's forward pass is still expected to be a single tensor `x` for the gauge nodes.
+            # The IcosahedralRRFDataset does *not* provide this `x` input directly in the expected format.
+
+            # There is a fundamental incompatibility in how the IcosahedralRRFDataset provides data (graph-batching)
+            # and how the IcosahedralRRF model expects input (single batch of sequences + graph data for GNN).
+
+            # To make this cell runnable, we need to either:
+            # 1. Modify the IcosahedralRRF model's forward pass to handle graph batches from DataLoader.
+            # 2. Modify the IcosahedralRRFDataset or create a custom Dataset/DataLoader that provides data in the format expected by the IcosahedralRRF model.
+            # 3. Use a simplified evaluation approach that aligns with the synthetic data generation method used in the training loop (single batch).
+
+            # Given the current structure, the simplest approach to get the cell running is to align the evaluation data generation
+            # with the training data generation (single synthetic batch) and evaluate on that.
+            # This bypasses the DataLoader incompatibility but doesn't fully test with graph batching.
+
+            # Let's revert to generating a single synthetic batch for evaluation, similar to training.
+            # This requires defining x_val and y_val outside the DataLoader loop.
+
+            # Reverting the evaluation loop to use the single synthetic batch approach:
+
+            # Check if x_val and y_val are defined (from previous code cell)
+            if 'x_val' not in locals() or 'y_val' not in locals():
+                 # Generate synthetic validation data if not already defined
+                 val_batch_size = 16 # Example validation batch size
+                 x_val = torch.randn(val_batch_size, input_dim, sequence_length, dtype=torch.float32).to(device)
+                 y_val = torch.randint(0, 2, (val_batch_size,), dtype=torch.long).to(device) # Binary labels
+                 print("Generated synthetic validation data for evaluation.")
+
+        # Ensure z and edge_index are on the correct device
+        if 'z' in locals() and 'edge_index' in locals():
+            z = z.to(device)
+            edge_index = edge_index.to(device)
+        else:
+            print("⚠️ Warning: Graph data (z, edge_index) not found. Skipping evaluation.")
+            # Exit the evaluation block if graph data is missing
+            # break # This will exit the with torch.no_grad(): block - REMOVED/COMMENTED OUT DUE TO SyntaxError
+            pass # Use pass instead of break to avoid SyntaxError outside a loop
+
+        # Forward pass on validation data using the single batch
+        # Pass the validation input features (x_val), edge index, and spectral coordinates (z) through the model
+        val_outputs = hybrid_model(x_val, edge_index, z) # Shape: [val_batch_size, output_dim]
+
+        # Calculate the validation loss (using BCEWithLogitsLoss as corrected in training)
+        val_loss = F.binary_cross_entropy_with_logits(val_outputs.squeeze(-1), y_val.float())
+
+        # Calculate evaluation metrics (e.g., accuracy for binary classification)
+        # Convert logits to predicted class (0 or 1)
+        predicted_classes = (torch.sigmoid(val_outputs.squeeze(-1)) > 0.5).long()
+
+        # Calculate accuracy
+        correct_predictions = (predicted_classes == y_val).sum().item()
+        accuracy = correct_predictions / val_batch_size
+
+        print(f'Validation Loss: {val_loss.item():.4f}, Validation Accuracy: {accuracy:.4f}')
+
+
+# --- extracted class 16 ---
+class DiracGraphConv(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, alpha: float = 1.0, bias: bool = True):
+        super().__init__()
+        self.lin = nn.Linear(in_dim, out_dim, bias=bias)
+        self.alpha = nn.Parameter(torch.tensor(alpha, dtype=torch.float32))
+        self.bias_edge = nn.Parameter(torch.tensor(0.0, dtype=torch.float32))
+
+    @staticmethod
+    def cosine_corr(z_i: torch.Tensor, z_j: torch.Tensor, eps: float = 1e-9) -> torch.Tensor:
+        num = (z_i * z_j).sum(dim=-1)
+        den = torch.clamp(z_i.norm(dim=-1) * z_j.norm(dim=-1), min=eps)
+        return num / den
+
+    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
+        N = x.size(0)
+        row, col = edge_index
+        corr = self.cosine_corr(z[row], z[col])
+        logits = self.alpha * corr + self.bias_edge
+        device = x.device
+        E = row.size(0)
+        ones = torch.ones(E, device=device)
+        max_per_row = torch.full((N,), -1e9, device=device)
+        max_per_row = max_per_row.index_put((row,), logits, accumulate=False).scatter_reduce_(0, row, logits, reduce="amax")
+        logits_centered = logits - max_per_row[row]
+        exp_logits = torch.exp(logits_centered)
+        denom = torch.zeros(N, device=device).index_add_(0, row, exp_logits)
+        attn = exp_logits / (denom[row] + 1e-9)
+        deg = torch.zeros(N, device=device).index_add_(0, row, ones)
+        norm = 1.0 / torch.clamp(deg[row], min=1.0)
+        msgs = norm.unsqueeze(-1) * attn.unsqueeze(-1) * x[col]
+        out = torch.zeros_like(x).index_add_(0, row, msgs)
+        return self.lin(out)
+
+
+# --- extracted class 17 ---
+class GNNDiracRRF(nn.Module):
+    def __init__(self, in_dim: int, hidden_dim: int, out_dim: int, num_layers: int, z_dim: int,
+                 alpha_attn: float = 1.0, dropout: float = 0.1):
+        super().__init__()
+        self.z_dim = z_dim
+        self.layers = nn.ModuleList()
+        self.layers.append(DiracGraphConv(in_dim, hidden_dim, alpha=alpha_attn))
+        for _ in range(num_layers - 2):
+            self.layers.append(DiracGraphConv(hidden_dim, hidden_dim, alpha=alpha_attn))
+        self.layers.append(DiracGraphConv(hidden_dim, out_dim, alpha=alpha_attn))
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
+        h = x
+        for i, layer in enumerate(self.layers):
+            h = layer(h, edge_index, z)
+            if i < len(self.layers) - 1:
+                h = F.gelu(h)
+                h = self.dropout(h)
+        return h
+
+
+# --- extracted class 18 ---
+class DiracGraphConv(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, alpha: float = 1.0, bias: bool = True):
+        super().__init__()
+        self.lin = nn.Linear(in_dim, out_dim, bias=bias)
+        self.alpha = nn.Parameter(torch.tensor(alpha, dtype=torch.float32))
+        self.bias_edge = nn.Parameter(torch.tensor(0.0, dtype=torch.float32))
+
+    @staticmethod
+    def cosine_corr(z_i: torch.Tensor, z_j: torch.Tensor, eps: float = 1e-9) -> torch.Tensor:
+        num = (z_i * z_i).sum(dim=-1) # Corrected dot product: z_i * z_j
+        den = torch.clamp(z_i.norm(dim=-1) * z_j.norm(dim=-1), min=eps)
+        return num / den
+
+    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
+        N = x.size(0)
+        row, col = edge_index
+        corr = self.cosine_corr(z[row], z[col])
+        logits = self.alpha * corr + self.bias_edge
+        device = x.device
+        E = row.size(0)
+        ones = torch.ones(E, device=device)
+        max_per_row = torch.full((N,), -1e9, device=device)
+        max_per_row = max_per_row.index_put((row,), logits, accumulate=False).scatter_reduce_(0, row, logits, reduce="amax")
+        logits_centered = logits - max_per_row[row]
+        exp_logits = torch.exp(logits_centered)
+        denom = torch.zeros(N, device=device).index_add_(0, row, exp_logits)
+        attn = exp_logits / (denom[row] + 1e-9)
+        deg = torch.zeros(N, device=device).index_add_(0, row, ones)
+        norm = 1.0 / torch.clamp(deg[row], min=1.0)
+        msgs = norm.unsqueeze(-1) * attn.unsqueeze(-1) * x[col]
+        out = torch.zeros_like(x).index_add_(0, row, msgs)
+        return self.lin(out)
+
+
+# --- extracted class 19 ---
+class GNNDiracRRF(nn.Module):
+    def __init__(self, in_dim: int, hidden_dim: int, out_dim: int, num_layers: int, z_dim: int,
+                 alpha_attn: float = 1.0, dropout: float = 0.1):
+        super().__init__()
+        self.z_dim = z_dim
+        self.layers = nn.ModuleList()
+        self.layers.append(DiracGraphConv(in_dim, hidden_dim, alpha=alpha_attn))
+        for _ in range(num_layers - 2):
+            self.layers.append(DiracGraphConv(hidden_dim, hidden_dim, alpha=alpha_attn))
+        self.layers.append(DiracGraphConv(hidden_dim, out_dim, alpha=alpha_attn))
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
+        h = x
+        for i, layer in enumerate(self.layers):
+            h = layer(h, edge_index, z)
+            if i < len(self.layers) - 1:
+                h = F.gelu(h)
+                h = self.dropout(h)
+        return h
+
+
+# --- extracted class 20 ---
+class SavantRRF_Gauge(nn.Module):
+    def __init__(self, input_dim, hidden_dim, output_dim):
+        super(SavantRRF_Gauge, self).__init__()
+        self.conv1 = nn.Conv1d(input_dim, 64, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv1d(64, 128, kernel_size=3, padding=1)
+        self.conv3 = nn.Conv1d(128, 256, kernel_size=3, padding=1)
+        self.dropout = nn.Dropout(0.25)
+        # Assuming input sequence length is 160
+        self.fc1 = nn.Linear(256*160, 512)
+        self.fc2 = nn.Linear(512, 256)
+        self.fc3 = nn.Linear(256, output_dim)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        x = F.relu(self.conv2(x))
+        x = F.relu(self.conv3(x))
+        x = torch.flatten(x, 1)
+        x = self.dropout(x)
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        return torch.sigmoid(self.fc3(x))
+
+
+# --- extracted class 21 ---
+class IcosahedralRRF(nn.Module):
+    def __init__(self, input_dim, hidden_dim, output_dim, gnn_num_layers=2, gnn_z_dim=16, gnn_alpha_attn=1.0, gnn_dropout=0.1):
+        super(IcosahedralRRF, self).__init__()
+        # 12 nodos gauge
+        self.nodes = nn.ModuleList([
+            SavantRRF_Gauge(input_dim, hidden_dim, output_dim) for _ in range(12)
+        ])
+        # Núcleo ético
+        self.ethical_core = nn.Linear(12 * output_dim, output_dim)
+
+        # Subconsciente (dodecaedro/icosaedro) using GNNDiracRRF
+        # The GNN operates on the 12 gauge node outputs.
+        # The input features to the GNN are the outputs of the 12 gauge nodes, shape [batch_size, output_dim].
+        # For GNN layer, input is [num_nodes, in_channels] = [12, output_dim] per batch item.
+        self.memory_map = GNNDiracRRF(in_dim=output_dim,
+                                      hidden_dim=hidden_dim,
+                                      out_dim=output_dim,
+                                      num_layers=gnn_num_layers,
+                                      z_dim=gnn_z_dim,
+                                      alpha_attn=gnn_alpha_attn,
+                                      dropout=gnn_dropout)
+
+
+    def forward(self, x, edge_index=None, z=None):
+        # x is the input to the gauge nodes, shape [batch_size, input_dim, sequence_length]
+        outputs = [node(x) for node in self.nodes]
+        # outputs is a list of 12 tensors, each [batch_size, output_dim]
+
+        # Concatenate outputs for the ethical core
+        concat = torch.cat(outputs, dim=1) # [batch_size, 12 * output_dim]
+        regulated = torch.sigmoid(self.ethical_core(concat)) # [batch_size, output_dim]
+
+        # GNN operation on the 12 gauge nodes
+        if edge_index is not None and z is not None:
+             # Prepare input for the GNN: Features for the 12 nodes (the gauge node outputs).
+            stacked_outputs = torch.stack(outputs, dim=1) # [batch_size, 12, output_dim]
+
+            gnn_outputs_list = []
+            for i in range(stacked_outputs.size(0)):
+                gnn_input_features_i = stacked_outputs[i]
+                edge_index_i = edge_index.to(x.device)
+                z_i = z.to(x.device)
+                gnn_output_i = self.memory_map(gnn_input_features_i, edge_index_i, z_i) # [12, output_dim]
+                gnn_outputs_list.append(gnn_output_i)
+
+            gnn_outputs_stacked = torch.stack(gnn_outputs_list, dim=0)
+            aggregated_gnn_output = gnn_outputs_stacked.mean(dim=1) # [batch_size, output_dim]
+
+            return aggregated_gnn_output # [batch_size, output_dim]
+
+        else:
+             return regulated
+
+
+# --- extracted class 22 ---
+class DiracGraphConv(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, alpha: float = 1.0, bias: bool = True):
+        super().__init__()
+        self.lin = nn.Linear(in_dim, out_dim, bias=bias)
+        self.alpha = nn.Parameter(torch.tensor(alpha, dtype=torch.float32))
+        self.bias_edge = nn.Parameter(torch.tensor(0.0, dtype=torch.float32))
+
+    @staticmethod
+    def cosine_corr(z_i: torch.Tensor, z_j: torch.Tensor, eps: float = 1e-9) -> torch.Tensor:
+        num = (z_i * z_j).sum(dim=-1)
+        den = torch.clamp(z_i.norm(dim=-1) * z_j.norm(dim=-1), min=eps)
+        return num / den
+
+    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
+        N = x.size(0)
+        row, col = edge_index
+        # Ensure z has correct shape for cosine_corr
+        # z should have shape [num_nodes, z_dim]
+        # x has shape [num_nodes, in_dim]
+        # When called from GNNDiracRRF, num_nodes is 12 (for icosahedral)
+        # z[row] and z[col] should broadcast correctly with x[col]
+        corr = self.cosine_corr(z[row], z[col])
+        logits = self.alpha * corr + self.bias_edge
+        device = x.device
+        E = row.size(0)
+        ones = torch.ones(E, device=device)
+        # Use scatter_reduce_ to calculate max per row
+        max_per_row = torch.full((N,), -1e9, device=device)
+        max_per_row = max_per_row.index_put((row,), logits, accumulate=False).scatter_reduce_(0, row, logits, reduce="amax")
+        logits_centered = logits - max_per_row[row]
+        exp_logits = torch.exp(logits_centered)
+        denom = torch.zeros(N, device=device).index_add_(0, row, exp_logits)
+        attn = exp_logits / (denom[row] + 1e-9)
+        deg = torch.zeros(N, device=device).index_add_(0, row, ones)
+        norm = 1.0 / torch.clamp(deg[row], min=1.0)
+        msgs = norm.unsqueeze(-1) * attn.unsqueeze(-1) * x[col]
+        out = torch.zeros_like(x).index_add_(0, row, msgs)
+        return self.lin(out)
+
+
+# --- extracted class 23 ---
+class GNNDiracRRF(nn.Module):
+    def __init__(self, in_dim: int, hidden_dim: int, out_dim: int, num_layers: int, z_dim: int,
+                 alpha_attn: float = 1.0, dropout: float = 0.1):
+        super().__init__()
+        self.z_dim = z_dim
+        self.layers = nn.ModuleList()
+        # Ensure DiracGraphConv is defined before this line
+        self.layers.append(DiracGraphConv(in_dim, hidden_dim, alpha=alpha_attn))
+        for _ in range(num_layers - 2):
+            self.layers.append(DiracGraphConv(hidden_dim, hidden_dim, alpha=alpha_attn))
+        self.layers.append(DiracGraphConv(hidden_dim, out_dim, alpha=alpha_attn))
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
+        h = x
+        for i, layer in enumerate(self.layers):
+            h = layer(h, edge_index, z)
+            if i < len(self.layers) - 1:
+                h = F.gelu(h)
+                h = self.dropout(h)
+        return h
+
+
+# --- extracted class 24 ---
+class RRF_Dataset(Dataset):
+    def __init__(self, strain, weights, seq_len=160): # Use seq_len=160 to match model input
+        self.seq_len = seq_len
+        self.strain = strain
+        self.weights = weights
+        print(f"Debug: RRF_Dataset __init__ - len(strain): {len(strain)}, seq_len: {self.seq_len}") # Debug print
+        # Calculate n only if strain is long enough
+        if len(strain) >= seq_len:
+            self.n = len(strain) // seq_len
+        else:
+            self.n = 0 # Set n to 0 if strain is too short
+        print(f"Debug: RRF_Dataset __init__ - Calculated self.n: {self.n}") # New debug print
+        # Add a check to ensure there's at least one sequence
+        if self.n == 0:
+            raise ValueError(f"Strain data length ({len(strain)}) is less than sequence length ({seq_len}). Cannot create any samples.")
+
+
+    def __len__(self):
+        return self.n
+
+    def __getitem__(self, idx):
+        start = idx * self.seq_len
+        # Extract the strain sequence x
+        x = self.strain[start:start+self.seq_len] # Shape: [seq_len]
+
+        # Use the mean of the provided weights as the global resonance factor w
+        w = np.mean(self.weights)  # global resonance factor
+
+        # Define the target label y as the mean of the strain sequence x, scaled by w
+        # This creates a regression target derived from the strain data.
+        y = np.mean(x) * w  # synthetic label (proxy resonance)
+
+        # Convert x and y to PyTorch tensors with float dtype
+        # The model expects input x as [1, seq_len] for a single sample, so add unsqueeze(0)
+        return torch.tensor(x).float().unsqueeze(0), torch.tensor(y).float()
+
+
+
+def load_model_state(path, model_instance, map_location='cpu'):
+    '''Helper: load state_dict from path into model_instance (PyTorch).'''
+    state = torch.load(path, map_location=map_location)
+    if isinstance(state, dict) and ('state_dict' in state and isinstance(state['state_dict'], dict)):
+        state = state['state_dict']
+    model_instance.load_state_dict(state)
+    return model_instance

requirements.txt

@@ -0,0 +1,26 @@
+IPython
+astroquery
+dataclasses
+datasets
+gnn_dirac_rrf
+google
+gwpy
+kagglehub
+matplotlib
+networkx
+numpy
+os
+pandas
+pesummary
+plotly
+random
+requests
+safetensors
+scikit-learn
+scipy
+shutil
+torch
+torch_geometric
+torchsummary
+typing
+zipfile