File size: 2,480 Bytes
236083b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
import os
import time
import traceback
from pathlib import Path
from litgpt.constants import _LITDATA_AVAILABLE
from litgpt.tokenizer import Tokenizer
from litgpt.utils import CLI, extend_checkpoint_dir
if _LITDATA_AVAILABLE:
from litdata.processing.data_processor import DataChunkRecipe
else:
DataChunkRecipe = object
class StarcoderDataRecipe(DataChunkRecipe):
is_generator = True
def __init__(self, tokenizer: Tokenizer, chunk_size: int):
super().__init__(chunk_size)
self.tokenizer = tokenizer
def prepare_structure(self, input_dir):
files = Path(input_dir).rglob("*.parquet")
return [str(file) for file in files]
def prepare_item(self, item_metadata):
import pyarrow.parquet as pq
filepath = item_metadata
start = time.time()
try:
parquet_file = pq.ParquetFile(filepath)
# reduce RAM usage
for batch in parquet_file.iter_batches(batch_size=8192, columns=["content"]):
for text in batch.to_pandas()["content"]:
yield self.tokenizer.encode(text, bos=False, eos=True)
except Exception:
print(traceback.format_exc())
print(f"Error reading {filepath}")
return
parquet_file.close()
end = time.time()
print(f"Took {end - start:.2f} seconds total", filepath)
def prepare(
input_dir: Path = Path("data/starcoderdata"),
output_dir: Path = Path("data/starcoder"),
tokenizer_path: Path = Path("checkpoints/Llama-2-7b-hf/"),
chunk_size: int = (2049 * 8192),
fast_dev_run: bool = False,
) -> None:
from litdata.processing.data_processor import DataProcessor
tokenizer_path = extend_checkpoint_dir(tokenizer_path)
tokenizer = Tokenizer(tokenizer_path)
data_recipe = StarcoderDataRecipe(tokenizer=tokenizer, chunk_size=chunk_size)
data_processor = DataProcessor(
input_dir=str(input_dir),
output_dir=str(output_dir),
fast_dev_run=fast_dev_run,
num_workers=os.cpu_count(),
num_downloaders=1,
)
start_time = time.time()
data_processor.run(data_recipe)
elapsed_time = time.time() - start_time
print(f"Time taken: {elapsed_time:.2f} seconds")
if __name__ == "__main__":
CLI(prepare)
|