debug

alexandrainst · Sep 27, 2024 · 523f999 · 523f999
1 parent 00fffed
commit 523f999
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/coral/finetune.py b/src/coral/finetune.py
@@ -15,6 +15,7 @@
 from .data import load_data_for_finetuning
 from .data_models import ModelSetup
 from .model_setup import load_model_setup
+from .ngram import train_and_store_ngram_model
 from .utils import push_model_to_hub
 
 logger = logging.getLogger(__package__)
@@ -71,8 +72,8 @@ def finetune(config: DictConfig) -> None:
 
     # model.save_pretrained(save_directory=config.model_dir)
 
-    # if hasattr(config.model, "use_decoder") and config.model.use_decoder:
-    #     train_and_store_ngram_model(config=config)
+    if hasattr(config.model, "use_decoder") and config.model.use_decoder:
+        train_and_store_ngram_model(config=config)
 
     if config.push_to_hub:
         push_model_to_hub(