Spaces:
Paused
Paused
File size: 2,506 Bytes
0d20347 | 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 | """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
|