DreamOn for code infilling task

src/xlm/commands/lightning_train.py

@@ -18,6 +18,10 @@
 from lightning.pytorch.loggers import Logger
 from lightning import Callback
 from transformers.modeling_utils import no_init_weights
+from xlm.utils.model_loading import (
+    load_model_weights_into_model,
+    _get_model_only_checkpoint_path,
+)
 from xlm.utils.rank_zero import RankedLogger
 
 logger = RankedLogger(__name__, rank_zero_only=True)
@@ -111,21 +115,16 @@ def train(cfg: DictConfig):
                 ckpt_path = os.path.join(cfg.checkpointing_dir, "last.ckpt")
             if not os.path.isfile(ckpt_path):
                 ckpt_path = None
-    # check if we have model only checkpoint
-    model_only_ckpt_path = None
-    if cfg.get("model_only_checkpoint_path", None) is not None:
-        if ckpt_path is not None:
+
+    if ckpt_path is not None:
+        has_model_only = cfg.get("model_only_checkpoint_path", None) is not None
+        has_hub = OmegaConf.select(cfg, "hub.repo_id", default=None) is not None
+        if has_model_only or has_hub:
             logger.error(
-                "model_only_checkpoint_path and resume_from_checkpoint cannot both be provided."
-                " We will use the resume_from_checkpoint path "
-                f"{ckpt_path} for the model weights as well."
+                "Resume checkpoint is set; model-only / Hub weight sources are ignored. "
+                f"Using {ckpt_path} for model weights."
             )
-        else:
-            if not os.path.isfile(cfg.model_only_checkpoint_path):
-                raise ValueError(
-                    f"The model only checkpoint path {cfg.model_only_checkpoint_path} does not exist."
-                )
-            model_only_ckpt_path = cfg.model_only_checkpoint_path
+    model_only_ckpt_path =  _get_model_only_checkpoint_path(cfg, "", True, ckpt_path)
 
     logger.info(f"Instantiating trainer <{cfg.trainer._target_}>")
     trainer = hydra.utils.instantiate(
@@ -169,14 +168,17 @@ def train(cfg: DictConfig):
             _recursive_=False,
         )
     if model_only_ckpt_path is not None:
-        message = lightning_module.model.load_state_dict(
-            torch.load(model_only_ckpt_path)
-        )
         logger.warning(
-            "Loading weights for `model` from a pretrained model at "
-            f"{model_only_ckpt_path} before call to `trainer.fit` => before `.setup()` and `.configure_model()`"
+            "Loading weights for `model` from "
+            f"{model_only_ckpt_path} before `trainer.fit` (before `.setup()` / `.configure_model()`)."
+        )
+        load_model_weights_into_model(
+            lightning_module.model,
+            model_only_ckpt_path,
+            map_location="cpu",
+            strict=True,
+            weights_only=True,
         )
-        logger.warning(message)
 
     # train
     if cfg.job_type == "train":

src/xlm/configs/lightning_train/datasets/humaneval_infill_test.yaml

@@ -0,0 +1,14 @@
+defaults:
+  - default_eval
+
+collator: ???
+load_func: xlm.tasks.humaneval_infill_task.load_func
+load_func_kwargs:
+  benchmark_name: ???
+preprocess_function: xlm.tasks.humaneval_infill_task.humaneval_infill_preprocess_fn
+full_name: null
+columns_to_keep:
+  - canonical_solution
+  - prompt
+  - suffix
+  - task_id

src/xlm/configs/lightning_train/datasets/opencoder_train.yaml

@@ -0,0 +1,24 @@
+_target_: xlm.datamodule.LocalDatasetManager
+ds_type: "parquet"
+collator: ??? # model specific collator
+full_name: opencoder/train
+full_name_debug: opencoder/train
+load_kwargs:
+    data_files: ???
+preprocess_function: xlm.tasks.opencoder.opencoder_preprocess_fn
+preprocess_function_kwargs:
+    prompt_key: prompt
+    response_key: response
+on_the_fly_processor: xlm.datamodule.token_ids_to_input_ids
+on_the_fly_group_processor: null
+stages:
+    - fit
+dataloader_kwargs:
+    batch_size: ${per_device_batch_size} # per device, depends on the device type
+    num_workers: ${num_dataloader_workers} 
+    shuffle: null # can't specify shuffle for IterableDataset
+    pin_memory: True
+    persistent_workers: False # INFO: see: https://github.com/huggingface/datasets/issues/7447
+    prefetch_factor: ${dataloader_prefetch_factor}
+    drop_last: True
+model_name: null

src/xlm/configs/lightning_train/datasets/opencoder_val.yaml

@@ -0,0 +1,25 @@
+_target_: xlm.datamodule.LocalDatasetManager
+ds_type: "parquet"
+collator: ??? # model specific collator
+full_name: opencoder/val
+full_name_debug: opencoder/val
+load_kwargs:
+    data_files: ???
+preprocess_function: xlm.tasks.opencoder.opencoder_preprocess_fn
+preprocess_function_kwargs:
+    prompt_key: prompt
+    response_key: response
+on_the_fly_processor: xlm.datamodule.token_ids_to_input_ids
+on_the_fly_group_processor: null
+stages:
+    - fit
+    - validate
+dataloader_kwargs:
+    batch_size: ${per_device_batch_size} # per device, depends on the device type
+    num_workers: ${num_dataloader_workers} 
+    shuffle: null # can't specify shuffle for IterableDataset
+    pin_memory: True
+    persistent_workers: False # INFO: see: https://github.com/huggingface/datasets/issues/7447
+    prefetch_factor: ${dataloader_prefetch_factor}
+    drop_last: True
+model_name: null

src/xlm/datamodule.py

@@ -609,7 +609,15 @@ def from_txt(cls, txt_file: Union[str, os.PathLike], **kwargs):
                 vocab.append(line.strip())
         return cls(vocab=vocab, **kwargs)
 
-
+class SimpleSpaceTokenizerWithDeletion(SimpleSpaceTokenizer):
+
+    def __init__(self, vocab: Sequence[str], **kwargs):
+        super().__init__(vocab=vocab, **kwargs)
+        del_token_id = len(self._vocab_str_to_int)
+        self._vocab_str_to_int["[DEL]"] = del_token_id
+        self._vocab_int_to_str[del_token_id] = "[DEL]"
+        setattr(self, "delete_token", "[DEL]")
+        setattr(self, "delete_token_id", del_token_id)
 class SimpleSpaceTokenizerWithCyclicPads(SimpleSpaceTokenizer):
     """SimpleSpaceTokenizer with cyclic pad tokens (pad_0..pad_{n-1})."""
 
@@ -1280,6 +1288,21 @@ def _download(self, num_proc: Optional[int] = None) -> datasets.Dataset:
                     num_proc=num_proc,
                 )["train"]
             return ds
+        elif self.ds_type == "parquet":
+            load_kwargs_copy = self.load_kwargs.copy()
+            if "data_files" in load_kwargs_copy:
+                data_files = load_kwargs_copy.pop("data_files")
+            else:
+                file_name = f"{self._split_to_download}.parquet"
+                _path = Path(self.full_name).parent
+                data_files = str(_path / file_name)
+            ds = datasets.load_dataset(
+                "parquet",
+                data_files=data_files,
+                **load_kwargs_copy,
+                num_proc=num_proc,
+            )['train']
+            return ds
         else:
             raise ValueError(f"Unsupported dataset type: {self.ds_type}")
 
@@ -1313,9 +1336,13 @@ def __init__(
         stages: Optional[
             List[Literal["fit", "validate", "test", "predict"]]
         ] = None,
+        load_func: Optional[str] = None,
+        load_func_kwargs: Optional[Dict[str, Any]] = None,
     ):
         self.collator = collator
         self.full_name = full_name
+        self.load_func = load_func
+        self.load_func_kwargs = load_func_kwargs or {}
         self.dataloader_kwargs = dataloader_kwargs
         self.preprocess_function = preprocess_function
         self.preprocess_function_kwargs = preprocess_function_kwargs or {}
@@ -1404,7 +1431,13 @@ def prepare_data(
         logger.info(
             f"EvalDatasetManager: preparing {self.full_name} (no manual cache)"
         )
-        ds = self._download(num_proc=num_proc)
+        if self.load_func:
+            load_fn: Callable[..., Any] = get_function(
+                self.load_func
+            )
+            ds = load_fn(**self.load_func_kwargs)
+        else:
+            ds = self._download(num_proc=num_proc)
         ds = self._preprocess(ds, tokenizer, num_proc=num_proc)
         return ds
 

src/xlm/external_models.py

@@ -27,7 +27,7 @@
 ENV_XLM_MODELS_PATH = "XLM_MODELS_PATH"  # dir containing external models
 ENV_XLM_MODELS_PACKAGES = "XLM_MODELS_PACKAGES"  # installed python packages containing external models, comma separated list of package names
 CORE_XLM_MODELS = (
-    "arlm:mlm:ilm:mdlm:flexmdm"  # core models available in xlm-models package
+    "arlm:mlm:ilm:mdlm:flexmdm:dream:dreamon"  # core models available in xlm-models package
 )