MCPcopy Create free account
hub / github.com/Open-Quant/openquant / build_index

Function build_index

scripts/afml_semantic_index.py:352–422  ·  view source on GitHub ↗
(args: argparse.Namespace)

Source from the content-addressed store, hash-verified

350
351
352def build_index(args: argparse.Namespace) -> None:
353 pdf_path = Path(args.pdf_path)
354 out_dir = Path(args.out_dir)
355 index_dir = out_dir / "index"
356 pages_dir = out_dir / "pages"
357
358 out_dir.mkdir(parents=True, exist_ok=True)
359 index_dir.mkdir(parents=True, exist_ok=True)
360 pages_dir.mkdir(parents=True, exist_ok=True)
361
362 pages = extract_pages(pdf_path)
363 repeated = detect_repeated_margin_lines(pages)
364 cleaned_pages = clean_pages(pages, repeated)
365 sections = build_sections(cleaned_pages)
366
367 chunks = build_chunks(
368 sections,
369 chunk_size=args.chunk_size,
370 chunk_overlap=args.chunk_overlap,
371 min_chars=args.min_chars,
372 )
373
374 if not chunks:
375 raise RuntimeError("No chunks were produced. Try lowering --min-chars.")
376
377 model = load_sentence_transformer(args.model, args.allow_online_model_fetch)
378 chunk_texts = [c["text"] for c in chunks]
379 embeddings = model.encode(
380 chunk_texts,
381 batch_size=args.batch_size,
382 normalize_embeddings=True,
383 convert_to_numpy=True,
384 show_progress_bar=True,
385 ).astype(np.float32)
386
387 with (index_dir / "chunks.jsonl").open("w", encoding="utf-8") as f:
388 for c in chunks:
389 f.write(json.dumps(c, ensure_ascii=False) + "\n")
390
391 np.save(index_dir / "embeddings.npy", embeddings)
392
393 manifest = {
394 "pdf_path": str(pdf_path),
395 "chunk_count": len(chunks),
396 "embedding_dim": int(embeddings.shape[1]),
397 "model": args.model,
398 "chunk_size": args.chunk_size,
399 "chunk_overlap": args.chunk_overlap,
400 "min_chars": args.min_chars,
401 "pages": len(cleaned_pages),
402 "sections": len(sections),
403 "chapters": chapter_list_from_sections(sections),
404 "repeated_margin_lines_removed": sorted(repeated),
405 "segmentation": "hierarchical(part/chapter/section/snippet)",
406 }
407 (index_dir / "manifest.json").write_text(
408 json.dumps(manifest, indent=2, ensure_ascii=False) + "\n",
409 encoding="utf-8",

Callers 1

mainFunction · 0.85

Calls 7

extract_pagesFunction · 0.85
clean_pagesFunction · 0.85
build_sectionsFunction · 0.85
build_chunksFunction · 0.85

Tested by

no test coverage detected