76 lines
2.3 KiB
Python
76 lines
2.3 KiB
Python
"""Tokenize JSONL files and pack them into HDF5 storage."""
|
|
import json
|
|
import os
|
|
import logging
|
|
from typing import List
|
|
from pathlib import Path
|
|
|
|
from tqdm import tqdm
|
|
|
|
from .processors import BaseProcessor
|
|
from .packing import SequencePacker
|
|
from .io import IOHandler
|
|
from .utils import error_handler
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@error_handler()
|
|
def cache_jsonl(
|
|
files: List[str],
|
|
output_dir: str,
|
|
processor: BaseProcessor,
|
|
*,
|
|
pack_size: int = -1,
|
|
pad_value: int = 1,
|
|
) -> List[str]:
|
|
"""
|
|
Tokenize JSONL files and pack them into HDF5 storage.
|
|
|
|
Args:
|
|
files: List of JSONL file paths
|
|
output_dir: H5 output directory
|
|
processor: Initialized Processor instance
|
|
pack_size: Packing length, <=0 means no packing
|
|
pad_value: Padding value
|
|
|
|
Returns:
|
|
List of generated H5 file paths
|
|
"""
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
output_files: List[str] = []
|
|
|
|
for file_path in files:
|
|
file_name = Path(file_path).stem
|
|
|
|
arrows = []
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
for line_num, line in enumerate(tqdm(f, desc=f"Processing {file_name}", leave=False), start=1):
|
|
try:
|
|
arrow = processor.process(json.loads(line))
|
|
except json.JSONDecodeError as e:
|
|
logger.warning(f"JSON decode error in {file_path} line {line_num}: {e}. Skipping line.")
|
|
continue
|
|
except Exception as e:
|
|
logger.warning(f"Unexpected error processing line {line_num} in {file_path}: {e}. Skipping line.")
|
|
continue
|
|
if arrow is not None:
|
|
arrows.append(arrow)
|
|
|
|
package = {key: [a[key] for a in arrows] for key in processor.output_keys}
|
|
|
|
output = {}
|
|
for key in processor.output_keys:
|
|
if pack_size > 0:
|
|
packer = SequencePacker(pack_size, pad_value) # independent instance per key
|
|
output[key] = packer.pack(package[key])
|
|
else:
|
|
output[key] = package[key]
|
|
|
|
IOHandler.save_h5(output_dir, file_name, output)
|
|
h5_path = os.path.join(output_dir, f"{file_name}.h5")
|
|
output_files.append(h5_path)
|
|
logger.info(f"Saved {h5_path}")
|
|
|
|
return output_files
|