File size: 3,078 Bytes
262eb2c | 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 59 60 61 62 63 64 65 66 67 68 69 70 | from __future__ import annotations
import csv
import json
from collections.abc import Iterator
from pathlib import Path
SUPPORTED_EXTENSIONS = {".txt", ".md", ".jsonl", ".json", ".csv", ".parquet"}
def discover_files(path: str | Path) -> list[Path]:
root = Path(path).expanduser().resolve()
if root.is_file():
files = [root]
elif root.is_dir():
files = [item for item in root.rglob("*") if item.is_file()]
else:
raise FileNotFoundError(root)
supported = sorted(item for item in files if item.suffix.lower() in SUPPORTED_EXTENSIONS)
if not supported:
raise ValueError(f"No supported data files found under {root}")
return supported
def iter_documents(files: list[Path], text_field: str = "text") -> Iterator[str]:
for path in files:
suffix = path.suffix.lower()
if suffix in {".txt", ".md"}:
text = path.read_text(encoding="utf-8", errors="replace")
if text.strip():
yield text
elif suffix == ".jsonl":
with path.open("r", encoding="utf-8", errors="replace") as handle:
for line_number, line in enumerate(handle, 1):
if not line.strip():
continue
value = json.loads(line)
text = value.get(text_field) if isinstance(value, dict) else None
if isinstance(text, str) and text.strip():
yield text
elif text is None:
raise ValueError(f"Missing string field {text_field!r} in {path}:{line_number}")
elif suffix == ".json":
value = json.loads(path.read_text(encoding="utf-8", errors="replace"))
records = value if isinstance(value, list) else [value]
for index, record in enumerate(records):
text = record.get(text_field) if isinstance(record, dict) else None
if isinstance(text, str) and text.strip():
yield text
elif text is None:
raise ValueError(f"Missing string field {text_field!r} in {path}, record {index}")
elif suffix == ".csv":
with path.open("r", encoding="utf-8", errors="replace", newline="") as handle:
for index, record in enumerate(csv.DictReader(handle), 2):
text = record.get(text_field)
if isinstance(text, str) and text.strip():
yield text
elif text is None:
raise ValueError(f"Missing field {text_field!r} in {path}:{index}")
elif suffix == ".parquet":
try:
import pyarrow.parquet as pq
except ImportError as error:
raise RuntimeError("Parquet scanning requires pyarrow") from error
table = pq.read_table(path, columns=[text_field])
for text in table[text_field].to_pylist():
if isinstance(text, str) and text.strip():
yield text
|