olmocr/olmocr/train/core/compression.py
2025-01-27 18:30:41 +00:00

17 lines
418 B
Python

from smart_open import register_compressor
__all__ = ["mk_compression"]
def mk_compression():
def _handle_zst(file_obj, mode):
try:
import zstandard as zstd
except ImportError:
raise ImportError("zstandard is required for zstd support")
return zstd.open(file_obj, mode)
register_compressor(".zstd", _handle_zst)
register_compressor(".zst", _handle_zst)