Merge pull request #84 from daniel-unyi-42/main

Fix type of 'is_token_based' parameter

Author: EliHei2

src/segger/models/segger_model.py

@@ -31,25 +31,25 @@ def forward(self, x_dict, edge_index_dict):
 class Segger(nn.Module):
     def __init__(
         self,
-        is_token_based: int,
         num_node_features: dict[str, int],
         init_emb: int = 16,
         hidden_channels: int = 32,
         num_mid_layers: int = 3,
         out_channels: int = 32,
         heads: int = 3,
+        is_token_based: bool = True,
     ):
         """
         Initializes the Segger model.
 
         Args:
-            is_token_based (int)   : Whether the model is using token-based embeddings or scRNAseq embeddings.
             num_node_features (dict[str, int]): Number of node features for each node type.
             init_emb (int)         : Initial embedding size for both 'tx' and boundary (non-token) nodes.
             hidden_channels (int)  : Number of hidden channels.
             num_mid_layers (int)   : Number of hidden layers (excluding first and last layers).
             out_channels (int)     : Number of output channels.
             heads (int)            : Number of attention heads.
+            is_token_based (bool)  : Whether the model is using token-based embeddings or scRNAseq embeddings.
         """
         super().__init__()
 

src/segger/training/train.py

@@ -61,22 +61,20 @@ def __init__(self, learning_rate: float = 1e-3, **kwargs):
 
     def from_new(
         self,
-        is_token_based: int,
         num_node_features: dict[str, int],
         init_emb: int,
         hidden_channels: int,
         out_channels: int,
         heads: int,
         num_mid_layers: int,
         aggr: str,
+        is_token_based: bool = True,
     ):
         """
         Initializes the LitSegger module with new parameters.
 
         Parameters
         ----------
-        is_token_based : int
-            Whether the model is using token-based embeddings or scRNAseq embeddings.
         num_node_features : dict[str, int]
             Number of node features for each node type.
         init_emb : int
@@ -87,22 +85,22 @@ def from_new(
             Number of output channels.
         heads : int
             Number of attention heads.
-        aggr : str
-            Aggregation method for heterogeneous graph conversion.
         num_mid_layers: int
             Number of hidden layers (excluding first and last layers).
-        metadata : Union[Tuple, Metadata]
-            Metadata for heterogeneous graph structure.
+        aggr : str
+            Aggregation method for heterogeneous graph conversion.
+        is_token_based : bool
+            Whether the model is using token-based embeddings or scRNAseq embeddings.
         """
         # Create the Segger model (ensure num_tx_tokens is passed here)
         self.model = Segger(
-            is_token_based=is_token_based,
             num_node_features=num_node_features,
             init_emb=init_emb,
             hidden_channels=hidden_channels,
             out_channels=out_channels,
             heads=heads,
             num_mid_layers=num_mid_layers,
+            is_token_based=is_token_based,
         )
         # Save hyperparameters
         self.save_hyperparameters()