`src/classify.py`

85 lines
"""
Doc classifier — local Ollama (default gemma3:4b).

Runs before extract.py so the expensive model only sees filings that actually
contain assets (Schedule A/B, 363 motion, auctioneer retention, APA, abandonment).
"""
from __future__ import annotations

import argparse
import os
from pathlib import Path

import psycopg

from llm import generate_text

PG_DSN = os.environ.get("BANKRUPT_LEADS_PG", "postgresql:///bankrupt_leads")
MODEL = os.environ.get("CLASSIFY_MODEL", "gemma3:4b")

CATEGORIES = {
    "schedule_ab",
    "363_motion",
    "auctioneer_retention",
    "apa",
    "abandonment",
    "other",
}

PROMPT_TEMPLATE = """Classify this US bankruptcy court filing into exactly ONE category:

- schedule_ab           Schedule A/B (real & personal property inventory)
- 363_motion            Motion to sell assets under 11 U.S.C. section 363
- auctioneer_retention  Application to employ auctioneer or liquidator
- apa                   Asset Purchase Agreement (or notice attaching one)
- abandonment           Notice of abandonment of property under section 554
- other                 everything else (procedural, claims, plan, etc.)

Respond with ONLY the category string. No prose, no punctuation.

FILING TEXT:
{text}
"""


def classify(text: str) -> str:
    prompt = PROMPT_TEMPLATE.format(text=text[:20_000])
    raw = generate_text(MODEL, prompt, num_ctx=8192).strip().lower()
    # Model sometimes wraps in quotes or adds a period — strip.
    raw = raw.strip("`\"' .\n")
    return raw if raw in CATEGORIES else "other"


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--limit", type=int, default=50)
    args = ap.parse_args()

    with psycopg.connect(PG_DSN) as conn:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT filing_id, text_path FROM bankrupt_leads.filings
                WHERE classified_at IS NULL AND text_path IS NOT NULL
                LIMIT %s
                """,
                (args.limit,),
            )
            rows = cur.fetchall()

        for filing_id, text_path in rows:
            text = Path(text_path).read_text(encoding="utf-8", errors="replace")
            doc_type = classify(text)
            with conn.cursor() as cur:
                cur.execute(
                    "UPDATE bankrupt_leads.filings SET doc_type = %s, classified_at = now() "
                    "WHERE filing_id = %s",
                    (doc_type, filing_id),
                )
            conn.commit()
            print(f"{filing_id}: {doc_type}")


if __name__ == "__main__":
    main()