Upload ml/06_extract_keywords.py with huggingface_hub

ml/06_extract_keywords.py

@@ -0,0 +1,185 @@
+#!/usr/bin/env python3
+"""
+Phase 3: TF-IDF Keyword Extraction
+
+Extracts top keywords per document using TF-IDF vectorization.
+Processes in collection-level batches to get meaningful IDF weights.
+Stores results in document_keywords table.
+
+Runs on: Hetzner CPU (scikit-learn)
+"""
+
+import logging
+import sys
+from collections import defaultdict
+
+import numpy as np
+from sklearn.feature_extraction.text import TfidfVectorizer
+
+import psycopg2.extras
+from db import get_conn
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(levelname)-8s  %(message)s")
+log = logging.getLogger(__name__)
+
+TOP_K = 15  # keywords per document
+BATCH_SIZE = 5000  # docs per TF-IDF batch
+MIN_DF = 3  # minimum document frequency
+MAX_DF = 0.85  # max document frequency (fraction)
+INSERT_BATCH = 1000
+
+
+def get_collections(conn):
+    with conn.cursor() as cur:
+        cur.execute("SELECT DISTINCT source_section FROM documents ORDER BY source_section")
+        return [r[0] for r in cur.fetchall()]
+
+
+def get_doc_texts(conn, collection, offset, limit):
+    """Get concatenated OCR text (first 5 pages) for a batch of documents."""
+    with conn.cursor() as cur:
+        cur.execute("""
+            SELECT d.id, string_agg(p.ocr_text, ' ' ORDER BY p.page_number)
+            FROM documents d
+            JOIN pages p ON p.document_id = d.id AND p.page_number <= 5
+            WHERE d.source_section = %s
+              AND d.id NOT IN (SELECT DISTINCT document_id FROM document_keywords)
+            GROUP BY d.id
+            ORDER BY d.id
+            OFFSET %s LIMIT %s
+        """, (collection, offset, limit))
+        return cur.fetchall()
+
+
+def extract_keywords_batch(doc_ids, texts):
+    """Run TF-IDF on a batch and return top keywords per document."""
+    if not texts:
+        return {}
+
+    vectorizer = TfidfVectorizer(
+        max_features=10000,
+        min_df=min(MIN_DF, max(1, len(texts) // 10)),
+        max_df=MAX_DF,
+        stop_words='english',
+        ngram_range=(1, 2),
+        token_pattern=r'(?u)\b[a-zA-Z][a-zA-Z]{2,}\b',  # 3+ letter words only
+        sublinear_tf=True,
+    )
+
+    try:
+        tfidf_matrix = vectorizer.fit_transform(texts)
+    except ValueError as e:
+        log.warning(f"TF-IDF failed for batch: {e}")
+        return {}
+
+    feature_names = vectorizer.get_feature_names_out()
+    results = {}
+
+    for i, doc_id in enumerate(doc_ids):
+        row = tfidf_matrix[i].toarray().flatten()
+        top_indices = row.argsort()[-TOP_K:][::-1]
+        keywords = []
+        for idx in top_indices:
+            score = row[idx]
+            if score > 0:
+                keywords.append((feature_names[idx], float(score)))
+        results[doc_id] = keywords
+
+    return results
+
+
+def flush_keywords(conn, rows):
+    with conn.cursor() as cur:
+        psycopg2.extras.execute_batch(
+            cur,
+            """INSERT INTO document_keywords (document_id, keyword, score, method)
+               VALUES (%s, %s, %s, 'tfidf')
+               ON CONFLICT DO NOTHING""",
+            rows,
+            page_size=1000,
+        )
+    conn.commit()
+
+
+def main():
+    conn = get_conn()
+    collections = get_collections(conn)
+
+    total_docs = 0
+    total_keywords = 0
+
+    for collection in collections:
+        # Count remaining docs for this collection
+        with conn.cursor() as cur:
+            cur.execute("""
+                SELECT COUNT(DISTINCT d.id) FROM documents d
+                JOIN pages p ON p.document_id = d.id AND p.page_number <= 5
+                WHERE d.source_section = %s
+                  AND d.id NOT IN (SELECT DISTINCT document_id FROM document_keywords)
+            """, (collection,))
+            remaining = cur.fetchone()[0]
+
+        if remaining == 0:
+            log.info(f"  {collection}: already done, skipping")
+            continue
+
+        log.info(f"Processing {collection}: {remaining} documents")
+        offset = 0
+        insert_buffer = []
+
+        while True:
+            rows = get_doc_texts(conn, collection, 0, BATCH_SIZE)  # offset=0 since we exclude already-done
+            if not rows:
+                break
+
+            doc_ids = [r[0] for r in rows]
+            texts = [r[1] or '' for r in rows]
+
+            # Mark all docs as processed (even short/failed ones) with a placeholder
+            skip_ids = []
+
+            # Filter out very short texts
+            valid = [(did, txt) for did, txt in zip(doc_ids, texts) if len(txt.strip()) > 50]
+            short_ids = [did for did, txt in zip(doc_ids, texts) if len(txt.strip()) <= 50]
+            skip_ids.extend(short_ids)
+
+            if valid:
+                v_ids, v_texts = zip(*valid)
+                keywords = extract_keywords_batch(list(v_ids), list(v_texts))
+
+                if not keywords:
+                    # TF-IDF failed for entire batch
+                    skip_ids.extend(list(v_ids))
+                else:
+                    for doc_id, kws in keywords.items():
+                        if kws:
+                            for kw, score in kws:
+                                insert_buffer.append((doc_id, kw, score))
+                        else:
+                            skip_ids.append(doc_id)
+            else:
+                skip_ids.extend(doc_ids)
+
+            # Insert placeholder for skipped docs so they don't get re-fetched
+            if skip_ids:
+                placeholder_rows = [(did, '_no_keywords_', 0.0) for did in skip_ids]
+                flush_keywords(conn, placeholder_rows)
+
+            if len(insert_buffer) >= INSERT_BATCH:
+                flush_keywords(conn, insert_buffer)
+                total_keywords += len(insert_buffer)
+                insert_buffer = []
+
+            total_docs += len(rows)
+            log.info(f"  {collection}: {total_docs} docs processed, {total_keywords} keywords stored (skipped {len(skip_ids)})")
+
+        if insert_buffer:
+            flush_keywords(conn, insert_buffer)
+            total_keywords += len(insert_buffer)
+
+    conn.close()
+    log.info(f"Done. {total_docs} documents, {total_keywords} keywords extracted.")
+
+
+if __name__ == "__main__":
+    main()