Add datasets to requirements (Lightning-AI#130)

Fristy89 · Jun 9, 2023 · 8bf9664 · 8bf9664
1 parent 675e4ee
commit 8bf9664
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/quantize/gptq.py b/quantize/gptq.py
@@ -9,6 +9,7 @@
 from pathlib import Path
 from typing import Optional
 
+from datasets import load_dataset
 import torch
 
 # support running without installing as a package
@@ -211,7 +212,6 @@ def quantize(self):
 
 
 def get_sample_data():
-    from datasets import load_dataset
 
     traindata = load_dataset(
         "allenai/c4", "allenai--c4", data_files={"train": "en/c4-train.00000-of-01024.json.gz"}, split="train"

diff --git a/requirements.txt b/requirements.txt
@@ -3,4 +3,6 @@ lightning @ git+https://github.com/Lightning-AI/lightning@master
 tokenizers
 jsonargparse[signatures]  # CLI
 bitsandbytes  # quantize
+datasets  # quantize/gptq.py
 zstandard  # prepare_redpajama.py
+