cogcore/datasets/disk_cache.py

"""English implementation note."""
from __future__ import annotations
import gzip
import json
from pathlib import Path
from typing import Iterator, List, Optional, Union
from .loader import StreamingCorpusLoader, BundleStreamingLoader
from .registry import DatasetSpec
DOC_SEPARATOR = '\n\n\n---DOC---\n\n\n'

def cache_to_disk(loader: Union[StreamingCorpusLoader, BundleStreamingLoader], output_path: Union[str, Path], compress: bool=True, overwrite: bool=False, progress_every: int=1000) -> dict:
    """English implementation note."""
    output_path = Path(output_path)
    if compress and output_path.suffix != '.gz':
        output_path = output_path.with_suffix(output_path.suffix + '.gz')
    if output_path.exists() and (not overwrite):
        raise FileExistsError(f'{output_path} already exists. Use overwrite=True.')
    output_path.parent.mkdir(parents=True, exist_ok=True)
    open_fn = gzip.open if compress else open
    n_docs = 0
    total_chars = 0
    with open_fn(output_path, 'wt', encoding='utf-8') as f:
        for doc in loader:
            if n_docs > 0:
                f.write(DOC_SEPARATOR)
            f.write(doc)
            n_docs += 1
            total_chars += len(doc)
            if progress_every > 0 and n_docs % progress_every == 0:
                print(f'  cached {n_docs:,} docs, {total_chars:,} chars', flush=True)
    metadata = {'n_documents': n_docs, 'total_chars': total_chars, 'compress': compress, 'separator': DOC_SEPARATOR, 'output_path': str(output_path)}
    meta_path = output_path.with_suffix(output_path.suffix + '.meta.json')
    with open(meta_path, 'w', encoding='utf-8') as f:
        json.dump(metadata, f, indent=2, ensure_ascii=False)
    return metadata

def iter_cached_corpus(path: Union[str, Path]) -> Iterator[str]:
    """English implementation note."""
    path = Path(path)
    if not path.exists():
        raise FileNotFoundError(f'Cache file not found: {path}')
    open_fn = gzip.open if path.suffix == '.gz' else open
    with open_fn(path, 'rt', encoding='utf-8') as f:
        text = f.read()
    for doc in text.split(DOC_SEPARATOR):
        doc = doc.strip()
        if doc:
            yield doc

def load_cached_corpus(path: Union[str, Path], max_documents: Optional[int]=None) -> List[str]:
    """English implementation note."""
    out: List[str] = []
    for doc in iter_cached_corpus(path):
        out.append(doc)
        if max_documents is not None and len(out) >= max_documents:
            break
    return out