Merge pull request #199 from databio/bug_fix_r2v_scembed

fix tokenizer creation issue with ScEmbed (#198)

.gitignore

@@ -197,4 +197,5 @@ qdrant_storage/
 
 local_cache
 
-lightning_logs
+lightning_logs
+data/

geniml/region2vec/main.py

@@ -140,11 +140,12 @@ def _load_local_model(self, model_path: str, vocab_path: str, config_path: str):
         :param str model_path: Path to the model checkpoint.
         :param str vocab_path: Path to the vocabulary file.
         """
-        _model, tokenizer, config = load_local_region2vec_model(
-            model_path, vocab_path, config_path
-        )
+        _model, config = load_local_region2vec_model(model_path, config_path)
+        tokenizer = TreeTokenizer(vocab_path)
+
         self._model = _model
         self.tokenizer = tokenizer
+
         self.trained = True
         if POOLING_METHOD_KEY in config:
             self.pooling_method = config[POOLING_METHOD_KEY]

geniml/region2vec/utils.py

@@ -452,18 +452,15 @@ def export_region2vec_model(
 
 def load_local_region2vec_model(
     model_path: str,
-    vocab_path: str,
     config_path: str,
-) -> Tuple[Region2Vec, TreeTokenizer, dict]:
+) -> Tuple[Region2Vec, dict]:
     """
     Load a region2vec model from a local directory
 
     :param str model_path: The path to the model checkpoint file
     :param str config_path: The path to the model config file
     :param str vocab_path: The path to the model vocabulary file
     """
-    # init the tokenizer - only one option for now
-    tokenizer = TreeTokenizer(vocab_path)
 
     # load the model state dict (weights)
     params = torch.load(model_path)
@@ -491,7 +488,7 @@ def load_local_region2vec_model(
 
     model.load_state_dict(params)
 
-    return model, tokenizer, config
+    return model, config
 
 
 class Region2VecDataset:

geniml/scembed/main.py

@@ -131,11 +131,12 @@ def _load_local_model(self, model_path: str, vocab_path: str, config_path: str):
         :param str model_path: Path to the model checkpoint.
         :param str vocab_path: Path to the vocabulary file.
         """
-        _model, tokenizer, config = load_local_region2vec_model(
-            model_path, vocab_path, config_path
-        )
+        _model, config = load_local_region2vec_model(model_path, config_path)
+        tokenizer = AnnDataTokenizer(vocab_path, verbose=True)
+
         self._model = _model
         self.tokenizer = tokenizer
+
         if POOLING_METHOD_KEY in config:
             self.pooling_method = config[POOLING_METHOD_KEY]