Set max_genes to 200 to avoid OOM issues

helical/models/c2s/config.py

@@ -57,8 +57,8 @@ class Cell2SenConfig:
         If True, the attention implementation will be set to "flash_attention_2".
         If False, the attention implementation will be set to "sdpa".
 
-    max_genes: int = None
-        Maximum number of genes to use for the model. Default is None.
+    max_genes: int = 200
+        Maximum number of genes to use for the model. Default is 200.
         If None, all genes will be used.
         If a number is provided, the genes will be sorted by expression level and the top max_genes will be used.
 
@@ -73,7 +73,7 @@ def __init__(
         organism: str = None,
         perturbation_column: str = None,
         max_new_tokens: int = 200,
-        max_genes: int = None,
+        max_genes: int = 200,
         return_fit: bool = False,
         dtype: str = "bfloat16", 
         model_size: str = "2B",  
@@ -138,4 +138,4 @@ def __init__(
             "use_flash_attn": use_flash_attn,
             "max_genes": max_genes,
             "device": device,
-        }
+        }

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "helical"
-version = "1.4.18"
+version = "1.4.19"
 authors = [
   { name="Helical Team", email="support@helical-ai.com" },
 ]