From ccd4741084e2a5b1ba8f7fb1e8d042f55c0390de Mon Sep 17 00:00:00 2001
From: Gautam Kumar <gakumar@fb.com>
Date: Tue, 3 Sep 2024 12:40:06 -0700
Subject: [PATCH] Adding custom dataset file (#659)

---
 src/llama_recipes/datasets/custom_dataset.py | 37 ++++++++++++++++++++
 1 file changed, 37 insertions(+)
 create mode 100644 src/llama_recipes/datasets/custom_dataset.py

diff --git a/src/llama_recipes/datasets/custom_dataset.py b/src/llama_recipes/datasets/custom_dataset.py
new file mode 100644
index 000000000..4bcf0ed6c
--- /dev/null
+++ b/src/llama_recipes/datasets/custom_dataset.py
@@ -0,0 +1,37 @@
+import importlib
+from pathlib import Path
+
+def load_module_from_py_file(py_file: str) -> object:
+    """
+    This method loads a module from a py file which is not in the Python path
+    """
+    module_name = Path(py_file).name
+    loader = importlib.machinery.SourceFileLoader(module_name, py_file)
+    spec = importlib.util.spec_from_loader(module_name, loader)
+    module = importlib.util.module_from_spec(spec)
+
+    loader.exec_module(module)
+
+    return module
+
+
+def get_custom_dataset(dataset_config, tokenizer, split: str):
+    if ":" in dataset_config.file:
+        module_path, func_name = dataset_config.file.split(":")
+    else:
+        module_path, func_name = dataset_config.file, "get_custom_dataset"
+
+    if not module_path.endswith(".py"):
+        raise ValueError(f"Dataset file {module_path} is not a .py file.")
+
+    module_path = Path(module_path)
+    if not module_path.is_file():
+        raise FileNotFoundError(f"Dataset py file {module_path.as_posix()} does not exist or is not a file.")
+
+    module = load_module_from_py_file(module_path.as_posix())
+    try:
+        return getattr(module, func_name)(dataset_config, tokenizer, split)
+    except AttributeError as e:
+        print(f"It seems like the given method name ({func_name}) is not present in the dataset .py file ({module_path.as_posix()}).")
+        raise e
+