Spaces:
Running
Running
| """Join manifest, dense vectors, and sparse vectors into Pinecone documents.""" | |
| from pathlib import Path | |
| import json | |
| import pandas as pd | |
| PROCESSED_DIR = Path("data/comics/processed") | |
| MANIFEST_PATH = PROCESSED_DIR / "panels_manifest.parquet" | |
| DENSE_PATH = PROCESSED_DIR / "image_dense_vectors.parquet" | |
| SPARSE_PATH = PROCESSED_DIR / "text_sparse_vectors.parquet" | |
| OUTPUT_PATH = PROCESSED_DIR / "pinecone_documents.jsonl" | |
| def to_list(val) -> list: | |
| """Convert numpy array or other sequence to a plain Python list.""" | |
| import numpy as np | |
| if isinstance(val, np.ndarray): | |
| return val.tolist() | |
| return list(val) | |
| def valid_sparse(sparse) -> bool: | |
| if not isinstance(sparse, dict): | |
| return False | |
| indices = sparse.get("indices") | |
| values = sparse.get("values") | |
| return indices is not None and len(indices) > 0 and values is not None and len(values) > 0 | |
| def main(): | |
| manifest = pd.read_parquet(MANIFEST_PATH) | |
| dense = pd.read_parquet(DENSE_PATH) | |
| sparse = pd.read_parquet(SPARSE_PATH) | |
| # Inner join on dense — only panels that were successfully embedded | |
| df = manifest.merge(dense, on="panel_id", how="inner") | |
| # Left join on sparse — panels without text keep their document slot | |
| df = df.merge(sparse, on="panel_id", how="left") | |
| dropped = len(manifest) - len(df) | |
| if dropped: | |
| print(f"Warning: {dropped} panels dropped (missing dense vector)") | |
| count = 0 | |
| with open(OUTPUT_PATH, "w") as f: | |
| for row in df.itertuples(index=False): | |
| doc: dict = { | |
| "_id": row.panel_id, | |
| "image_dense": to_list(row.image_dense), | |
| "ocr_text": row.ocr_text_clean or "", | |
| "search_text": row.search_text or "", | |
| "comic_id": row.comic_id, | |
| "book_id": row.book_id, | |
| "page_id": row.page_id, | |
| "page_num": int(row.page_num) if row.page_num is not None and not pd.isna(row.page_num) else None, | |
| "panel_num": int(row.panel_num) if row.panel_num is not None and not pd.isna(row.panel_num) else None, | |
| "image_path": row.image_path, | |
| "source": "COMICS", | |
| "is_ad_page": bool(row.is_ad_page), | |
| } | |
| sparse_vec = getattr(row, "text_sparse", None) | |
| if valid_sparse(sparse_vec): | |
| doc["text_sparse"] = { | |
| "indices": to_list(sparse_vec["indices"]), | |
| "values": to_list(sparse_vec["values"]), | |
| } | |
| f.write(json.dumps(doc) + "\n") | |
| count += 1 | |
| print(f"Wrote {count} documents → {OUTPUT_PATH}") | |
| with_sparse = sum(1 for row in df.itertuples(index=False) if valid_sparse(getattr(row, "text_sparse", None))) | |
| print(f" With sparse vector: {with_sparse}") | |
| print(f" Dense only: {count - with_sparse}") | |
| if __name__ == "__main__": | |
| main() | |