cleaning up after cursor's bumbling

Author: dlwh

src/levanter/data/dataset.py

@@ -123,8 +123,17 @@ def map_batches(self, fn: MapFunction[Sequence[U]], *extra_args, **extra_kwargs)
         return BatchMappedAsyncDataset(self, fn, *extra_args, **extra_kwargs)
 
     def slice_dataset(self, start_index: Optional[int] = None, end_index: Optional[int] = None):
+        """
+        Slices the dataset from `start_index` to `end_index`.
+        """
         return SlicedAsyncDataset(self, start_index, end_index)
 
+    def take(self, n: int):
+        """
+        Alias for `slice_dataset(end_index=n)`.
+        """
+        return self.slice_dataset(end_index=n)
+
     def shuffle(self, key: PRNGKey):
         import levanter.data.permutation as permutation
 

src/levanter/main/eval_lm.py

@@ -47,26 +47,25 @@ def main(config: EvalLmConfig):
 
     Batch = Axis("batch", config.trainer.eval_batch_size)
     Pos = config.model.Pos
-    KeyPos = config.model.KeyPos
 
     if config.eval_on_train:
-        raw_dataset = config.data.train_set(
+        ds = config.data.train_set(
             Pos,
             config.trainer.batch_schedule,
             key=jax.random.PRNGKey(0),
-            KPos=KeyPos,
         )
     else:
-        raw_dataset = config.data.validation_set(Pos, KPos=KeyPos)
+        ds = config.data.validation_set(Pos)  # type: ignore
+        assert ds is not None, "No validation set found"
 
-    if raw_dataset is None:
+    if ds is None:
         raise ValueError("no dataset found!")
 
     if config.max_batches is not None:
-        raw_dataset = raw_dataset.take(config.max_batches * config.trainer.eval_batch_size)
+        ds = ds.take(config.max_batches * config.trainer.eval_batch_size)
 
     eval_loader = DataLoader(
-        raw_dataset,
+        ds,
         Batch,
         max_buffered_batches=None,
         mesh=config.trainer.device_mesh,

src/levanter/main/lora_lm.py

@@ -10,7 +10,7 @@
 import levanter
 from levanter import callbacks
 from levanter.compat.hf_checkpoints import HFCheckpointConverter
-from levanter.data.text import CausalLmDataset, LMDatasetConfig
+from levanter.data.text import LMDatasetConfig
 from levanter.lora import (
     LoraConfig,
     lora_trainable_params_filter,
@@ -80,23 +80,15 @@ def main(config: LoraLmConfig):
         parameter_axis_mapping = config.trainer.parameter_axis_mapping
 
         train_dataset = config.data.train_set(
-            Pos.size,
+            Pos,
             batch_schedule=config.trainer.batch_schedule,
-            QPos=Pos,
-            KPos=KeyPos,
+            key=data_key,
         )
 
         if train_dataset is None:
             raise ValueError("No training set!")
 
-        eval_datasets = {
-            name: config.data.validation_set(
-                Pos.size,
-                QPos=Pos,
-                KPos=KeyPos,
-            )
-            for name in config.data.validation_splits
-        }
+        eval_datasets = config.data.validation_sets(Pos)
 
         if len(eval_datasets) == 0:
             logger.warning("No evaluation datasets provided.")
@@ -135,19 +127,25 @@ def loraize_hf_model(model):
         logger.info(f"Trainable parameter count: {just_lora_params}")
         logger.info(f"Fraction of parameters that are trainable: {just_lora_params * 1.0 / all_param_count:.3e}")
 
-        for name, eval_dataset in eval_datasets.items():
-            eval_dataset = CausalLmDataset(eval_dataset, Pos, KeyPos, ignore_index=config.data.ignore_token_id)
-            trainer.add_eval_hook(eval_dataset, name=name)
+        max_eval_examples_per_ds = config.trainer.max_eval_batches
+        if max_eval_examples_per_ds is not None:
+            max_eval_examples_per_ds *= config.trainer.eval_batch_size
 
-        # boilerplate hooks and such
-        if len(eval_datasets) == 0:
-            logger.warning("No evaluation datasets provided.")
+        tagged_eval_datasets = config.data.tagged_eval_sets(Pos)
 
-        for name, eval_dataset in eval_datasets.items():
-            eval_dataset = CausalLmDataset(
-                eval_dataset, Pos, KeyPos, ignore_index=config.data.ignore_token_id, eos_id=tokenizer.eos_token_id
+        if len(tagged_eval_datasets) == 0:
+            logger.warning("No evaluation datasets provided.")
+        else:
+            cb = levanter.eval.cb_tagged_lm_evaluate(
+                trainer.EvalBatch,
+                tagged_eval_datasets,
+                tokenizer,
+                trainer.device_mesh,
+                trainer.compute_axis_mapping,
+                max_eval_examples_per_ds,
+                mp=config.trainer.mp,
             )
-            trainer.add_eval_hook(eval_dataset, name=name)
+            trainer.add_hook(cb, every=config.trainer.steps_per_eval)
 
         trainer.add_hook(callbacks.log_performance_stats(Pos.size, trainer.config.train_batch_size), every=1)
         if config.peft_save_path is not None:

src/levanter/main/viz_logprobs.py

@@ -53,9 +53,8 @@ def main(config: VizLmConfig):
     # some axes we use outside the model proper
     EvalBatch = config.trainer.EvalBatch
     Pos = config.model.Pos
-    KeyPos = config.model.KeyPos
 
-    validation_sets = config.data.validation_sets(Pos, KPos=KeyPos)
+    validation_sets = config.data.validation_sets(Pos)
 
     # some axes we use outside the model proper
     Pos = config.model.Pos

tests/test_text.py

@@ -20,9 +20,8 @@ def test_dont_blow_up_without_validation_set():
         )
 
         Pos = hax.Axis("Pos", 10)
-        KPos = hax.Axis("KPos", 10)
         # mostly just making sure this doesn't blow up
-        assert config.validation_set(Pos, KPos=KPos) is None
+        assert config.validation_set(Pos) is None
 
 
 def test_lm_example_handles_ignore_id():