33 lines
1.1 KiB
Python
33 lines
1.1 KiB
Python
import os
|
|
from typing import List, Dict, Any
|
|
from datetime import datetime
|
|
|
|
import numpy as np
|
|
import pyarrow as pa
|
|
import pyarrow.parquet as pq
|
|
|
|
from config import EMB_DIR, MANIFEST_PATH
|
|
|
|
def _ensure_float32(vectors: List[List[float]]) -> List[List[float]]:
|
|
return [np.asarray(v, dtype=np.float32).tolist() for v in vectors]
|
|
|
|
def write_parquet(
|
|
rows: List[Dict[str, Any]],
|
|
parquet_out_path: str
|
|
) -> None:
|
|
"""Write rows with schema to Parquet (ZSTD)."""
|
|
table = pa.Table.from_pylist(rows)
|
|
pq.write_table(table, parquet_out_path, compression="zstd" )
|
|
print(f"[INFO] Parquet saved: {parquet_out_path}")
|
|
|
|
def append_manifest(entries: List[Dict[str, Any]]) -> None:
|
|
os.makedirs(os.path.dirname(MANIFEST_PATH), exist_ok=True)
|
|
with open(MANIFEST_PATH, "a", encoding="utf-8") as f:
|
|
for e in entries:
|
|
f.write(__to_jsonl_line(e) + "\n")
|
|
print(f"[INFO] Manifest appended: {MANIFEST_PATH}")
|
|
|
|
def __to_jsonl_line(obj: Dict[str, Any]) -> str:
|
|
# Basic compact JSON; avoid pretty for smaller size
|
|
import json
|
|
return json.dumps(obj, ensure_ascii=False, separators=(",", ":")) |