Add the custom model and dataset loaders to prompt debugging notebook.

PiperOrigin-RevId: 676102850
PAIR-code · Sep 18, 2024 · 338c6b1 · 338c6b1
1 parent b87d3d4
commit 338c6b1
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/lit_nlp/examples/prompt_debugging/notebook.py b/lit_nlp/examples/prompt_debugging/notebook.py
@@ -59,5 +59,12 @@ def make_notebook_widget(
       ),
       layouts=layouts.PROMPT_DEBUGGING_LAYOUTS,
       default_layout=layouts.LEFT_RIGHT,
+      model_loaders=models.get_model_loaders(
+          dl_framework=dl_framework,
+          dl_runtime=dl_runtime,
+          batch_size=batch_size,
+          max_length=models.DEFAULT_SEQUENCE_LENGTH,
+      ),
+      dataset_loaders=datasets.get_dataset_loaders(),
       **kwargs,
   )