# Copyright 2022 Cerebras Systems.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
This module contains helper functions and classes to read data from different
formats, process them, and save in HDF5 format. It supports JSONL, GZipped JSON,
Parquet, ZST compressed JSONL, and TAR archives of ZST compressed JSONL files.
Classes:
DataFrame:
An object to hold and process data with the ability to serialize itself
into an HDF5 format.
Reader:
Provides a mechanism to read data from multiple file formats, process it,
and yield in manageable chunks.
"""
import gzip
import io
import json
import logging
import math
import numbers
import tarfile
from collections import defaultdict
from contextlib import contextmanager
from typing import Any, Callable, Dict, Iterator, List, Optional
import jsonlines
import pyarrow.parquet as pq
import zstandard
from cerebras.modelzoo.data_preparation.data_preprocessing.vsl_finetuning_token_generator import (
VSLFinetuningTokenGenerator,
)
from cerebras.modelzoo.data_preparation.data_preprocessing.vsl_pretraining_token_generator import (
VSLPretrainingTokenGenerator,
)
logger = logging.getLogger("data_reader")
logger.setLevel(logging.INFO)
[docs]def set_doc_idx(df, file_idx, start_doc_idx, end_doc_idx) -> None:
"""
This is used to set metadata for a given dataframe
Args:
file_idx: The file index of the current dataframe
start_doc_idx: The starting doc index of the current dataframe
end_doc_idx: The ending doc index of the current dataframe
"""
df.file_idx = file_idx
df.start_doc_idx = start_doc_idx
df.end_doc_idx = end_doc_idx
assert (
df.end_doc_idx >= df.start_doc_idx
), "Dataframe's ending document idx must not be less than it's starting document idx"
[docs]def get_data_size(data: Any) -> int:
"""
Compute the size of the given data.
Args:
data (Any): Data whose size needs to be determined.
Returns:
int: Size of the given data in bytes.
"""
if data is None:
return 0
if isinstance(data, str):
return len(data.encode("utf-8"))
elif isinstance(data, bytes):
return len(data)
elif isinstance(data, (tuple, list)):
# Calculate the size by summing up the sizes of individual elements.
return sum(get_data_size(item) for item in data)
elif isinstance(data, dict):
# Calculate the size of the dictionary by summing the sizes of keys and values.
return sum(
get_data_size(key) + get_data_size(value)
for key, value in data.items()
)
elif isinstance(data, int):
# If you wish to support integers as well, represent them in bytes.
return 1 if data == 0 else math.ceil(data.bit_length() / 8)
# Handle other data types like floats or custom objects here if needed.
# If the data type is not supported, you could raise a TypeError.
raise TypeError("Unsupported data type for size calculation")
[docs]@contextmanager
def optional_lock(lock):
if lock:
with lock:
yield
else:
yield
[docs]def find_last_paragraph_or_sentence_end(buffer: str) -> int:
"""
Find the last end of a paragraph (denoted by '\n\n') or a sentence in the buffer.
Args:
buffer (str): The text buffer.
Returns:
int: The position of the last end of the paragraph or sentence.
"""
# Check for the end of a paragraph first
para_end_pos = buffer.rfind('\n\n')
if para_end_pos != -1:
return para_end_pos + 2
# If no end of paragraph is found, check for the end of a sentence
sentence_endings = ['.', '!', '?']
for i in range(len(buffer) - 1, -1, -1):
if buffer[i] in sentence_endings:
return i + 1
return None
[docs]def split_entry_by_paragraph_or_sentence(
entry: str, entry_size: int, chunk_size: int
) -> Iterator[str]:
"""
Split a large entry into chunks by sentence or paragraph end.
Args:
entry (str): The text entry.
entry_size (int): Size of the input entry.
chunk_size (int): The desired chunk size.
Returns:
Iterator[str]: Yields chunks of the text.
"""
while entry_size > chunk_size:
end_pos = find_last_paragraph_or_sentence_end(entry[:chunk_size])
if end_pos:
yield entry[:end_pos]
entry = entry[end_pos:]
entry_size -= end_pos
else:
# Fallback: chunk at chunk_size if no sentence-ending is found
yield entry[:chunk_size]
entry = entry[chunk_size:]
entry_size -= chunk_size
if entry:
yield entry
[docs]class DataFrame:
def __init__(self, keys: Optional[Dict] = None, format_hook_fn=None):
"""
Initialize the DataFrame object.
Args:
keys (Dict): Keys for the data entries.
"""
self.data_keys = keys
self.format_hook_fn = format_hook_fn
self.raw_data = [] ## Stores a list of dictionaries
self.tokenized_data = defaultdict(list)
self.file_idx = None ## stores the file idx from the list of files to which current dataframe belongs
self.start_doc_idx = None ## stores the starting doc index of the current df in the current file
self.end_doc_idx = None ## stores the ending doc index of the current df in the current file
self.size = 0
self.checkpoint_data = []
self.data_stats = defaultdict(int)
[docs] def add(self, value: Dict[str, Any]) -> None:
"""
Add an entry to the DataFrame.
Args:
value (Union[Dict[str, Any], Any]): Entry to be added.
"""
if not isinstance(value, dict):
value = {self.data_keys.values()[0]: value}
self.raw_data.append(self.format_hook_fn(value))
self.size += get_data_size(value)
[docs] def clear(self) -> None:
"""
Clear the raw data after tokenizing.
"""
self.raw_data.clear()
[docs] def tokenize(self, token_generator: Any) -> None:
"""
Tokenize the data values.
Args:
token_generator: Token generator to be used for processing the data.
"""
for doc in self.raw_data:
tokenized_doc, data_stats = token_generator.encode(doc)
for key in data_stats:
self.data_stats[key] += data_stats[key]
if not tokenized_doc:
continue
else:
for key, value in tokenized_doc.items():
self.tokenized_data[key].append(value)
if self.tokenized_data and isinstance(
token_generator,
(VSLPretrainingTokenGenerator, VSLFinetuningTokenGenerator),
):
self.tokenized_data = token_generator.append_within_max_length(
self.tokenized_data
)
self.tokenized_data, data_stats = token_generator.process_chunks(
self.tokenized_data
)
for key in data_stats:
self.data_stats[key] += data_stats[key]
self.clear()
def __repr__(self) -> str:
"""
String representation of the DataFrame object.
Returns:
str: Description of the DataFrame.
"""
output = [f"DataFrame(size={self.size}):"]
for values in self.tokenized_data:
if len(values) >= 6:
# Get the first and last three values
start_values = values[:3]
end_values = values[-3:]
output.append(
f"[{', '.join(map(str, start_values))}, ... , {', '.join(map(str, end_values))}] ({len(values)} entries)"
)
elif values:
# If there are less than 6 values, print them all
output.append(
f"[{', '.join(map(str, values))}] ({len(values)} entries)"
)
else:
output.append(f"[] (0 entries)")
# Adding statistics to the representation
output.append("\nStatistics:")
for stat, value in self.data_stats.items():
output.append(f"{stat}: {value}")
return "\n".join(output)
[docs]class Reader:
def __init__(
self,
file_list: List[str],
max_chunk_size: int,
keys: Optional[Dict],
format_hook_fn: Callable,
) -> None:
"""
Initialize the Reader instance.
Args:
file_list (List[str]): List of file paths to be read.
max_chunk_size (int): Maximum chunk size for accumulated data.
keys (Optional[Dict]): Dictionary containing the type of key and it's name.
"""
self.file_list = file_list
self.max_chunk_size = max_chunk_size
self.keys = keys
self.format_hook_fn = format_hook_fn
[docs] def handle_jsonl(
self,
jsonl_reader: Any,
start_doc_idx: int,
get_meta: bool,
autojoin_paragraphs: bool,
para_joiner: str,
) -> Iterator[Dict[str, Any]]:
"""
Handle JSONL data and yield processed entries.
Args:
jsonl_reader (Any): The JSONL reader object.
start_doc_idx (int): Contains the current document starting index
get_meta (bool): Flag to determine if meta data should be extracted.
autojoin_paragraphs (bool): Flag to auto join paragraphs.
para_joiner (str): Paragraph joiner string.
Returns:
Iterator[Dict[str, Any]]: Yields processed data entries.
"""
for idx, ob in enumerate(
jsonl_reader.iter(type=dict, skip_invalid=True)
):
if (
idx < start_doc_idx
): ## resume streaming data from the current document starting index
continue
if isinstance(ob, str):
assert not get_meta
yield {"text": ob, "doc_idx": idx}
continue
entry = {}
for key, value in self.keys.items():
if value in ob:
text = ob[value]
if not text:
entry[value] = None
continue
# ## Special Case: If the data is an integer typecast it to a string
if isinstance(text, numbers.Number):
text = str(text)
entry[value] = text
else:
entry[value] = None
if get_meta and "meta" in ob:
entry["meta"] = ob["meta"]
entry["doc_idx"] = idx
yield entry
[docs] def accumulate_and_yield(
self, data_gen: Iterator[Dict[str, Any]], file_idx
) -> Iterator[Any]:
"""
Accumulate data and yield in chunks.
Args:
data_gen (Iterator[Dict[str, Any]]): Generator yielding data entries.
file_idx (int): Current file index
Returns:
Iterator[Any]: Yields accumulated data chunks.
"""
df = DataFrame(self.keys, self.format_hook_fn)
start_doc_idx = None
previous_doc_idx = -1
for entry in data_gen:
doc_idx = entry.pop("doc_idx", 0)
entry.pop("meta", None)
if start_doc_idx is None:
start_doc_idx = doc_idx
entry_size = sum(get_data_size(val) for val in entry.values())
# If there's only one key and its size exceeds the chunk size
if len(entry) == 1 and entry_size > self.max_chunk_size:
if df.size > 0:
set_doc_idx(df, file_idx, start_doc_idx, previous_doc_idx)
yield df
df = DataFrame(self.keys, self.format_hook_fn)
start_doc_idx = doc_idx
key = next(iter(entry))
for chunk in split_entry_by_paragraph_or_sentence(
entry[key], entry_size, self.max_chunk_size
):
new_entry = {key: chunk}
df.add(new_entry)
set_doc_idx(df, file_idx, start_doc_idx, start_doc_idx)
yield df
df = DataFrame(self.keys, self.format_hook_fn)
start_doc_idx = None
continue
elif df.size + entry_size > self.max_chunk_size and df.size != 0:
set_doc_idx(df, file_idx, start_doc_idx, previous_doc_idx)
yield df
start_doc_idx = doc_idx
df = DataFrame(self.keys, self.format_hook_fn)
df.add(entry)
previous_doc_idx = doc_idx
if df.size > 0:
set_doc_idx(df, file_idx, start_doc_idx, previous_doc_idx)
assert (
df.end_doc_idx >= df.start_doc_idx
), "Dataframe's ending document idx must not be less than it's starting document idx"
yield df
[docs] def read_txt(self, file: str, checkpoint_args: tuple) -> Iterator[Any]:
"""
Read and process text file.
Args:
file (str): Path to the .txt file.
checkpoint_args (tuple): Contains the current file starting index , current document starting index
Returns:
Iterator[Any]: Yields processed data lines.
"""
current_file_idx, start_doc_idx = checkpoint_args
with open(file, "r") as fh:
text = fh.read()
entry = {self.keys["text_key"]: text}
entry["doc_idx"] = start_doc_idx
df = DataFrame(self.keys, self.format_hook_fn)
set_doc_idx(df, current_file_idx, start_doc_idx, start_doc_idx)
df.add(entry)
yield df
[docs] def read_jsongz(
self,
file: str,
checkpoint_args: tuple,
) -> Iterator[Any]:
"""
Read and process gzipped JSON file.
Args:
file (str): Path to the .json.gz file.
checkpoint_args (tuple): Contains the current file starting index , current document starting index
Returns:
Iterator[Any]: Yields processed data entries.
"""
current_file_idx, start_doc_idx = checkpoint_args
with gzip.open(file, "rb") as f:
text_key = self.keys["text_key"]
data_gen = (
{
text_key: json.loads(line.decode("utf-8").strip())[
text_key
],
"doc_idx": idx,
}
for idx, line in enumerate(f)
if idx >= start_doc_idx
)
yield from self.accumulate_and_yield(data_gen, current_file_idx)
[docs] def read_jsonl(
self,
file: str,
checkpoint_args: tuple,
get_meta: bool = False,
autojoin_paragraphs: bool = True,
para_joiner: str = "\n\n",
) -> Iterator[Any]:
"""
Read and process JSONL file.
Args:
file (str): Path to the .jsonl file.
checkpoint_args (tuple): Contains the current file starting index , current document starting index
get_meta (bool): Flag to determine if meta data should be extracted.
autojoin_paragraphs (bool): Flag to auto join paragraphs.
para_joiner (str): Paragraph joiner string.
Returns:
Iterator[Any]: Yields processed data entries.
"""
current_file_idx, start_doc_idx = checkpoint_args
with open(file, "r", errors='ignore') as fh:
rdr = jsonlines.Reader(fh)
data_gen = self.handle_jsonl(
rdr, start_doc_idx, get_meta, autojoin_paragraphs, para_joiner
)
yield from self.accumulate_and_yield(data_gen, current_file_idx)
[docs] def read_jsonl_zst(
self,
file: str,
checkpoint_args: tuple,
get_meta: bool = False,
autojoin_paragraphs: bool = True,
para_joiner: str = "\n\n",
) -> Iterator[Any]:
"""
Read and process ZST compressed JSONL file.
Args:
file (str): Path to the .jsonl.zst file.
checkpoint_args (tuple): Contains the current file starting index , current document starting index
get_meta (bool): Flag to determine if meta data should be extracted.
autojoin_paragraphs (bool): Flag to auto join paragraphs.
para_joiner (str): Paragraph joiner string.
Returns:
Iterator[Any]: Yields processed data entries.
"""
current_file_idx, start_doc_idx = checkpoint_args
with open(file, "rb") as fh:
cctx = zstandard.ZstdDecompressor()
reader = io.BufferedReader(cctx.stream_reader(fh))
rdr = jsonlines.Reader(reader)
data_gen = self.handle_jsonl(
rdr, start_doc_idx, get_meta, autojoin_paragraphs, para_joiner
)
yield from self.accumulate_and_yield(data_gen, current_file_idx)
[docs] def read_jsonl_tar(
self,
file: str,
checkpoint_args: tuple,
get_meta: bool = False,
autojoin_paragraphs: bool = True,
para_joiner: str = "\n\n",
) -> Iterator[Any]:
"""
Read and process TAR archive containing ZST compressed JSONL files.
Args:
file (str): Path to the .jsonl.zst.tar file.
checkpoint_args (tuple): Contains the current file starting index , current document starting index
get_meta (bool): Flag to determine if meta data should be extracted.
autojoin_paragraphs (bool): Flag to auto join paragraphs.
para_joiner (str): Paragraph joiner string.
Returns:
Iterator[Any]: Yields processed data entries.
"""
current_file_idx, start_doc_idx = checkpoint_args
with tarfile.open(file, "r") as archive:
for member in archive:
with archive.extractfile(member) as f:
cctx = zstandard.ZstdDecompressor()
reader = io.BufferedReader(cctx.stream_reader(f))
rdr = jsonlines.Reader(reader)
data_gen = self.handle_jsonl(
rdr,
start_doc_idx,
get_meta,
autojoin_paragraphs,
para_joiner,
)
yield from self.accumulate_and_yield(
data_gen, current_file_idx
)
[docs] def read_parquet(self, file: str, checkpoint_args: tuple) -> Iterator[Any]:
"""
Read and process Parquet file.
Args:
file (str): Path to the .parquet file.
checkpoint_args (tuple): Contains the current file starting index , current document starting index
Returns:
Iterator[Any]: Yields processed data rows.
"""
current_file_idx, start_doc_idx = checkpoint_args
parquet_file = pq.ParquetFile(file)
def entry_gen(start_doc_idx) -> Iterator[Dict[str, Any]]:
global_doc_idx = 0
for row_group_index in range(parquet_file.num_row_groups):
table = parquet_file.read_row_group(row_group_index)
columns = {
value: table.column(value)
for key, value in self.keys.items()
if value != None
}
for i in range(table.num_rows):
if global_doc_idx < start_doc_idx:
global_doc_idx += 1
continue
else:
entry = {
key: (
str(col[i].as_py())
if isinstance(col[i].as_py(), numbers.Number)
else col[i].as_py()
)
for key, col in columns.items()
}
entry["doc_idx"] = global_doc_idx
yield entry
global_doc_idx += 1
yield from self.accumulate_and_yield(
entry_gen(start_doc_idx), current_file_idx
)
[docs] def read_fasta(
self, file: str, checkpoint_args: tuple
) -> Iterator[Dict[str, Any]]:
"""
Read and process Fasta file without using BioPython.
Args:
file (str): Path to the .fasta file.
checkpoint_args (tuple): Contains the current file starting index, current document starting index
Returns:
Iterator[Dict[str, Any]]: Yields processed data rows.
"""
current_file_idx, start_doc_idx = checkpoint_args
def entry_gen(start_doc_idx):
with open(file, 'r') as fasta_file:
record_id = None
sequence_lines = []
idx = -1 # Initialize sequence index
for line in fasta_file:
line = line.strip()
if not line:
continue # Skip empty lines
if line.startswith(">"):
if record_id is not None and idx >= start_doc_idx:
# Yield the previous record
yield {
"text": ''.join(sequence_lines),
"doc_idx": idx,
}
record_id = line[
1:
] # Remove the ">" symbol and store the record ID
sequence_lines = (
[]
) # Reset the sequence for a new record
idx += 1 # Increment sequence index when a new record is found
else:
sequence_lines.append(line)
# Don't forget to yield the last record in the file
if record_id is not None and idx >= start_doc_idx:
yield {"text": ''.join(sequence_lines), "doc_idx": idx}
yield from self.accumulate_and_yield(
entry_gen(start_doc_idx), current_file_idx
)
[docs] def stream_data(
self, checkpoint_args, get_meta: bool = False
) -> Iterator[Any]:
"""
Stream and process data from multiple file formats.
Args:
get_meta (bool): Flag to determine if meta data should be extracted.
checkpoint_args (tuple): Contains the current file starting index , current document starting index
Returns:
Iterator[Any]: Yields processed data chunks.
"""
file_idx, start_doc_idx = checkpoint_args
zipped_file_list = list(zip(range(len(self.file_list)), self.file_list))
file_list = zipped_file_list[file_idx:]
for idx, f in file_list:
checkpoint_args = (idx, start_doc_idx)
if f.endswith(".jsonl"):
yield from self.read_jsonl(f, checkpoint_args, get_meta)
elif f.endswith(".jsonl.zst"):
yield from self.read_jsonl_zst(f, checkpoint_args, get_meta)
elif f.endswith(".jsonl.zst.tar"):
yield from self.read_jsonl_tar(f, checkpoint_args, get_meta)
elif f.endswith(".txt"):
assert not get_meta
yield from self.read_txt(f, checkpoint_args)
elif f.endswith(".json.gz"):
assert not get_meta
yield from self.read_jsongz(f, checkpoint_args)
elif f.endswith(".parquet"):
assert not get_meta
yield from self.read_parquet(f, checkpoint_args)
elif f.endswith(".fasta"):
assert not get_meta
yield from self.read_fasta(f, checkpoint_args)
else:
logger.warning(
f"Skipping {f} as streaming for that filetype is not implemented"
)