| |
| """Re-ingest papers that failed (chunk_count == 0).""" |
|
|
| import hashlib |
| import json |
| import os |
|
|
| import fitz |
| import httpx |
|
|
| PAPERS_DIR = "/root/backend/data/papers" |
| METADATA_FILE = f"{PAPERS_DIR}/papers_metadata.json" |
| RAG_API = "http://localhost:8000/api/v1/rag/ingest" |
| CHUNK_SIZE = 2500 |
| CHUNK_OVERLAP = 300 |
|
|
|
|
| def main(): |
| with open(METADATA_FILE) as f: |
| metadata = json.load(f) |
|
|
| to_redo = [k for k, v in metadata.items() if not v.get("chunk_count", 0)] |
| print(f"Re-ingesting {len(to_redo)} failed papers") |
|
|
| total_chunks = 0 |
| errors = 0 |
|
|
| for arxiv_id in to_redo: |
| pdf_path = os.path.join(PAPERS_DIR, f"{arxiv_id}.pdf") |
| if not os.path.exists(pdf_path): |
| print(f" SKIP {arxiv_id}: no PDF") |
| continue |
| meta = metadata.get(arxiv_id, {}) |
|
|
| doc = fitz.open(pdf_path) |
| text = "" |
| for page in doc: |
| text += page.get_text() + "\n" |
| doc.close() |
|
|
| if len(text.strip()) < 200: |
| print(f" SKIP {arxiv_id}: too short") |
| continue |
|
|
| chunks = [] |
| start = 0 |
| while start < len(text): |
| end = min(start + CHUNK_SIZE, len(text)) |
| chunk = text[start:end].strip() |
| if len(chunk) > 100: |
| chunks.append(chunk) |
| start += CHUNK_SIZE - CHUNK_OVERLAP |
|
|
| title = meta.get("title", arxiv_id) |
| success = 0 |
| for i, chunk in enumerate(chunks): |
| payload = { |
| "collection": meta.get("collection", "forensic_reports"), |
| "content": chunk, |
| "metadata": { |
| "arxiv_id": arxiv_id, |
| "title": title, |
| "authors": ", ".join(meta.get("authors", [])[:3]) |
| if isinstance(meta.get("authors"), list) |
| else str(meta.get("authors", "")), |
| "attack_types": ", ".join(meta.get("attack_types", [])) |
| if isinstance(meta.get("attack_types"), list) |
| else str(meta.get("attack_types", "")), |
| "chunk_index": i, |
| "total_chunks": len(chunks), |
| "source": "research_paper", |
| "doc_id": hashlib.sha256(f"{arxiv_id}:chunk:{i}".encode()).hexdigest()[:16], |
| }, |
| } |
| try: |
| resp = httpx.post(RAG_API, json=payload, timeout=30) |
| if resp.status_code == 200: |
| success += 1 |
| else: |
| errors += 1 |
| except Exception: |
| errors += 1 |
|
|
| total_chunks += success |
| metadata[arxiv_id]["ingested"] = True |
| metadata[arxiv_id]["chunk_count"] = success |
| print(f" {arxiv_id}: {success}/{len(chunks)} chunks — {title[:50]}") |
|
|
| with open(METADATA_FILE, "w") as f: |
| json.dump(metadata, f, indent=2) |
|
|
| print(f"\nDONE: {total_chunks} chunks, {errors} errors") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
