structured-data-extraction 0.1.0__py3-none-any.whl

structured_data_extraction/__init__.py

@@ -0,0 +1,8 @@
+"""
+Loads environment variables from the nearest .env file on import so that
+config.py and all other modules see the user's settings immediately.
+"""
+from dotenv import load_dotenv, find_dotenv
+
+
+load_dotenv(find_dotenv(usecwd=True))

structured_data_extraction/classify.py

@@ -0,0 +1,148 @@
+"""Section classifier : re-labels Docling items with a supervised RandomForest.
+
+Docling's built-in labels are unreliable on dense datasheets: multi-sentence
+captions get tagged as "text", and sub-headers are often labelled "list_item".
+A small RandomForest (trained on hand-labelled datasheet items) predicts a corrected section type for every item.
+
+The classifier uses TF-IDF features (word- and character-level), 30 hand-crafted
+layout/typography features from features.py, and one-hot context from the
+surrounding Docling labels. Its output — sections.json — is the canonical
+content stream that every downstream step (embedding, search, extraction) reads.
+
+The trained model (~89 MB: one joblib per vectoriser plus the forest) ships
+inside the package under models/section_classifier/.
+"""
+import argparse
+import json
+import logging
+from collections import Counter
+from pathlib import Path
+import joblib
+from . import config
+from .features import build_features, context_rows
+
+logger = logging.getLogger(__name__)
+
+# Section classifier is a random forest model trained on labeled sections from the Docling dataset.
+# It uses TF-IDF features of the text and character n-grams, as well as the Docling label of each section, 
+# to predict the section type (table, text, list, etc)
+class SectionClassifier:
+
+    def __init__(self, model_dir=None):
+        d = Path(model_dir) if model_dir else Path(config.SECTION_CLASSIFIER_DIR)
+        self.rf = joblib.load(d / "rf_classifier.joblib")
+        self.word_vec = joblib.load(d / "tfidf_vectorizer.joblib")
+        self.char_vec = joblib.load(d / "tfidf_char_vectorizer.joblib")
+        self.dl_cats = joblib.load(d / "docling_cats.joblib")
+        self.dl_idx = {c: i for i, c in enumerate(self.dl_cats)}
+
+    def classify(self, chunks):
+        """
+        Predict a section type (table, text, list, ...) for each chunk.
+
+        Each chunk is a dict with "section_id", "docling_label", and "text";
+        "page" and "bbox" are passed through if present.
+        Returns the chunks sorted by section_id, each with "section_type" added.
+        """
+        if not chunks:
+            return []
+        ordered, rows = context_rows(chunks, text_key="text")
+        X = build_features(rows, self.dl_idx, self.word_vec, self.char_vec, fit=False)
+        preds = self.rf.predict(X)
+        return [dict(c, section_type=str(p)) for c, p in zip(ordered, preds)]
+
+
+_classifier = None
+
+
+def get_classifier():
+    """Load the ~89 MB random forest + its vectorizers once and reuse across PDFs."""
+    global _classifier
+    if _classifier is None:
+        _classifier = SectionClassifier()
+    return _classifier
+
+
+def classify_sections(pdf_stem):
+    """
+    Disk wrapper around SectionClassifier.classify for the batch pipeline.
+    Reads {stem}_items.json, runs the items with text through classify(),
+    and writes {stem}_sections.json (skips if it already exists).
+
+    Returns the list of classified sections — each a dict with
+    "type", "content", "docling_label", "page", "bbox", and "section_idx".
+    """
+    items_path = Path(config.INTERMEDIATE_DIR) / f"{pdf_stem}_items.json"
+    out_path = Path(config.INTERMEDIATE_DIR) / f"{pdf_stem}_sections.json"
+
+    if not items_path.exists():
+        raise FileNotFoundError(f"Items not found: {items_path}. Run convert.py first.")
+
+    if out_path.exists() and not config.FORCE:
+        logger.info(f"Skipping {pdf_stem} — sections already exist")
+        return json.loads(out_path.read_text(encoding="utf-8"))
+
+    items = json.loads(items_path.read_text(encoding="utf-8"))
+    logger.info(f"Classifying {len(items)} items for {pdf_stem}...")
+
+    model_dir = Path(config.SECTION_CLASSIFIER_DIR)
+    if not model_dir.exists():
+        raise FileNotFoundError(f"Classifier not found at {model_dir}.")
+
+    # only classify items that have text
+    text_chunks = [
+        {
+            "section_id": i,
+            "docling_label": item.get("label", "text"),
+            "text": text,
+            "page": item.get("page"),
+            "bbox": item.get("bbox"),
+        }
+        for i, item in enumerate(items)
+        if (text := (item.get("text") or "").strip())
+    ]
+
+    predicted = get_classifier().classify(text_chunks)
+
+    # build output keeping original item order
+    by_idx = {chunk["section_id"]: chunk for chunk in predicted}
+    sections = []
+    for i, item in enumerate(items):
+        if i in by_idx:
+            chunk = by_idx[i]
+            sections.append({
+                "type": chunk["section_type"],
+                "content": chunk["text"],
+                "docling_label": chunk["docling_label"],
+                "page": chunk["page"],
+                "bbox": chunk["bbox"],
+                "section_idx": i,
+            })
+        elif item.get("label") == "picture":
+            sections.append({
+                "type": "picture",
+                "content": "",
+                "docling_label": item.get("label"),
+                "page": item.get("page"),
+                "bbox": item.get("bbox"),
+                "section_idx": i,
+            })
+
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    out_path.write_text(json.dumps(sections, indent=2, ensure_ascii=False), encoding="utf-8")
+    logger.info(f"Classified {len(sections)} sections → {out_path}")
+    return sections
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Classify sections")
+    parser.add_argument("--stem", required=True)
+    args = parser.parse_args()
+
+    result = classify_sections(args.stem)
+    counts = Counter(s["type"] for s in result)
+    print(f"Classified {len(result)} sections: {dict(counts)}")
+
+
+if __name__ == "__main__":
+    main()

structured_data_extraction/config.py

@@ -0,0 +1,62 @@
+"""Global configuration: directory layout, device selection, and tuning knobs.
+
+Directory paths default to subdirectories of the
+working directory (pdf/, markdown/, intermediate/, output/).
+
+The embedding device is auto-detected: CUDA if available, then MPS (Apple
+Silicon), then CPU. Override with EMBEDDING_DEVICE in .env.
+
+Tuning constants (similarity threshold, search top-K, vision complexity
+threshold, etc.) have sensible defaults but can be overridden per-run in .env.
+"""
+import os
+from pathlib import Path
+from importlib import resources
+
+
+
+CWD = Path.cwd()
+PDF_DIR = CWD / "pdf"
+MARKDOWN_DIR = CWD / "markdown"
+INTERMEDIATE_DIR = CWD / "intermediate"
+OUTPUT_DIR = CWD / "output"
+
+
+SECTION_CLASSIFIER_DIR = resources.files("structured_data_extraction") / "models" / "section_classifier"
+
+
+def _auto_device():
+    import torch
+    if torch.cuda.is_available():
+        return "cuda"
+    if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+
+EMBEDDING_DEVICE = os.environ.get("EMBEDDING_DEVICE") or _auto_device()
+
+
+# The models to run for extraction, one or many 
+# The first one is the "primary": used wherever a single model is needed 
+EXTRACTION_MODELS = [m.strip() for m in os.environ["EXTRACTION_MODELS"].split(",") if m.strip()]
+assert EXTRACTION_MODELS, "EXTRACTION_MODELS in .env must list at least one model"
+
+# When True (set by main.py --force), every cached pipeline step recomputes instead of skipping.
+FORCE = False
+
+# Section-type preference when deduping records that describe the same fact: table beats text beats list.
+SECTION_TYPE_PRIORITY = {"table": 0, "text": 1, "list": 2}
+
+# Tuning constants — overridable in .env; the values here are the defaults if unset.
+def _env_bool(name, default):
+    v = os.environ.get(name)
+    return default if v is None else v.strip().lower() in ("1", "true", "yes", "on")
+
+VOCAB_TOP_N = int(os.environ.get("VOCAB_TOP_N", 50))
+VOCAB_MAX_NGRAM = int(os.environ.get("VOCAB_MAX_NGRAM", 4))
+SIMILARITY_THRESHOLD = float(os.environ.get("SIMILARITY_THRESHOLD", 0.55))
+SEARCH_TOP_K = int(os.environ.get("SEARCH_TOP_K", 25))
+VISION_COMPLEXITY_THRESHOLD = float(os.environ.get("VISION_COMPLEXITY_THRESHOLD", 0.15))
+VISION_PARALLELISM = int(os.environ.get("VISION_PARALLELISM", 4))
+EXPAND_KEYS = _env_bool("EXPAND_KEYS", True)

structured_data_extraction/convert.py

@@ -0,0 +1,165 @@
+"""PDF-to-structured-items conversion via Docling.
+
+Converts a PDF into two artifacts: a full-document Markdown file and a
+per-item JSON list (_items.json) where each item carries its Docling label,
+text, page number, and bounding box.
+
+OCR is disabled because the target documents are born-digital; Docling's table
+structure model runs in FAST mode (roughly 2x faster than ACCURATE) because the
+downstream vision-refinement step in tables.py re-reads complex tables anyway,
+so the speed/accuracy trade-off favours throughput here.
+
+Bounding boxes are stored in Docling's native BOTTOMLEFT coordinate origin.
+Any consumer that needs top-left origin (e.g. PyMuPDF rendering in tables.py)
+converts on read.
+"""
+import argparse
+import json
+import logging
+import re
+from pathlib import Path
+from docling.document_converter import DocumentConverter, PdfFormatOption
+from docling.datamodel.base_models import InputFormat
+from docling.datamodel.pipeline_options import PdfPipelineOptions, TableFormerMode
+from docling.backend.pypdfium2_backend import PyPdfiumDocumentBackend
+from . import config
+
+logger = logging.getLogger(__name__)
+
+_converter = None
+
+
+def _get_converter():
+    """Build the Docling converter on first use — it loads vision models, so importing
+    this module shouldn't pay that cost."""
+    global _converter
+    if _converter is None:
+        #Custom Docling pipeline configuration for PDF conversion
+        pipeline_opts = PdfPipelineOptions()
+        pipeline_opts.do_ocr = False
+        pipeline_opts.do_table_structure = True
+        pipeline_opts.table_structure_options.mode = TableFormerMode.FAST #reduced runtime by 2x compared to accurate mode,similar output quality
+                                                                          #since we pay the price of vision refinement anyways at least we have speed
+        _converter = DocumentConverter(
+            format_options={
+                InputFormat.PDF: PdfFormatOption(
+                    pipeline_options=pipeline_opts,
+                    backend=PyPdfiumDocumentBackend, # since we don't need OCR, we can use the faster PyPdfium backend instead of Tesseract
+                )
+            }
+        )
+    return _converter
+
+
+_CTRL_RE = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f]") # control characters that break JSON and Markdown rendering
+
+
+def _scrub(text):
+    """Remove control characters from text, which can break JSON and Markdown rendering."""
+    return _CTRL_RE.sub("", text) if text else text
+
+
+def _get_item_text(item, label, doc):
+    """Turn one docling item into Markdown text.
+
+    Docling's whole-doc export gives one big string, we need the text for
+    each item on its own so it stays lined up with that item's label, page, and
+    bbox. So process each item by type: tables go through docling, list items
+    get their marker in front, a title gets one '#' and a section_header gets
+    '#' x its docling level, everything else is just its plain text. Gives back
+    "" when an item has no text.
+
+    The header depth feeds the section classifier (features.regex_features reads it
+    back off the '#' prefix), so keep emitting the level even though the markdown
+    export wouldn't strictly need it.
+    """
+    # tables store content differently, need export_to_markdown
+    if label == "table":
+        try:
+            return item.export_to_markdown(doc)
+        except TypeError:
+            return item.export_to_markdown()  # older docling: no doc arg
+
+    raw = getattr(item, "text", "") or ""
+    if label == "list_item":
+        marker = getattr(item, "marker", "") or ""
+        return f"{marker} {raw}".strip() if raw else ""
+    if label == "section_header":
+        lvl = getattr(item, "level", None)
+        depth = min(max(int(lvl) if lvl else 2, 1), 6)  # docling level, default 2, clamped to 1..6
+        return ("#" * depth + " " + raw) if raw else ""
+    if label == "title":
+        return ("# " + raw) if raw else ""
+    return raw
+
+
+def _get_bbox(item):
+    if not hasattr(item, 'prov') or not item.prov:
+        return None, None
+    prov = item.prov[0]
+    page = prov.page_no
+    bbox = None
+    if hasattr(prov, 'bbox') and prov.bbox:
+        bbox = {
+            "l": prov.bbox.l, "t": prov.bbox.t,
+            "r": prov.bbox.r, "b": prov.bbox.b,
+            "coord_origin": str(prov.bbox.coord_origin),
+        }
+    return page, bbox
+
+
+def convert_pdf(pdf_path):
+    pdf_path = Path(pdf_path)
+    if not pdf_path.exists():
+        raise FileNotFoundError(f"PDF not found: {pdf_path}")
+
+    out_md = Path(config.MARKDOWN_DIR) / f"{pdf_path.stem}.md"
+    out_items = Path(config.INTERMEDIATE_DIR) / f"{pdf_path.stem}_items.json"
+    out_md.parent.mkdir(parents=True, exist_ok=True)
+    out_items.parent.mkdir(parents=True, exist_ok=True)
+
+    if out_md.exists() and out_items.exists() and not config.FORCE:
+        logger.info(f"Skipping {pdf_path.name} — already exists")
+        return str(out_md)
+
+    logger.info(f"Converting {pdf_path.name}...")
+    try:
+        result = _get_converter().convert(str(pdf_path))
+        doc = result.document
+
+        md_text = _scrub(doc.export_to_markdown())
+        out_md.write_text(md_text, encoding="utf-8")
+
+        # collect items from docling
+        items = []
+        for item, _ in doc.iterate_items():
+            label = str(item.label) if hasattr(item, 'label') else "unknown"
+            text = _scrub(_get_item_text(item, label, doc))
+            page, bbox = _get_bbox(item)
+            items.append({"label": label, "text": text, "page": page, "bbox": bbox})
+
+        out_items.write_text(json.dumps(items, indent=2, ensure_ascii=False), encoding="utf-8")
+
+     
+        logger.info(f"Saved markdown ({len(md_text)} chars) + {len(items)} items → {out_md}")
+        return str(out_md)
+
+    except Exception as e:
+        logger.error(f"Failed to convert {pdf_path.name}: {e}")
+        return None
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Convert PDF to Markdown")
+    parser.add_argument("--input", required=True, help="Path to PDF file")
+    args = parser.parse_args()
+
+    result = convert_pdf(args.input)
+    if result:
+        print(f"Done: {result}")
+    else:
+        print("Conversion failed.")
+
+
+if __name__ == "__main__":
+    main()

structured_data_extraction/embed.py

@@ -0,0 +1,270 @@
+"""Embedding of document sections, table rows, and user keys.
+
+Encodes every section and individual table row into dense vectors with a
+SentenceTransformer model (default: nomic-embed-text-v1.5), then encodes the
+user's extraction keys into the same space so search.py can match them via
+cosine similarity.
+
+nomic requires task-type prefixes: sections are prefixed with "search_document:"
+and keys with "search_query:". Without these prefixes similarity scores collapse
+and SIMILARITY_THRESHOLD stops discriminating relevant from irrelevant hits.
+
+Sections are stacked before rows in the embedding matrix, search.py relies on
+that ordering to tell section-level hits from row-level hits using num_sections
+as the boundary index.
+
+Table rows are parsed into per-row text (with group headers prepended) so that
+search can match an individual parameter row inside a large table rather than
+only matching the table as a whole.
+"""
+import argparse
+import json
+import logging
+import os
+import re
+from pathlib import Path
+import numpy as np
+from sentence_transformers import SentenceTransformer
+from . import config
+from .paths import inter, read_json, write_json
+
+logger = logging.getLogger(__name__)
+
+
+EMBED_BATCH_SIZE = int(os.environ.get("EMBED_BATCH_SIZE", 4))
+
+_model_cache = {}
+
+
+def _load_model(model_name=None):
+    """Load the embedding model (from the given name or env), caching it so the hundreds-of-MB
+    SentenceTransformer is loaded once and reused across PDFs."""
+    name = model_name or os.environ["EMBEDDING_MODEL"]
+    if name not in _model_cache:
+        logger.info(f"Loading embedding model: {name} on {config.EMBEDDING_DEVICE}")
+        _model_cache[name] = SentenceTransformer(name, device=config.EMBEDDING_DEVICE, trust_remote_code=True)
+    return _model_cache[name]
+
+
+def _expand_key(key, vocab):
+    """Return a tuple of (expanded_key, original_key). If EXPAND_KEYS is False, returns (key, key)."""
+    if not config.EXPAND_KEYS:
+        return key, key
+
+    key_lower = key.lower()
+    key_words = key_lower.split()
+    expansions = []
+    for item in vocab:
+        term_lower = item["term"].lower()
+        term_words = term_lower.split()
+        if key_lower in term_lower or all(word in term_words for word in key_words):
+            expansions.append(item["term"])
+
+    if expansions:
+        return " ".join([key, " ".join(expansions[:3])]), key
+    return key, key
+
+
+def _is_group_header(cells):
+    """Determine if a row is likely a group header based on its content."""
+    if len(cells) <= 1 or not cells[0]:
+        return False
+    empty_count = sum(1 for cell in cells[1:] if not cell)
+    standalone_nums = sum(
+        1 for cell in cells[1:]
+        if cell and re.match(r'^[\±\-\+]?\d+\.?\d*\s*[a-zA-Zµ°/²³%]*$', cell.strip()) # matches standalone numbers with optional units
+    )
+    has_equals = '=' in cells[0]
+    return empty_count >= len(cells[1:]) * 0.6 and standalone_nums < 2 and not has_equals # if most cells are empty, few standalone numbers,
+                                                                                        # and no equals sign in the first cell, it's likely a group header
+
+
+def _parse_html_rows(content):
+    """Parse table rows from HTML content, return a list of row dictionaries."""
+    rows = []
+    for n, tr in enumerate(re.findall(r"<tr\b[^>]*>(.*?)</tr>", content, re.S | re.I)): # find all <tr>...</tr> blocks
+        cells = [re.sub(r"<[^>]+>", "", c).strip()
+                 for c in re.findall(r"<t[dh]\b[^>]*>(.*?)</t[dh]>", tr, re.S | re.I)] # find all <td> or <th> blocks within the <tr>
+        dedup = []
+        for c in cells:                       # collapse repeats from rowspan/colspan expansion
+            if not dedup or dedup[-1] != c:
+                dedup.append(c)
+        row_text = " | ".join(c for c in dedup if c)
+        if row_text.strip():
+            rows.append({"row_number": n, "row_text": row_text, "group_header": ""})
+    return rows
+
+
+def _parse_table_rows(content):
+    """Parse table rows from markdown content, return a list of row dictionaries."""
+    if "<tr" in content.lower():              # pre-embed refined tables are HTML, not markdown pipes
+        return _parse_html_rows(content)
+    lines = content.strip().split('\n')
+    rows = []
+    group_header = ""
+    col_headers = []
+    row_num = 0
+
+    for line in lines:
+        stripped = line.strip()
+        if not stripped or not stripped.startswith('|'):
+            continue
+        if re.match(r'^\|[\s\-:|]+\|$', stripped):
+            continue
+
+        cells = [cell.strip() for cell in stripped.split('|')]
+        if cells and cells[0] == '':
+            cells = cells[1:]
+        if cells and cells[-1] == '':
+            cells = cells[:-1]
+        if not cells:
+            continue
+
+        
+        non_empty = [cell for cell in cells if cell]
+        if non_empty and all(cell in ('Min', 'Typ', 'Max') for cell in non_empty) and len(non_empty) >= 3:
+            col_headers = cells
+            continue
+
+        if _is_group_header(cells):
+            group_header = cells[0]
+            row_text = ' | '.join(cell for cell in cells if cell)
+            if row_text.strip():
+                rows.append({"row_number": row_num, "row_text": row_text, "group_header": ""})
+                row_num += 1
+            continue
+
+        if col_headers and len(col_headers) == len(cells):
+            annotated = []
+            for i, cell in enumerate(cells):
+                if cell and col_headers[i]:
+                    annotated.append(f"[{col_headers[i]}] {cell}")
+                else:
+                    annotated.append(cell)
+            row_text = ' | '.join(annotated)
+        else:
+            row_text = ' | '.join(cells)
+
+        if row_text.strip():
+            rows.append({"row_number": row_num, "row_text": row_text, "group_header": group_header})
+            row_num += 1
+
+    return rows
+
+
+def _build_section_texts(sections, tables):
+    """Build per-section text list, replacing table sections with their cleaned markdown content."""
+    sec_texts = [f"{section['type']}: {section['content']}" for section in sections]
+    table_by_idx = {table["section_idx"]: table for table in tables}
+    for i, section in enumerate(sections):
+        if section.get("type") == "table" and i in table_by_idx:
+            cleaned = table_by_idx[i].get("content", "")
+            if cleaned:
+                sec_texts[i] = f"table: {cleaned}"
+    return sec_texts, table_by_idx
+
+
+def _parse_all_rows(sections, table_by_idx):
+    """Parse table rows from all table sections, return (row_texts, row_map)."""
+    row_texts = []
+    row_map = []
+    for i, section in enumerate(sections):
+        if section.get("type") != "table" or i not in table_by_idx:
+            continue
+        cleaned = table_by_idx[i].get("content", "")
+        if not cleaned:
+            continue
+        for row_info in _parse_table_rows(cleaned):
+            txt = row_info["row_text"]
+            if row_info["group_header"]:
+                txt = f"{row_info['group_header']} > {txt}"
+            row_texts.append(f"table_row: {txt}")
+            row_map.append({
+                "section_idx": i,
+                "row_number": row_info["row_number"],
+                "row_text": row_info["row_text"],
+                "group_header": row_info["group_header"],
+            })
+    return row_texts, row_map
+
+
+def _expand_all_keys(user_keys, vocab):
+    expansions = [_expand_key(key, vocab) for key in user_keys]
+    expanded = [pair[0] for pair in expansions]
+    originals = [pair[1] for pair in expansions]
+    return expanded, originals
+
+
+def _save_outputs(stem, out_dir, all_embs, key_embs, all_texts, expanded, originals, row_map,
+                  num_sections):
+    """Write the embedding + text artifacts for this stem."""
+    np.save(str(out_dir / f"{stem}_section_embeddings.npy"), all_embs)
+    np.save(str(out_dir / f"{stem}_key_embeddings.npy"), key_embs)
+    write_json(out_dir / f"{stem}_section_texts.json", all_texts)
+    write_json(out_dir / f"{stem}_key_texts.json", expanded)
+    write_json(out_dir / f"{stem}_original_keys.json", originals)
+    write_json(out_dir / f"{stem}_row_map.json", {"num_sections": num_sections, "rows": row_map})
+
+
+def embed_sections_and_keys(stem, user_keys=None, model_name=None):
+    out_dir = Path(config.INTERMEDIATE_DIR)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    sec_names = ("section_embeddings.npy", "section_texts.json", "row_map.json")
+    key_names = ("key_embeddings.npy", "key_texts.json")
+    sec_cached = all(inter(stem, n).exists() for n in sec_names) and not config.FORCE
+    key_cached = all(inter(stem, n).exists() for n in key_names) and not config.FORCE
+
+    if sec_cached and key_cached:
+        logger.info(f"Skipping embeddings — already cached for {stem}")
+        return
+
+    sections = read_json(inter(stem, "sections.json"))
+    tables = read_json(inter(stem, "tables.json")) if inter(stem, "tables.json").exists() else []
+    vocab = read_json(inter(stem, "vocab.json")) if inter(stem, "vocab.json").exists() else []
+
+    if user_keys is None:
+        user_keys = []
+
+    model = _load_model(model_name)
+
+    if sec_cached:
+        logger.info(f"Reusing cached section embeddings for {stem}")
+        all_embs = np.load(str(inter(stem, "section_embeddings.npy")))
+        all_texts = read_json(inter(stem, "section_texts.json"))
+        rm = read_json(inter(stem, "row_map.json"))
+        num_sections = rm.get("num_sections", len(sections))
+        row_map = rm.get("rows", [])
+    else:
+        sec_texts, table_by_idx = _build_section_texts(sections, tables)
+        row_texts, row_map = _parse_all_rows(sections, table_by_idx)
+        logger.info(f"Parsed {len(row_map)} table rows")
+        all_texts = sec_texts + row_texts
+        num_sections = len(sec_texts)
+        logger.info(f"Embedding {len(sec_texts)} sections + {len(row_texts)} rows...")
+        all_embs = model.encode([f"search_document: {t}" for t in all_texts],
+                                show_progress_bar=False, batch_size=EMBED_BATCH_SIZE)
+
+    expanded, originals = _expand_all_keys(user_keys, vocab)
+    logger.info(f"Embedding {len(expanded)} keys...")
+    key_embs = (model.encode([f"search_query: {k}" for k in expanded],
+                             show_progress_bar=False, batch_size=EMBED_BATCH_SIZE)
+                if expanded else np.array([]))
+
+    _save_outputs(stem, out_dir, all_embs, key_embs, all_texts, expanded, originals,
+                  row_map, num_sections)
+    logger.info(f"Saved embeddings: {all_embs.shape[0]} content ({num_sections} sections + {len(row_map)} rows), {len(expanded)} keys")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Embed sections and keys")
+    parser.add_argument("--stem", required=True)
+    parser.add_argument("--keys", nargs="+")
+    parser.add_argument("--model")
+    args = parser.parse_args()
+    embed_sections_and_keys(args.stem, user_keys=args.keys, model_name=args.model)
+    print("Done")
+
+
+if __name__ == "__main__":
+    main()