Upload notebooks/04_forensic/43_metadata_extraction.ipynb with huggingface_hub

notebooks/04_forensic/43_metadata_extraction.ipynb

@@ -0,0 +1,242 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 43 - Metadata Extraction\n",
+    "\n",
+    "Pipeline notebook that extracts PDF metadata and scans for classification stamps.\n",
+    "\n",
+    "**Document features stored:**\n",
+    "- `pdf_metadata` -- JSONB with author, title, subject, creator, producer, dates\n",
+    "- `classification_stamps` -- JSONB list of stamps found (CONFIDENTIAL, CLASSIFIED, etc.)\n",
+    "\n",
+    "Incremental: skips documents that already have `pdf_metadata` in `document_features`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": [
+     "parameters"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "# Parameters\n",
+    "source_section = None\n",
+    "batch_size = 5000"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys, warnings, re, json, time\n",
+    "sys.path.insert(0, '/opt/epstein_env/research')\n",
+    "warnings.filterwarnings('ignore')\n",
+    "\n",
+    "import fitz\n",
+    "import pandas as pd\n",
+    "from pathlib import Path\n",
+    "\n",
+    "from research_lib.config import RAW_DIR, COLLECTIONS\n",
+    "from research_lib.db import fetch_df, upsert_feature\n",
+    "from research_lib.incremental import (\n",
+    "    start_run, finish_run, get_unprocessed_documents,\n",
+    ")\n",
+    "\n",
+    "print('Libraries loaded.')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ---- Classification stamp patterns ----\n",
+    "STAMP_PATTERNS = [\n",
+    "    'CONFIDENTIAL',\n",
+    "    'CLASSIFIED',\n",
+    "    'SEALED',\n",
+    "    'TOP SECRET',\n",
+    "    'RESTRICTED',\n",
+    "    'UNCLASSIFIED',\n",
+    "    'DECLASSIFIED',\n",
+    "    'FOR OFFICIAL USE ONLY',\n",
+    "]\n",
+    "\n",
+    "# Compile combined regex for efficiency\n",
+    "_stamp_re = re.compile(\n",
+    "    r'\\b(' + '|'.join(re.escape(s) for s in STAMP_PATTERNS) + r')\\b',\n",
+    "    re.IGNORECASE,\n",
+    ")\n",
+    "\n",
+    "\n",
+    "def extract_pdf_metadata(pdf_path):\n",
+    "    \"\"\"Extract metadata dict from a PDF file.\"\"\"\n",
+    "    try:\n",
+    "        doc = fitz.open(pdf_path)\n",
+    "        meta = doc.metadata or {}\n",
+    "        result = {\n",
+    "            'author': meta.get('author', ''),\n",
+    "            'title': meta.get('title', ''),\n",
+    "            'subject': meta.get('subject', ''),\n",
+    "            'creator': meta.get('creator', ''),\n",
+    "            'producer': meta.get('producer', ''),\n",
+    "            'creationDate': meta.get('creationDate', ''),\n",
+    "            'modDate': meta.get('modDate', ''),\n",
+    "        }\n",
+    "        doc.close()\n",
+    "        return result\n",
+    "    except Exception:\n",
+    "        return {}\n",
+    "\n",
+    "\n",
+    "def scan_for_stamps(document_id):\n",
+    "    \"\"\"Check the first page OCR text for classification stamps.\"\"\"\n",
+    "    first_page = fetch_df(\n",
+    "        'SELECT ocr_text FROM pages WHERE document_id = %s ORDER BY page_number LIMIT 1',\n",
+    "        [document_id],\n",
+    "    )\n",
+    "    if first_page.empty or not first_page.iloc[0]['ocr_text']:\n",
+    "        return []\n",
+    "    text = first_page.iloc[0]['ocr_text']\n",
+    "    found = list(set(m.upper() for m in _stamp_re.findall(text)))\n",
+    "    return sorted(found)\n",
+    "\n",
+    "\n",
+    "print('Functions defined.')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ---- Identify unprocessed documents ----\n",
+    "PIPELINE = 'metadata_extraction'\n",
+    "run_id = start_run(PIPELINE, source_section=source_section, parameters={\n",
+    "    'batch_size': batch_size,\n",
+    "})\n",
+    "\n",
+    "docs_df = get_unprocessed_documents(\n",
+    "    PIPELINE, source_section=source_section,\n",
+    "    feature_table='document_features', feature_name='pdf_metadata',\n",
+    ")\n",
+    "print(f'Documents to process: {len(docs_df)}')\n",
+    "if len(docs_df) > 0:\n",
+    "    print(docs_df['source_section'].value_counts().to_string())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ---- Process in batches ----\n",
+    "total_processed = 0\n",
+    "stamp_counter = {s: 0 for s in STAMP_PATTERNS}\n",
+    "meta_coverage = {'has_author': 0, 'has_title': 0, 'has_date': 0}\n",
+    "total_batches = (len(docs_df) + batch_size - 1) // batch_size\n",
+    "\n",
+    "for batch_idx in range(total_batches):\n",
+    "    start = batch_idx * batch_size\n",
+    "    end = min(start + batch_size, len(docs_df))\n",
+    "    batch = docs_df.iloc[start:end]\n",
+    "    print(f'\\nBatch {batch_idx + 1}/{total_batches}: documents {start}-{end - 1}')\n",
+    "\n",
+    "    t0 = time.time()\n",
+    "    meta_rows = []\n",
+    "    stamp_rows = []\n",
+    "\n",
+    "    for _, row in batch.iterrows():\n",
+    "        doc_id = int(row['id'])\n",
+    "\n",
+    "        # Extract PDF metadata\n",
+    "        pdf_path = Path(row.get('file_path', '')) if row.get('file_path') else None\n",
+    "        if pdf_path is None or not pdf_path.exists():\n",
+    "            pdf_path = RAW_DIR / row['source_section'] / row['filename']\n",
+    "\n",
+    "        meta = extract_pdf_metadata(str(pdf_path)) if pdf_path.exists() else {}\n",
+    "        meta_rows.append((doc_id, 'pdf_metadata', None, json.dumps(meta)))\n",
+    "\n",
+    "        # Track coverage\n",
+    "        if meta.get('author'):\n",
+    "            meta_coverage['has_author'] += 1\n",
+    "        if meta.get('title'):\n",
+    "            meta_coverage['has_title'] += 1\n",
+    "        if meta.get('creationDate') or meta.get('modDate'):\n",
+    "            meta_coverage['has_date'] += 1\n",
+    "\n",
+    "        # Scan for classification stamps\n",
+    "        stamps = scan_for_stamps(doc_id)\n",
+    "        stamp_rows.append((doc_id, 'classification_stamps', None, json.dumps(stamps)))\n",
+    "        for s in stamps:\n",
+    "            if s in stamp_counter:\n",
+    "                stamp_counter[s] += 1\n",
+    "\n",
+    "    # Upsert features\n",
+    "    n_meta = upsert_feature(\n",
+    "        'document_features',\n",
+    "        ['document_id', 'feature_name'],\n",
+    "        ['feature_value', 'feature_json'],\n",
+    "        meta_rows,\n",
+    "    )\n",
+    "    n_stamp = upsert_feature(\n",
+    "        'document_features',\n",
+    "        ['document_id', 'feature_name'],\n",
+    "        ['feature_value', 'feature_json'],\n",
+    "        stamp_rows,\n",
+    "    )\n",
+    "    elapsed = time.time() - t0\n",
+    "    print(f'  {n_meta} metadata + {n_stamp} stamp rows in {elapsed:.1f}s')\n",
+    "    total_processed += len(batch)\n",
+    "\n",
+    "finish_run(run_id, documents_processed=total_processed)\n",
+    "print(f'\\nRun {run_id} complete: {total_processed} documents processed.')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ---- Summary: metadata coverage ----\n",
+    "total = max(total_processed, 1)\n",
+    "print('Metadata Coverage:')\n",
+    "for field, count in meta_coverage.items():\n",
+    "    print(f'  {field}: {count} / {total} ({count / total * 100:.1f}%)')\n",
+    "\n",
+    "print('\\nClassification Stamp Frequency:')\n",
+    "for stamp, count in sorted(stamp_counter.items(), key=lambda x: -x[1]):\n",
+    "    if count > 0:\n",
+    "        print(f'  {stamp}: {count}')\n",
+    "\n",
+    "if all(v == 0 for v in stamp_counter.values()):\n",
+    "    print('  No classification stamps found.')"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}