From 1bf7d2cfa94ce34f3abcb75df1aba973b3d2f8d4 Mon Sep 17 00:00:00 2001 From: raymondcen Date: Thu, 21 May 2026 23:09:23 -0700 Subject: [PATCH 1/4] fix: consolidate both entry points --- README.md | 3 + documentation/CONTRIBUTING.md | 23 +- documentation/DESIGN_DECISIONS.md | 2 +- src/pipeline/classify_extract.py | 825 ++++++++++++++++++++++++------ src/pipeline/extract_from_txt.py | 470 ----------------- 5 files changed, 688 insertions(+), 635 deletions(-) delete mode 100644 src/pipeline/extract_from_txt.py diff --git a/README.md b/README.md index ff13809..097ca8d 100644 --- a/README.md +++ b/README.md @@ -143,6 +143,9 @@ pip install -e ".[dev]" # Classify and extract from a folder of PDFs python src/pipeline/classify_extract.py path/to/pdfs/ +# Extract from pre-classified .txt files (skip XGBoost) +python src/pipeline/classify_extract.py path/to/txts/ --skip-classifier + # Adjust the LLM model or confidence threshold python src/pipeline/classify_extract.py path/to/pdfs/ --llm-model qwen2.5:7b --confidence-threshold 0.70 ``` diff --git a/documentation/CONTRIBUTING.md b/documentation/CONTRIBUTING.md index 0856e74..61ccadc 100644 --- a/documentation/CONTRIBUTING.md +++ b/documentation/CONTRIBUTING.md @@ -96,6 +96,25 @@ All contributors must follow the Oregon State University Student Code of Conduct | `--num-ctx` | `4096` | Ollama context window size (tokens) | | `--workers` | `1` | Parallel worker processes (`1` = sequential) | +* ### Extracting from pre-classified .txt files + When PDFs have already been classified (or labels come from `labels.json`), skip + XGBoost and run clean → filter → extract on `.txt` files: + ```bash + python src/pipeline/classify_extract.py path/to/txts/ --skip-classifier + + # Optional: same flags as the former extract_from_txt.py + python src/pipeline/classify_extract.py path/to/txts/ --skip-classifier --chunked + python src/pipeline/classify_extract.py path/to/txts/ --skip-classifier --labels labels.json + ``` + | Flag | Default | Description | + |------|---------|-------------| + | `--skip-classifier` | off | Process `.txt` files without classification | + | `--labels` | — | `labels.json` path; only stems labelled `useful` are processed | + | `--chunked` | off | Chunked extraction with per-chunk scoring and majority-vote merge | + | `--top-chunks` | `3` | Top chunks to extract (with `--chunked`) | + | `--chunk-size` | `4000` | Chunk size in characters (with `--chunked`) | + | `--chunk-overlap` | `500` | Overlap between chunks (with `--chunked`) | + * ### Sample Output Each PDF classified as "useful" produces a JSON file in `data/results/metrics/`: ```json @@ -389,8 +408,8 @@ Extraction fields are defined in two places: 2. **`src/extraction/llm_client.py`** - the system prompt that instructs the LLM. Add a description of the new field and its expected format to the prompt string. -3. **`src/pipeline/classify_extract.py`** and **`src/pipeline/extract_from_txt.py`** - update the `row` dict - and `fieldnames` list in the summary CSV writer to include the new column. +3. **`src/pipeline/classify_extract.py`** - update the summary `row` dict and `fieldnames` list + in the summary CSV writer to include the new column. After adding a field, run `pytest tests/test_llm_text.py` to verify that the prompt changes do not break existing extraction tests. diff --git a/documentation/DESIGN_DECISIONS.md b/documentation/DESIGN_DECISIONS.md index 77d2355..0c2fc30 100644 --- a/documentation/DESIGN_DECISIONS.md +++ b/documentation/DESIGN_DECISIONS.md @@ -152,7 +152,7 @@ The `src/` directory was reorganized once near the end of the project (PR #64, M | `src/llm/` | `src/extraction/` | | `src/preprocessing/` | `src/io/` | | `classify_extract.py` (root) | `src/pipeline/classify_extract.py` | -| `extract-from-txt.py` (root) | `src/pipeline/extract_from_txt.py` | +| `extract-from-txt.py` (root) | merged into `src/pipeline/classify_extract.py` (`--skip-classifier`) | `requirements.txt` was also removed at this point in favor of `pyproject.toml` as the single source of dependency truth. diff --git a/src/pipeline/classify_extract.py b/src/pipeline/classify_extract.py index adc7d63..cfd3d4a 100644 --- a/src/pipeline/classify_extract.py +++ b/src/pipeline/classify_extract.py @@ -1,15 +1,17 @@ """Classify-and-Extract Pipeline -Accepts a single PDF or a folder of PDFs, classifies each using the trained -XGBoost classifier, and passes only the files classified as "useful" to the -LLM for structured data extraction. +Single entry point for the full pipeline: PDF classify → extract, or +pre-classified .txt extract (with ``--skip-classifier``). Usage: + # PDF folder (classify then extract useful papers) + python src/pipeline/classify_extract.py path/to/folder/ + # Single PDF python src/pipeline/classify_extract.py path/to/file.pdf - # Folder of PDFs - python src/pipeline/classify_extract.py path/to/folder/ + # Pre-classified .txt files (no XGBoost step) + python src/pipeline/classify_extract.py path/to/txts/ --skip-classifier # Custom options python src/pipeline/classify_extract.py path/to/folder/ \\ @@ -21,30 +23,331 @@ --num-ctx 4096 Output: - - One JSON file per useful PDF (in --output-dir) containing extracted metrics. - - A summary CSV (pipeline_summary.csv) in --output-dir listing every PDF, - its classification, confidence, and extraction status. + - One JSON file per useful PDF or processed .txt (in --output-dir). + - A summary CSV in --output-dir/summaries/ listing each file and status. """ import argparse import csv +import json import logging import sys -from datetime import datetime from concurrent.futures import ProcessPoolExecutor, as_completed +from datetime import datetime from pathlib import Path from src.io.pdf_text_extraction import extract_text_from_pdf +from src.io.text_cleaner import clean_text +from src.io.section_filter import filter_relevant_sections from src.classifier.pdf_classifier import load_classifier, classify_text from src.extraction.llm_text import extract_key_sections from src.extraction.llm_client import extract_metrics_from_text, save_extraction_result +from src.extraction.chunked_extraction import extract_with_chunking from src.utils.logger import setup_logging log = logging.getLogger(__name__) +METRIC_FIELDNAMES = [ + "species_name", + "study_location", + "study_date", + "sample_size", + "num_empty_stomachs", + "num_nonempty_stomachs", + "fraction_feeding", +] +CLASSIFIER_FIELDNAMES = ["classification", "confidence", "pred_prob"] +PREPROCESS_FIELDNAMES = ["raw_chars", "cleaned_chars", "filtered_chars", "trimmed_chars"] + + +# --------------------------------------------------------------------------- +# Shared helpers +# --------------------------------------------------------------------------- + + +def _summary_fieldnames(*, include_classifier: bool, include_preprocess: bool) -> list[str]: + fields = ["filename"] + if include_classifier: + fields.extend(CLASSIFIER_FIELDNAMES) + if include_preprocess: + fields.extend(PREPROCESS_FIELDNAMES) + fields.append("extraction_status") + fields.extend(METRIC_FIELDNAMES) + return fields + + +def _new_summary_row(filename: str, *, include_classifier: bool, include_preprocess: bool) -> dict: + row = {"filename": filename, "extraction_status": ""} + if include_classifier: + row.update({k: "" for k in CLASSIFIER_FIELDNAMES}) + if include_preprocess: + row.update({k: "" for k in PREPROCESS_FIELDNAMES}) + row.update({k: "" for k in METRIC_FIELDNAMES}) + return row + + +def _populate_metrics_row(row: dict, metrics: dict) -> None: + row["species_name"] = metrics.get("species_name") or "" + row["study_location"] = metrics.get("study_location") or "" + row["study_date"] = metrics.get("study_date") or "" + row["sample_size"] = "" if metrics.get("sample_size") is None else metrics["sample_size"] + row["num_empty_stomachs"] = "" if metrics.get("num_empty_stomachs") is None else metrics["num_empty_stomachs"] + row["num_nonempty_stomachs"] = "" if metrics.get("num_nonempty_stomachs") is None else metrics["num_nonempty_stomachs"] + row["fraction_feeding"] = "" if metrics.get("fraction_feeding") is None else metrics["fraction_feeding"] + + +def _load_useful_stems(labels_path: Path) -> set[str]: + if not labels_path.exists(): + print(f"[ERROR] Labels file not found: {labels_path}", file=sys.stderr) + sys.exit(1) + with open(labels_path, encoding="utf-8") as f: + labels = json.load(f) + useful = {k for k, v in labels.items() if v == "useful"} + print(f"[INFO] Labels filter: {len(useful)} useful papers", file=sys.stderr) + return useful + + +def _collect_pdf_paths(input_path: Path) -> list[Path]: + if input_path.is_dir(): + pdf_paths = sorted(input_path.glob("*.pdf")) + if not pdf_paths: + print(f"[ERROR] No PDF files found in directory: {input_path}", file=sys.stderr) + sys.exit(1) + print(f"[INFO] Found {len(pdf_paths)} PDF(s) in {input_path}", file=sys.stderr) + return pdf_paths + if input_path.is_file() and input_path.suffix.lower() == ".pdf": + return [input_path] + print(f"[ERROR] Input must be a .pdf file or a directory of PDFs: {input_path}", file=sys.stderr) + sys.exit(1) + + +def _collect_txt_paths( + input_dir: Path, + *, + single_file: Path | None = None, + useful_stems: set | None = None, +) -> list[Path]: + if single_file is not None: + return [single_file] + txt_paths = sorted(input_dir.glob("*.txt")) + if useful_stems is not None: + txt_paths = [p for p in txt_paths if p.stem in useful_stems or p.name in useful_stems] + if not txt_paths: + print(f"[ERROR] No .txt files found in: {input_dir}", file=sys.stderr) + sys.exit(1) + print(f"[INFO] Found {len(txt_paths)} .txt file(s) to process", file=sys.stderr) + return txt_paths + + +def _cleaned_text_dirs(output_dir: Path) -> tuple[Path, Path, Path]: + base = output_dir.parent / "cleaned-text" + cleaner_dir = base / "text_cleaner" + filter_dir = base / "section_filter" + llm_dir = base / "llm_text" + for d in (cleaner_dir, filter_dir, llm_dir): + d.mkdir(parents=True, exist_ok=True) + return cleaner_dir, filter_dir, llm_dir + + +def _preprocess_text( + raw_text: str, + max_chars: int, + *, + stem: str, + save_intermediates: bool, + cleaner_dir: Path | None, + filter_dir: Path | None, + llm_dir: Path | None, +) -> tuple[str | None, str | None, dict, str | None]: + counts: dict = {"raw_chars": len(raw_text)} + ts = datetime.now().strftime("%Y%m%d_%H%M%S") + + if not raw_text.strip(): + return None, None, counts, "empty_text" + + cleaned = clean_text(raw_text) + counts["cleaned_chars"] = len(cleaned) + print(f" [INFO] After clean: {len(cleaned):,} chars", file=sys.stderr) + + if not cleaned.strip(): + return None, None, counts, "empty_after_clean" + + if save_intermediates and cleaner_dir is not None: + _write_intermediate(cleaner_dir / f"{stem}_{ts}.txt", cleaned, "Cleaner text") + + filtered = filter_relevant_sections(cleaned) + counts["filtered_chars"] = len(filtered) + print(f" [INFO] After filter: {len(filtered):,} chars", file=sys.stderr) + + if save_intermediates and filter_dir is not None: + _write_intermediate(filter_dir / f"{stem}_{ts}.txt", filtered, "Filter text") + + if len(filtered) > max_chars: + trimmed = extract_key_sections(filtered, max_chars) + print( + f" [INFO] After trim : {len(trimmed):,} chars (budget {max_chars:,})", + file=sys.stderr, + ) + else: + trimmed = filtered + + counts["trimmed_chars"] = len(trimmed) + + if save_intermediates and llm_dir is not None: + _write_intermediate(llm_dir / f"{stem}_{ts}.txt", trimmed, "LLM text") + + return filtered, trimmed, counts, None + + +def _write_intermediate(path: Path, text: str, label: str) -> None: + try: + path.write_text(text, encoding="utf-8") + print(f" [INFO] {label:<13}: {path.name}", file=sys.stderr) + except Exception as exc: + print(f" [WARN] Could not save {label.lower()}: {exc}", file=sys.stderr) + + +def _run_llm_extraction( + *, + filtered_text: str, + trimmed_text: str, + source_file: Path, + original_text: str, + output_dir: Path, + llm_model: str, + num_ctx: int, + chunked: bool, + model_dir: str, + top_chunks: int, + chunk_size: int, + chunk_overlap: int, +) -> tuple[dict | None, str | None]: + if chunked: + print(f" [INFO] Chunked extraction (top {top_chunks} chunks)…", file=sys.stderr) + try: + merged = extract_with_chunking( + text=filtered_text, + model_dir=model_dir, + llm_model=llm_model, + num_ctx=num_ctx, + top_n=top_chunks, + chunk_size=chunk_size, + overlap=chunk_overlap, + ) + except Exception as exc: + print(f" [ERROR] Chunked extraction failed: {exc}", file=sys.stderr) + log.error("Chunked extraction failed for %s: %s", source_file.name, exc) + return None, "extraction_failed" + + metrics_dir = output_dir / "metrics" + metrics_dir.mkdir(parents=True, exist_ok=True) + out_path = metrics_dir / f"{source_file.stem}_results.json" + result_obj = { + "source_file": source_file.name, + "file_type": source_file.suffix.lower(), + "extraction_mode": "chunked", + "metrics": merged, + } + with open(out_path, "w", encoding="utf-8") as f: + json.dump(result_obj, f, indent=2) + print(f" [SUCCESS] Results saved to {out_path}", file=sys.stderr) + return merged, None + + print(f" [INFO] Calling Ollama ({llm_model})…", file=sys.stderr) + try: + metrics = extract_metrics_from_text( + text=trimmed_text, + model=llm_model, + num_ctx=num_ctx, + ) + result = save_extraction_result( + metrics=metrics, + source_file=source_file, + original_text=original_text, + output_dir=output_dir, + ) + return result["metrics"], None + except Exception as exc: + print(f" [ERROR] LLM extraction failed: {exc}", file=sys.stderr) + log.error("LLM extraction failed for %s: %s", source_file.name, exc) + return None, "extraction_failed" + + +def _write_summary_csv( + summary_rows: list[dict], + output_dir: Path, + *, + include_classifier: bool, + include_preprocess: bool, + prefix: str = "pipeline", +) -> Path: + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + summaries_dir = output_dir / "summaries" + summaries_dir.mkdir(parents=True, exist_ok=True) + summary_path = summaries_dir / f"{prefix}_summary_{timestamp}.csv" + fieldnames = _summary_fieldnames( + include_classifier=include_classifier, + include_preprocess=include_preprocess, + ) + with open(summary_path, "w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore") + writer.writeheader() + writer.writerows(summary_rows) + return summary_path + + +def _print_classify_summary(summary_rows: list[dict], summary_path: Path) -> None: + total = len(summary_rows) + useful_count = sum(1 for r in summary_rows if r.get("classification") == "useful") + not_useful_count = sum(1 for r in summary_rows if r.get("classification") == "not useful") + extracted_count = sum(1 for r in summary_rows if r.get("extraction_status") == "success") + error_statuses = { + "text_extraction_failed", + "empty_text", + "empty_after_clean", + "extraction_failed", + "worker_failed", + } + error_count = sum(1 for r in summary_rows if r.get("extraction_status") in error_statuses) + + print("\n" + "=" * 50, file=sys.stderr) + print("PIPELINE COMPLETE", file=sys.stderr) + print("=" * 50, file=sys.stderr) + print(f" Total PDFs processed : {total}", file=sys.stderr) + print(f" Useful : {useful_count}", file=sys.stderr) + print(f" Not useful : {not_useful_count}", file=sys.stderr) + print(f" Successfully extracted: {extracted_count}", file=sys.stderr) + print(f" Errors : {error_count}", file=sys.stderr) + print(f" Summary CSV : {summary_path}", file=sys.stderr) + print("=" * 50, file=sys.stderr) + + if error_count > 0: + log.warning("Pipeline finished with %d error(s). See logs/fracfeed.log for details.", error_count) + + +def _print_txt_summary(summary_rows: list[dict], summary_path: Path) -> None: + total = len(summary_rows) + succeeded = sum(1 for r in summary_rows if r.get("extraction_status") == "success") + failed = total - succeeded + + print("\n" + "=" * 55, file=sys.stderr) + print("TXT EXTRACTION PIPELINE COMPLETE", file=sys.stderr) + print("=" * 55, file=sys.stderr) + print(f" Files processed : {total}", file=sys.stderr) + print(f" Successful : {succeeded}", file=sys.stderr) + print(f" Failed / skipped : {failed}", file=sys.stderr) + print(f" Summary CSV : {summary_path}", file=sys.stderr) + print("=" * 55, file=sys.stderr) + + +def _apply_char_counts(row: dict, counts: dict) -> None: + for key in PREPROCESS_FIELDNAMES: + if key in row and key in counts: + row[key] = counts[key] + # --------------------------------------------------------------------------- -# Pipeline +# Per-file processing # --------------------------------------------------------------------------- @@ -58,26 +361,18 @@ def _process_single_pdf( clf_model, vectorizer, encoder, + *, + chunked: bool, + model_dir: str, + top_chunks: int, + chunk_size: int, + chunk_overlap: int, ): """Classify one PDF and return a summary row dict.""" output_dir.mkdir(parents=True, exist_ok=True) - row = { - "filename": pdf_path.name, - "classification": "", - "confidence": "", - "pred_prob": "", - "extraction_status": "", - "species_name": "", - "study_location": "", - "study_date": "", - "sample_size": "", - "num_empty_stomachs": "", - "num_nonempty_stomachs": "", - "fraction_feeding": "", - } + row = _new_summary_row(pdf_path.name, include_classifier=True, include_preprocess=False) - # ── Step 1: Extract text ────────────────────────────────────────── try: original_text = extract_text_from_pdf(str(pdf_path)) except Exception as e: @@ -94,7 +389,6 @@ def _process_single_pdf( print(f" [INFO] {pdf_path.name}: {len(original_text)} chars", file=sys.stderr) - # ── Step 2: Classify ────────────────────────────────────────────── label, confidence, pred_prob = classify_text( text=original_text, model=clf_model, @@ -108,51 +402,166 @@ def _process_single_pdf( row["confidence"] = f"{confidence:.4f}" row["pred_prob"] = f"{pred_prob:.4f}" - # ── Step 3: Extract ─────────────────────────────────────────────── - if label == "useful": - print(f" [INFO] {pdf_path.name}: Running LLM extraction...", file=sys.stderr) + if label != "useful": + print(f" [INFO] {pdf_path.name}: Not useful — skipping LLM extraction.", file=sys.stderr) + row["extraction_status"] = "skipped_not_useful" + return row - text_for_llm = original_text - if len(text_for_llm) > max_chars: - text_for_llm = extract_key_sections(text_for_llm, max_chars) - print(f" [INFO] {pdf_path.name}: trimmed to {len(text_for_llm)} chars (budget {max_chars})", file=sys.stderr) + print(f" [INFO] {pdf_path.name}: Running LLM extraction...", file=sys.stderr) + row = _run_extraction_on_text( + row, + raw_text=original_text, + source_file=pdf_path, + original_text=original_text, + output_dir=output_dir, + llm_model=llm_model, + max_chars=max_chars, + num_ctx=num_ctx, + save_intermediates=False, + cleaner_dir=None, + filter_dir=None, + llm_dir=None, + chunked=chunked, + model_dir=model_dir, + top_chunks=top_chunks, + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + ) + return row - try: - metrics = extract_metrics_from_text( - text=text_for_llm, - model=llm_model, - num_ctx=num_ctx, - ) - result = save_extraction_result( - metrics=metrics, - source_file=pdf_path, - original_text=original_text, - output_dir=output_dir, - ) - m = result["metrics"] - row["extraction_status"] = "success" - row["species_name"] = m.get("species_name") or "" - row["study_location"] = m.get("study_location") or "" - row["study_date"] = m.get("study_date") or "" - row["sample_size"] = "" if m.get("sample_size") is None else m["sample_size"] - row["num_empty_stomachs"] = "" if m.get("num_empty_stomachs") is None else m["num_empty_stomachs"] - row["num_nonempty_stomachs"] = "" if m.get("num_nonempty_stomachs") is None else m["num_nonempty_stomachs"] - row["fraction_feeding"] = "" if m.get("fraction_feeding") is None else m["fraction_feeding"] - - except Exception as e: - print(f" [ERROR] LLM extraction failed ({pdf_path.name}): {e}", file=sys.stderr) - log.error("LLM extraction failed for %s: %s", pdf_path.name, e) - row["extraction_status"] = "extraction_failed" +def _process_single_txt( + txt_path: Path, + llm_model: str, + output_dir: Path, + max_chars: int, + num_ctx: int, + cleaner_dir: Path, + filter_dir: Path, + llm_dir: Path, + *, + chunked: bool, + model_dir: str, + top_chunks: int, + chunk_size: int, + chunk_overlap: int, +): + """Process one .txt file (skip-classifier mode) and return a summary row dict.""" + row = _new_summary_row(txt_path.name, include_classifier=False, include_preprocess=True) - else: - print(f" [INFO] {pdf_path.name}: Not useful — skipping LLM extraction.", file=sys.stderr) - row["extraction_status"] = "skipped_not_useful" + try: + raw_text = txt_path.read_text(encoding="utf-8") + except Exception as exc: + print(f" [ERROR] Could not read file: {exc}", file=sys.stderr) + row["extraction_status"] = "read_failed" + return row + + row["raw_chars"] = len(raw_text) + print(f" [INFO] Raw size : {len(raw_text):,} chars", file=sys.stderr) + if not raw_text.strip(): + print(f" [WARN] File is empty — skipping.", file=sys.stderr) + row["extraction_status"] = "empty_file" + return row + + return _run_extraction_on_text( + row, + raw_text=raw_text, + source_file=txt_path, + original_text=raw_text, + output_dir=output_dir, + llm_model=llm_model, + max_chars=max_chars, + num_ctx=num_ctx, + save_intermediates=True, + cleaner_dir=cleaner_dir, + filter_dir=filter_dir, + llm_dir=llm_dir, + chunked=chunked, + model_dir=model_dir, + top_chunks=top_chunks, + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + ) + + +def _run_extraction_on_text( + row: dict, + *, + raw_text: str, + source_file: Path, + original_text: str, + output_dir: Path, + llm_model: str, + max_chars: int, + num_ctx: int, + save_intermediates: bool, + cleaner_dir: Path | None, + filter_dir: Path | None, + llm_dir: Path | None, + chunked: bool, + model_dir: str, + top_chunks: int, + chunk_size: int, + chunk_overlap: int, +) -> dict: + """Clean, filter, optionally trim, then run LLM extraction.""" + filtered, trimmed, counts, err = _preprocess_text( + raw_text, + max_chars, + stem=source_file.stem, + save_intermediates=save_intermediates and not chunked, + cleaner_dir=cleaner_dir, + filter_dir=filter_dir, + llm_dir=llm_dir if not chunked else None, + ) + _apply_char_counts(row, counts) + + if err == "empty_text": + row["extraction_status"] = "empty_text" if "classification" in row else "empty_file" + return row + if err == "empty_after_clean": + row["extraction_status"] = "empty_after_clean" + return row + + if chunked: + row["trimmed_chars"] = "" + + metrics, extract_err = _run_llm_extraction( + filtered_text=filtered, + trimmed_text=trimmed if trimmed is not None else filtered, + source_file=source_file, + original_text=original_text, + output_dir=output_dir, + llm_model=llm_model, + num_ctx=num_ctx, + chunked=chunked, + model_dir=model_dir, + top_chunks=top_chunks, + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + ) + if extract_err: + row["extraction_status"] = extract_err + return row + + row["extraction_status"] = "success" + populate_metrics_row(row, metrics) + print( + f" [OK] species={metrics.get('species_name')} " + f"n={metrics.get('sample_size')} " + f"date={metrics.get('study_date')}", + file=sys.stderr, + ) return row -def run_pipeline( +# --------------------------------------------------------------------------- +# Pipeline runners +# --------------------------------------------------------------------------- + + +def run_pdf_pipeline( input_path: Path, model_dir: str, llm_model: str, @@ -160,39 +569,22 @@ def run_pipeline( confidence_threshold: float, max_chars: int, num_ctx: int, - workers: int = 1, + workers: int, + useful_stems: set | None, + *, + chunked: bool, + top_chunks: int, + chunk_size: int, + chunk_overlap: int, ): - """Run classify → extract pipeline on one or more PDFs. - - For each PDF: - 1. Extract text via PyMuPDF / OCR (pdf_text_extraction.py) - 2. Classify with XGBoost (pdf_classifier.py) - 3. If 'useful': trim text to budget (llm_text.py), run LLM extraction - (llm_client.py), and save result JSON (llm_client.py) - 4. Append a row to the summary CSV regardless of classification outcome - - Args: - input_path: Path to a single PDF or a directory of PDFs. - model_dir: Directory containing classifier model artifacts. - llm_model: Ollama model name for extraction. - output_dir: Where to write JSON results and the summary CSV. - confidence_threshold: Classifier probability threshold for 'useful'. - max_chars: Max characters to send to the LLM. - num_ctx: Context window size for Ollama. - workers: Number of parallel worker processes (default: 1 = sequential). - """ - # ── Collect PDF paths ───────────────────────────────────────────────── - if input_path.is_dir(): - pdf_paths = sorted(input_path.glob("*.pdf")) + """Classify PDFs and extract metrics from those labelled useful.""" + pdf_paths = _collect_pdf_paths(input_path) + if useful_stems is not None: + pdf_paths = [p for p in pdf_paths if p.stem in useful_stems or p.name in useful_stems] if not pdf_paths: - print(f"[ERROR] No PDF files found in directory: {input_path}", file=sys.stderr) + print("[ERROR] No PDF files match the labels filter.", file=sys.stderr) sys.exit(1) - print(f"[INFO] Found {len(pdf_paths)} PDF(s) in {input_path}", file=sys.stderr) - elif input_path.is_file() and input_path.suffix.lower() == ".pdf": - pdf_paths = [input_path] - else: - print(f"[ERROR] Input must be a .pdf file or a directory of PDFs: {input_path}", file=sys.stderr) - sys.exit(1) + print(f"[INFO] After labels filter: {len(pdf_paths)} PDF(s)", file=sys.stderr) output_dir.mkdir(parents=True, exist_ok=True) @@ -206,6 +598,13 @@ def run_pipeline( print("[INFO] Classifier loaded.", file=sys.stderr) summary_rows = [] + extract_kw = dict( + chunked=chunked, + model_dir=model_dir, + top_chunks=top_chunks, + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + ) if workers > 1 and len(pdf_paths) > 1: print(f"[INFO] Using {workers} worker processes.", file=sys.stderr) @@ -222,6 +621,7 @@ def run_pipeline( clf_model, vectorizer, encoder, + **extract_kw, ): pdf_path for pdf_path in pdf_paths } @@ -231,7 +631,10 @@ def run_pipeline( row = future.result() except Exception as exc: print(f" [ERROR] Worker failed for {pdf_path.name}: {exc}", file=sys.stderr) - row = {"filename": pdf_path.name, "extraction_status": "worker_failed"} + row = _new_summary_row( + pdf_path.name, include_classifier=True, include_preprocess=False + ) + row["extraction_status"] = "worker_failed" summary_rows.append(row) else: for idx, pdf_path in enumerate(pdf_paths, start=1): @@ -246,56 +649,80 @@ def run_pipeline( clf_model, vectorizer, encoder, + **extract_kw, ) summary_rows.append(row) - # ── Write summary CSV ───────────────────────────────────────────────── - from datetime import datetime + summary_path = _write_summary_csv( + summary_rows, + output_dir, + include_classifier=True, + include_preprocess=False, + prefix="pipeline", + ) + _print_classify_summary(summary_rows, summary_path) - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - summaries_dir = output_dir / "summaries" - summaries_dir.mkdir(parents=True, exist_ok=True) - summary_path = summaries_dir / f"pipeline_summary_{timestamp}.csv" - - fieldnames = [ - "filename", - "classification", - "confidence", - "pred_prob", - "extraction_status", - "species_name", - "study_location", - "study_date", - "sample_size", - "num_empty_stomachs", - "num_nonempty_stomachs", - "fraction_feeding", - ] - with open(summary_path, "w", newline="", encoding="utf-8") as f: - writer = csv.DictWriter(f, fieldnames=fieldnames) - writer.writeheader() - writer.writerows(summary_rows) - # ── Final summary ───────────────────────────────────────────────────── - total = len(summary_rows) - useful_count = sum(1 for r in summary_rows if r["classification"] == "useful") - not_useful_count = sum(1 for r in summary_rows if r["classification"] == "not useful") - extracted_count = sum(1 for r in summary_rows if r["extraction_status"] == "success") - error_count = sum(1 for r in summary_rows if r["extraction_status"] in ("text_extraction_failed", "empty_text", "extraction_failed")) +def run_txt_pipeline( + input_path: Path, + output_dir: Path, + llm_model: str, + max_chars: int, + num_ctx: int, + useful_stems: set | None, + *, + chunked: bool, + model_dir: str, + top_chunks: int, + chunk_size: int, + chunk_overlap: int, +): + """Process .txt files without classification (all inputs treated as useful).""" + if input_path.is_file(): + if input_path.suffix.lower() != ".txt": + print(f"[ERROR] With --skip-classifier, input must be a .txt file: {input_path}", file=sys.stderr) + sys.exit(1) + single_file = input_path + input_dir = input_path.parent + elif input_path.is_dir(): + single_file = None + input_dir = input_path + else: + print(f"[ERROR] Input path not found: {input_path}", file=sys.stderr) + sys.exit(1) - print("\n" + "=" * 50, file=sys.stderr) - print("PIPELINE COMPLETE", file=sys.stderr) - print("=" * 50, file=sys.stderr) - print(f" Total PDFs processed : {total}", file=sys.stderr) - print(f" Useful : {useful_count}", file=sys.stderr) - print(f" Not useful : {not_useful_count}", file=sys.stderr) - print(f" Successfully extracted: {extracted_count}", file=sys.stderr) - print(f" Errors : {error_count}", file=sys.stderr) - print(f" Summary CSV : {summary_path}", file=sys.stderr) - print("=" * 50, file=sys.stderr) + txt_paths = _collect_txt_paths(input_dir, single_file=single_file, useful_stems=useful_stems) + output_dir.mkdir(parents=True, exist_ok=True) + cleaner_dir, filter_dir, llm_dir = _cleaned_text_dirs(output_dir) + summary_rows = [] - if error_count > 0: - log.warning("Pipeline finished with %d error(s). See logs/fracfeed.log for details.", error_count) + for idx, txt_path in enumerate(txt_paths, start=1): + print(f"\n[{idx}/{len(txt_paths)}] {txt_path.name}", file=sys.stderr) + row = _process_single_txt( + txt_path, + llm_model, + output_dir, + max_chars, + num_ctx, + cleaner_dir, + filter_dir, + llm_dir, + chunked=chunked, + model_dir=model_dir, + top_chunks=top_chunks, + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + ) + summary_rows.append(row) + + summary_path = _write_summary_csv( + summary_rows, + output_dir, + include_classifier=False, + include_preprocess=True, + prefix="txt_pipeline", + ) + _print_txt_summary(summary_rows, summary_path) # --------------------------------------------------------------------------- @@ -305,28 +732,35 @@ def run_pipeline( def main(): parser = argparse.ArgumentParser( - description=("Classify PDFs as useful/not-useful, then extract structured diet " "metrics from useful ones using an LLM."), + description=( + "Classify PDFs as useful/not-useful and extract structured diet metrics, " + "or extract from pre-classified .txt files with --skip-classifier." + ), formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: - Single PDF: - python src/pipeline/classify_extract.py paper.pdf - - Folder of PDFs: + PDF folder: python src/pipeline/classify_extract.py data/pdfs/ - Custom options: - python src/pipeline/classify_extract.py data/pdfs/ \\ - --model-dir src/classifier/models \\ - --output-dir results/ \\ - --llm-model qwen2.5:7b \\ - --confidence-threshold 0.70 + Pre-classified .txt folder: + python src/pipeline/classify_extract.py data/processed-text/ --skip-classifier + + Chunked .txt extraction: + python src/pipeline/classify_extract.py data/processed-text/ --skip-classifier --chunked + + Labels filter: + python src/pipeline/classify_extract.py data/processed-text/ --skip-classifier --labels labels.json """, ) parser.add_argument( "input", type=str, - help="Path to a single PDF file or a directory containing PDF files.", + help="Path to a PDF or .txt file, or a directory of PDFs or .txt files.", + ) + parser.add_argument( + "--skip-classifier", + action="store_true", + help="Skip XGBoost classification; process .txt files as useful (extract-from-txt mode).", ) parser.add_argument( "--model-dir", @@ -350,13 +784,13 @@ def main(): "--confidence-threshold", type=float, default=0.70, - help="Classifier probability threshold for 'useful' (default: 0.70).", + help="Classifier probability threshold for 'useful' (default: 0.70). Ignored with --skip-classifier.", ) parser.add_argument( "--max-chars", type=int, default=12000, - help="Max characters of text to send to the LLM (default: 12000).", + help="Max characters of text to send to the LLM after cleaning (default: 12000).", ) parser.add_argument( "--num-ctx", @@ -368,12 +802,43 @@ def main(): "--workers", type=int, default=1, - help="Number of parallel worker processes (default: 1 = sequential).", + help="Parallel worker processes for PDF mode only (default: 1 = sequential).", + ) + parser.add_argument( + "--labels", + type=str, + default=None, + help="Path to labels.json; only files labelled 'useful' are processed.", + ) + parser.add_argument( + "--chunked", + action="store_true", + default=False, + help=( + "Chunked extraction: split text, score chunks with XGBoost, " + "extract from top-N chunks, merge via majority voting." + ), + ) + parser.add_argument( + "--top-chunks", + type=int, + default=3, + help="Number of top-scoring chunks to extract from (default: 3). Only used with --chunked.", + ) + parser.add_argument( + "--chunk-size", + type=int, + default=4000, + help="Character size per chunk (default: 4000). Only used with --chunked.", + ) + parser.add_argument( + "--chunk-overlap", + type=int, + default=500, + help="Overlap between consecutive chunks (default: 500). Only used with --chunked.", ) args = parser.parse_args() - - # Configure persistent logging for this process — one call covers all modules setup_logging() input_path = Path(args.input) @@ -382,17 +847,53 @@ def main(): log.error("Input path not found: %s", input_path) sys.exit(1) - run_pipeline( - input_path=input_path, + useful_stems = _load_useful_stems(Path(args.labels)) if args.labels else None + + extract_kw = dict( + chunked=args.chunked, model_dir=args.model_dir, - llm_model=args.llm_model, - output_dir=Path(args.output_dir), - confidence_threshold=args.confidence_threshold, - max_chars=args.max_chars, - num_ctx=args.num_ctx, - workers=args.workers, + top_chunks=args.top_chunks, + chunk_size=args.chunk_size, + chunk_overlap=args.chunk_overlap, ) + if args.skip_classifier: + run_txt_pipeline( + input_path=input_path, + output_dir=Path(args.output_dir), + llm_model=args.llm_model, + max_chars=args.max_chars, + num_ctx=args.num_ctx, + useful_stems=useful_stems, + **extract_kw, + ) + else: + if input_path.suffix.lower() == ".txt": + print( + "[ERROR] .txt input requires --skip-classifier. " + "Use: python src/pipeline/classify_extract.py --skip-classifier", + file=sys.stderr, + ) + sys.exit(1) + if input_path.is_dir() and not list(input_path.glob("*.pdf")) and list(input_path.glob("*.txt")): + print( + "[ERROR] Directory contains .txt files only; use --skip-classifier for .txt extraction.", + file=sys.stderr, + ) + sys.exit(1) + run_pdf_pipeline( + input_path=input_path, + model_dir=args.model_dir, + llm_model=args.llm_model, + output_dir=Path(args.output_dir), + confidence_threshold=args.confidence_threshold, + max_chars=args.max_chars, + num_ctx=args.num_ctx, + workers=args.workers, + useful_stems=useful_stems, + **extract_kw, + ) + if __name__ == "__main__": main() diff --git a/src/pipeline/extract_from_txt.py b/src/pipeline/extract_from_txt.py deleted file mode 100644 index 9cee179..0000000 --- a/src/pipeline/extract_from_txt.py +++ /dev/null @@ -1,470 +0,0 @@ -"""Extract-from-TXT Pipeline - -Processes pre-classified useful .txt files through noise cleaning, section -filtering, text trimming, and LLM extraction — bypassing the XGBoost -classifier entirely. - -Every .txt file fed to this script is assumed to have already been confirmed -as useful (e.g. by the classifier in src/pipeline/classify_extract.py or by manual review). -The pipeline: - - 1. Read raw .txt file - 2. Strip noise (references, acknowledgements, affiliations, captions, …) - via src/io/text_cleaner.py - 3. Drop irrelevant paragraphs (taxonomy, morphometrics, stats methods, …) - via src/io/section_filter.py - 4. Trim to the character budget using section-priority ranking - via src/extraction/llm_text.py::extract_key_sections() - 5. Call Ollama for structured extraction via src/extraction/llm_client.py - 6. Save result JSON per file and a summary CSV - -Usage:: - - # Process the default directory (data/processed-text/) - python src/pipeline/extract_from_txt.py - - # Custom input directory - python src/pipeline/extract_from_txt.py --input-dir path/to/txt_files/ - - # Full options - python src/pipeline/extract_from_txt.py \\ - --input-dir data/processed-text/ \\ - --output-dir data/results/ \\ - --llm-model qwen2.5:7b \\ - --max-chars 10000 \\ - --num-ctx 8192 - -Output: - - data/cleaned-text/text_cleaner/_.txt noise-stripped text - - data/cleaned-text/section_filter/_.txt section-filtered text - - data/cleaned-text/llm_text/_.txt trimmed text passed to Ollama - - data/results/metrics/_results.json per file - - data/results/summaries/txt_pipeline_summary_.csv overall -""" - -import argparse -import csv -import sys -from datetime import datetime -from pathlib import Path - -from src.io.text_cleaner import clean_text -from src.io.section_filter import filter_relevant_sections -from src.extraction.llm_text import extract_key_sections -from src.extraction.llm_client import extract_metrics_from_text, save_extraction_result -from src.extraction.chunked_extraction import extract_with_chunking - - -# --------------------------------------------------------------------------- -# Core pipeline function -# --------------------------------------------------------------------------- - - -def run_txt_pipeline( - input_dir: Path, - output_dir: Path, - llm_model: str, - max_chars: int, - num_ctx: int, - single_file: Path = None, - useful_stems: set = None, - chunked: bool = False, - top_chunks: int = 3, - chunk_size: int = 4000, - chunk_overlap: int = 500, - model_dir: str = "src/classifier/models", -) -> None: - """Process every .txt file in *input_dir* through clean → filter → trim → extract. - - Args: - input_dir: Directory containing pre-classified useful .txt files. - Ignored when *single_file* is provided. - output_dir: Root output directory for JSON results and summary CSV. - llm_model: Ollama model name (e.g. ``"qwen2.5:7b"``). - max_chars: Character budget for the text sent to Ollama. - num_ctx: Context window size requested from Ollama. - single_file: If set, process only this one .txt file. - """ - if single_file is not None: - txt_paths = [single_file] - else: - txt_paths = sorted(input_dir.glob("*.txt")) - if useful_stems is not None: - txt_paths = [p for p in txt_paths if p.stem in useful_stems or p.name in useful_stems] - if not txt_paths: - print(f"[ERROR] No .txt files found in: {input_dir}", file=sys.stderr) - sys.exit(1) - - print(f"[INFO] Found {len(txt_paths)} .txt file(s) to process", file=sys.stderr) - output_dir.mkdir(parents=True, exist_ok=True) - cleaner_text_dir = output_dir.parent / "cleaned-text" / "text_cleaner" - filter_text_dir = output_dir.parent / "cleaned-text" / "section_filter" - llm_text_dir = output_dir.parent / "cleaned-text" / "llm_text" - cleaner_text_dir.mkdir(parents=True, exist_ok=True) - filter_text_dir.mkdir(parents=True, exist_ok=True) - llm_text_dir.mkdir(parents=True, exist_ok=True) - summary_rows = [] - - for idx, txt_path in enumerate(txt_paths, start=1): - print(f"\n[{idx}/{len(txt_paths)}] {txt_path.name}", file=sys.stderr) - - row: dict = { - "filename": txt_path.name, - "raw_chars": "", - "cleaned_chars": "", - "filtered_chars": "", - "trimmed_chars": "", - "extraction_status": "", - "species_name": "", - "study_location": "", - "study_date": "", - "sample_size": "", - "num_empty_stomachs": "", - "num_nonempty_stomachs": "", - "fraction_feeding": "", - } - - # ── Step 1: Read ──────────────────────────────────────────────────── - try: - raw_text = txt_path.read_text(encoding="utf-8") - except Exception as exc: - print(f" [ERROR] Could not read file: {exc}", file=sys.stderr) - row["extraction_status"] = "read_failed" - summary_rows.append(row) - continue - - row["raw_chars"] = len(raw_text) - print(f" [INFO] Raw size : {len(raw_text):,} chars", file=sys.stderr) - - if not raw_text.strip(): - print(f" [WARN] File is empty — skipping.", file=sys.stderr) - row["extraction_status"] = "empty_file" - summary_rows.append(row) - continue - - # ── Step 2: Clean ─────────────────────────────────────────────────── - cleaned = clean_text(raw_text) - row["cleaned_chars"] = len(cleaned) - print(f" [INFO] After clean: {len(cleaned):,} chars", file=sys.stderr) - - if not cleaned.strip(): - print(f" [WARN] Nothing left after cleaning — skipping.", file=sys.stderr) - row["extraction_status"] = "empty_after_clean" - summary_rows.append(row) - continue - - # ── Step 3: Save text_cleaner output ──────────────────────────────── - ts = datetime.now().strftime("%Y%m%d_%H%M%S") - cleaner_path = cleaner_text_dir / f"{txt_path.stem}_{ts}.txt" - try: - cleaner_path.write_text(cleaned, encoding="utf-8") - print(f" [INFO] Cleaner text : {cleaner_path.name}", file=sys.stderr) - except Exception as exc: - print(f" [WARN] Could not save cleaner text: {exc}", file=sys.stderr) - - # ── Step 4: Section filter ────────────────────────────────────────── - filtered = filter_relevant_sections(cleaned) - row["filtered_chars"] = len(filtered) - print(f" [INFO] After filter: {len(filtered):,} chars", file=sys.stderr) - - # ── Step 4b: Save section_filter output ───────────────────────────── - filter_path = filter_text_dir / f"{txt_path.stem}_{ts}.txt" - try: - filter_path.write_text(filtered, encoding="utf-8") - print(f" [INFO] Filter text : {filter_path.name}", file=sys.stderr) - except Exception as exc: - print(f" [WARN] Could not save filter text: {exc}", file=sys.stderr) - - # ── Step 5+6+7: Extract (chunked or single-pass) ──────────────────── - if chunked: - print(f" [INFO] Chunked extraction (top {top_chunks} chunks)…", file=sys.stderr) - row["trimmed_chars"] = "" - try: - merged = extract_with_chunking( - text=filtered, - model_dir=model_dir, - llm_model=llm_model, - num_ctx=num_ctx, - top_n=top_chunks, - chunk_size=chunk_size, - overlap=chunk_overlap, - ) - except Exception as exc: - print(f" [ERROR] Chunked extraction failed: {exc}", file=sys.stderr) - row["extraction_status"] = "extraction_failed" - summary_rows.append(row) - continue - - # save JSON - import json as _json - - metrics_dir = output_dir / "metrics" - metrics_dir.mkdir(parents=True, exist_ok=True) - out_path = metrics_dir / (txt_path.stem + "_results.json") - result_obj = { - "source_file": txt_path.name, - "file_type": txt_path.suffix.lower(), - "extraction_mode": "chunked", - "metrics": merged, - } - with open(out_path, "w", encoding="utf-8") as _f: - _json.dump(result_obj, _f, indent=2) - print(f"[SUCCESS] Results saved to {out_path}", file=sys.stderr) - - m = merged - else: - # ── Step 5: Trim to LLM budget ────────────────────────────────── - if len(filtered) > max_chars: - trimmed = extract_key_sections(filtered, max_chars) - print( - f" [INFO] After trim : {len(trimmed):,} chars " f"(budget {max_chars:,})", - file=sys.stderr, - ) - else: - trimmed = filtered - - row["trimmed_chars"] = len(trimmed) - - # ── Step 5b: Save llm_text output ─────────────────────────────── - llm_path = llm_text_dir / f"{txt_path.stem}_{ts}.txt" - try: - llm_path.write_text(trimmed, encoding="utf-8") - print(f" [INFO] LLM text : {llm_path.name}", file=sys.stderr) - except Exception as exc: - print(f" [WARN] Could not save LLM text: {exc}", file=sys.stderr) - - # ── Step 6: LLM extraction ────────────────────────────────────── - print(f" [INFO] Calling Ollama ({llm_model})…", file=sys.stderr) - try: - metrics = extract_metrics_from_text( - text=trimmed, - model=llm_model, - num_ctx=num_ctx, - ) - except Exception as exc: - print(f" [ERROR] Ollama extraction failed: {exc}", file=sys.stderr) - row["extraction_status"] = "extraction_failed" - summary_rows.append(row) - continue - - # ── Step 7: Save JSON ─────────────────────────────────────────── - try: - result = save_extraction_result( - metrics=metrics, - source_file=txt_path, - original_text=raw_text, - output_dir=output_dir, - ) - except Exception as exc: - print(f" [ERROR] Could not save result: {exc}", file=sys.stderr) - row["extraction_status"] = "save_failed" - summary_rows.append(row) - continue - - m = result["metrics"] - - row["extraction_status"] = "success" - row["species_name"] = m.get("species_name") or "" - row["study_location"] = m.get("study_location") or "" - row["study_date"] = m.get("study_date") or "" - row["sample_size"] = "" if m.get("sample_size") is None else m["sample_size"] - row["num_empty_stomachs"] = "" if m.get("num_empty_stomachs") is None else m["num_empty_stomachs"] - row["num_nonempty_stomachs"] = "" if m.get("num_nonempty_stomachs") is None else m["num_nonempty_stomachs"] - row["fraction_feeding"] = "" if m.get("fraction_feeding") is None else m["fraction_feeding"] - - print( - f" [OK] species={m.get('species_name')} " f"n={m.get('sample_size')} " f"date={m.get('study_date')}", - file=sys.stderr, - ) - - summary_rows.append(row) - - # ── Write summary CSV ─────────────────────────────────────────────────── - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - summaries_dir = output_dir / "summaries" - summaries_dir.mkdir(parents=True, exist_ok=True) - summary_path = summaries_dir / f"txt_pipeline_summary_{timestamp}.csv" - - fieldnames = [ - "filename", - "raw_chars", - "cleaned_chars", - "filtered_chars", - "trimmed_chars", - "extraction_status", - "species_name", - "study_location", - "study_date", - "sample_size", - "num_empty_stomachs", - "num_nonempty_stomachs", - "fraction_feeding", - ] - with open(summary_path, "w", newline="", encoding="utf-8") as f: - writer = csv.DictWriter(f, fieldnames=fieldnames) - writer.writeheader() - writer.writerows(summary_rows) - - # ── Final report ──────────────────────────────────────────────────────── - total = len(summary_rows) - succeeded = sum(1 for r in summary_rows if r["extraction_status"] == "success") - failed = total - succeeded - - print("\n" + "=" * 55, file=sys.stderr) - print("TXT EXTRACTION PIPELINE COMPLETE", file=sys.stderr) - print("=" * 55, file=sys.stderr) - print(f" Files processed : {total}", file=sys.stderr) - print(f" Successful : {succeeded}", file=sys.stderr) - print(f" Failed / skipped : {failed}", file=sys.stderr) - print(f" Summary CSV : {summary_path}", file=sys.stderr) - print("=" * 55, file=sys.stderr) - - -# --------------------------------------------------------------------------- -# CLI -# --------------------------------------------------------------------------- - - -def main() -> None: - parser = argparse.ArgumentParser( - description=("Extract structured predator-diet metrics from pre-classified " "useful .txt files using Ollama."), - formatter_class=argparse.RawDescriptionHelpFormatter, - epilog=""" -Examples: - Default (data/processed-text/ → data/results/): - python src/pipeline/extract_from_txt.py - - Custom directories: - python src/pipeline/extract_from_txt.py --input-dir data/useful-txt/ --output-dir out/ - - Different model / tighter budget: - python src/pipeline/extract_from_txt.py --llm-model mistral:7b --max-chars 4500 - """, - ) - parser.add_argument( - "--file", - type=str, - default=None, - help="Path to a single .txt file to process. Overrides --input-dir.", - ) - parser.add_argument( - "--input-dir", - type=str, - default="data/processed-text", - help="Directory of .txt files to process (default: data/processed-text).", - ) - parser.add_argument( - "--output-dir", - type=str, - default="data/results", - help="Root output directory for JSON results and CSV summary " "(default: data/results).", - ) - parser.add_argument( - "--llm-model", - type=str, - # default="qwen2.5:7b", - default="qwen2.5:7b", - help="Ollama model name (default: qwen2.5:7b).", - ) - parser.add_argument( - "--max-chars", - type=int, - default=10000, - help="Maximum characters to send to Ollama after cleaning (default: 10000).", - ) - parser.add_argument( - "--num-ctx", - type=int, - default=8192, - help="Ollama context window size (default: 8192).", - ) - parser.add_argument( - "--labels", - type=str, - default=None, - help="Path to labels.json. When provided, only files labelled 'useful' are processed.", - ) - parser.add_argument( - "--chunked", - action="store_true", - default=False, - help="Use chunked extraction: split text, score chunks with XGBoost, " "extract from top-N chunks, merge via majority voting.", - ) - parser.add_argument( - "--top-chunks", - type=int, - default=3, - help="Number of top-scoring chunks to extract from (default: 3). Only used with --chunked.", - ) - parser.add_argument( - "--chunk-size", - type=int, - default=4000, - help="Character size per chunk (default: 4000). Only used with --chunked.", - ) - parser.add_argument( - "--chunk-overlap", - type=int, - default=500, - help="Overlap between consecutive chunks (default: 500). Only used with --chunked.", - ) - parser.add_argument( - "--model-dir", - type=str, - default="src/classifier/models", - help="Directory containing XGBoost model artifacts (default: src/classifier/models). Only used with --chunked.", - ) - - args = parser.parse_args() - - # ── Load label filter ─────────────────────────────────────────────── - useful_stems = None - if args.labels: - import json - - labels_path = Path(args.labels) - if not labels_path.exists(): - print(f"[ERROR] Labels file not found: {labels_path}", file=sys.stderr) - sys.exit(1) - with open(labels_path, encoding="utf-8") as f: - labels = json.load(f) - useful_stems = {k for k, v in labels.items() if v == "useful"} - print(f"[INFO] Labels filter: {len(useful_stems)} useful papers", file=sys.stderr) - - single_file = None - if args.file: - single_file = Path(args.file) - if not single_file.exists(): - print(f"[ERROR] File not found: {single_file}", file=sys.stderr) - sys.exit(1) - if single_file.suffix.lower() != ".txt": - print(f"[ERROR] --file must point to a .txt file: {single_file}", file=sys.stderr) - sys.exit(1) - input_dir = single_file.parent - else: - input_dir = Path(args.input_dir) - if not input_dir.exists(): - print(f"[ERROR] Input directory not found: {input_dir}", file=sys.stderr) - sys.exit(1) - if not input_dir.is_dir(): - print(f"[ERROR] --input-dir must be a directory: {input_dir}", file=sys.stderr) - sys.exit(1) - - run_txt_pipeline( - input_dir=input_dir, - output_dir=Path(args.output_dir), - llm_model=args.llm_model, - max_chars=args.max_chars, - num_ctx=args.num_ctx, - single_file=single_file, - useful_stems=useful_stems, - chunked=args.chunked, - top_chunks=args.top_chunks, - chunk_size=args.chunk_size, - chunk_overlap=args.chunk_overlap, - model_dir=args.model_dir, - ) - - -if __name__ == "__main__": - main() From d0e0881b6ada385b63cba407dba088fa491db8f4 Mon Sep 17 00:00:00 2001 From: raymondcen Date: Thu, 21 May 2026 23:11:08 -0700 Subject: [PATCH 2/4] fixed NameError --- src/pipeline/classify_extract.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/pipeline/classify_extract.py b/src/pipeline/classify_extract.py index cfd3d4a..0ff2118 100644 --- a/src/pipeline/classify_extract.py +++ b/src/pipeline/classify_extract.py @@ -510,7 +510,7 @@ def _run_extraction_on_text( raw_text, max_chars, stem=source_file.stem, - save_intermediates=save_intermediates and not chunked, + save_intermediates=save_intermediates, cleaner_dir=cleaner_dir, filter_dir=filter_dir, llm_dir=llm_dir if not chunked else None, @@ -546,7 +546,7 @@ def _run_extraction_on_text( return row row["extraction_status"] = "success" - populate_metrics_row(row, metrics) + _populate_metrics_row(row, metrics) print( f" [OK] species={metrics.get('species_name')} " f"n={metrics.get('sample_size')} " @@ -896,4 +896,4 @@ def main(): if __name__ == "__main__": - main() + main() \ No newline at end of file From 503104e4950d62c4e754f74ec1c27c0abfce8d40 Mon Sep 17 00:00:00 2001 From: raymondcen Date: Thu, 21 May 2026 23:20:42 -0700 Subject: [PATCH 3/4] fix: test case calling python instead of sys --- data/processed-text/test.txt | 39 ++++++++++++++++++++++++++++++++++++ tests/test_pdf_extraction.py | 2 +- 2 files changed, 40 insertions(+), 1 deletion(-) diff --git a/data/processed-text/test.txt b/data/processed-text/test.txt index 59a85b5..eabbea0 100644 --- a/data/processed-text/test.txt +++ b/data/processed-text/test.txt @@ -1,8 +1,10 @@ +[PAGE 1] CS 331: Introduction to Artificial Intelligence Lecture 2: Agents Sandhya Saisubramanian +[PAGE 2] Announcements • Quiz 1 will be released on Oct 1 at 2 pm and will be active until Oct 3 11:59pm @@ -12,6 +14,7 @@ Announcements (instead of Thursday) 2 +[PAGE 3] Today’s Agenda What is an agent? What does it mean to design an intelligent agent? @@ -19,6 +22,7 @@ What does it mean to design an intelligent agent? How can these agents decide what to do/ how to act? Readings: Chapters 2 and 3 in course textbook +[PAGE 4] Few Real-World Examples At a high-level, all these systems perform the following: 1. Sense (perceive) @@ -26,6 +30,7 @@ At a high-level, all these systems perform the following: 3. Act (produce some output) 4 +[PAGE 5] Agent-Centric View of AI Reasoning Environment @@ -38,6 +43,7 @@ sensors and acts on that environment through actuators Agent 5 +[PAGE 6] Agent-Related Terms • Percept sequence (P): A complete history of everything the agent has ever perceived. Think of this as the state of the world from @@ -54,6 +60,7 @@ on a machine. implementation 6 +[PAGE 7] Example: Vacuum Cleaner Agent Sensors: Camera to detect obstacles and dirt @@ -61,6 +68,7 @@ Percept Sequence: Cleanliness status 7 +[PAGE 8] Example: Vacuum Cleaner Agent Percept Sequence Action @@ -93,6 +101,7 @@ wheels Agent function 8 +[PAGE 9] Agent Design and Implementation Before implementing an agent, we need to think carefully about: 1. @@ -109,6 +118,7 @@ How should the agent act? How will its behavior or performance be evaluated? 9 +[PAGE 10] PEAS Descriptions of Task Environments Performance, Environment, Actuators, Sensors Performance @@ -140,6 +150,7 @@ PEAS is the standard way of specifying an agent task environment before you even design the agent 10 +[PAGE 11] PEAS Descriptions of Task Environments Performance, Environment, Actuators, Sensors Performance @@ -166,6 +177,7 @@ answers Example: Medical Diagnosis Agent 11 +[PAGE 12] Properties of Environments Fully observable: can access complete state of environment at each point in time @@ -198,6 +210,7 @@ making/executing entities; cooperative or competitive 12 +[PAGE 13] How Should An Agent Act? Given: • A performance measure @@ -210,6 +223,7 @@ actions? How should the agent function be constructed? Intelligence emerges from how agents choose actions 13 +[PAGE 14] How Should An Agent Act? Given: • A performance measure @@ -223,6 +237,7 @@ Does not consider the agent’s knowledge, its percept sequence or the environment in which it is operating 14 +[PAGE 15] Rational Agent Rational agent: for each possible percept sequence, a rational agent should select an action that is expected to maximize its @@ -239,6 +254,7 @@ Actions agent can perform Agent’s percept sequence to date 15 +[PAGE 16] Types of Agents • Categorized based on their performance measure Types of @@ -253,6 +269,7 @@ Utility-directed agents 16 +[PAGE 17] Simple Reflex Agents 17 • Simple if-then: if condition then action; no internal model @@ -265,9 +282,11 @@ action ← RULE-ACTION[rule] return action Example: if temp < 68 → turn heater on +[PAGE 18] Simple Reflex Agents 18 +[PAGE 19] Simple Reflex Agents • Advantages: • Easy to implement @@ -278,6 +297,7 @@ observable • Infinite loops 19 +[PAGE 20] Model-based Reflex Agents • Maintain some internal state that keeps track of the part of the world it can’t see now (partial observability) @@ -295,9 +315,11 @@ The internal model (“map of which rooms are clean/dirty”) allows it to act sensibly despite partial observability. 20 +[PAGE 21] Model-based Reflex Agents 21 +[PAGE 22] Goal-directed Agents • Goal information guides agent’s actions (looks to the future) • Sometimes achieving goal is simple e.g. from a single action @@ -307,9 +329,11 @@ actions à accounts for future states Example: Navigation 22 +[PAGE 23] Goal-directed Agents 23 +[PAGE 24] Utility-directed Agents • What if there are many paths to the goal? à optimize trade-offs when multiple outcomes are possible @@ -324,9 +348,11 @@ Common types of utility definitions: Example: Self-driving car balancing speed vs. safety vs. comfort 24 +[PAGE 25] Utility-directed Agents 25 +[PAGE 26] Types of Agents • Categorized based on their performance measure • Simple reflux agents @@ -336,6 +362,7 @@ Types of Agents • Learning agents 26 +[PAGE 27] Learning Agents Successful agents split task of computing policy in 3 periods: 1. @@ -352,9 +379,11 @@ some emails, but improves by learning from user corrections (“mark as spam” / “not spam”) 27 +[PAGE 28] Learning Agents 28 +[PAGE 29] Learning Agents Think of this as outside the agent since you don’t @@ -363,6 +392,7 @@ the agent Maps percepts to actions 29 +[PAGE 30] Learning Agents Responsible for improving the agent’s behavior with experience @@ -376,6 +406,7 @@ percepts don’t tell the agent about its success/failure) 30 +[PAGE 31] In-Class Exercise: Designing Agents Design an intelligent vacuum cleaner for a house. To make it more interesting, we will analyze design choices for different @@ -388,6 +419,7 @@ At least need a model-based agent, if not more complex designs: must remember which rooms have been cleaned 31 +[PAGE 32] In-Class Exercise: Designing Agents Design an intelligent vacuum cleaner for a house. Scenario 3: Multi-room house with user-specific requests such as ”Clean the @@ -395,6 +427,7 @@ kitchen first and then the living room” Goal-based agents 32 +[PAGE 33] In-Class Exercise: Designing Agents Design an intelligent vacuum cleaner for a house. Scenario 4: Multi-room house and the agent must determine the order of @@ -405,6 +438,7 @@ must be cleaned such that the utility is maximized associated with each room) 33 +[PAGE 34] Designing Modern Agents 34 Example: @@ -416,6 +450,7 @@ based (uses an internal model of the environment) + learning Many of the modern AI agents are designed to learn from data, to adapt better to the environment and user +[PAGE 35] From Classical AI to Modern Agents 35 Agent Type @@ -451,6 +486,7 @@ AlphaGo, adaptive robots Improves behavior over time from data and feedback +[PAGE 36] Planning Agents • Need to plan a sequence of actions to reach the goal or complete a task @@ -460,6 +496,7 @@ some level of planning sequences 36 +[PAGE 37] Search Problem 37 • Perform a search in the solution space intelligently @@ -467,6 +504,7 @@ Search Problem Solution space: set of all possible solutions to the problem Solution +[PAGE 38] Formalizing the Search Problem • A finite set of states à set of all decision points • Initial state @@ -479,6 +517,7 @@ negative one-step cost of travelling from state s to s’. The cost function is only defined if s’ is a successor state of s. 38 +[PAGE 39] Formalizing the Search Problem: Example from Textbook 39 diff --git a/tests/test_pdf_extraction.py b/tests/test_pdf_extraction.py index 7d3e022..0ce7824 100644 --- a/tests/test_pdf_extraction.py +++ b/tests/test_pdf_extraction.py @@ -75,7 +75,7 @@ def test_main_cli(tmp_path): output_file = output_dir / "test.txt" result = subprocess.run( - ["python", "src/io/pdf_text_extraction.py", input_pdf], + [sys.executable, "src/io/pdf_text_extraction.py", input_pdf], capture_output=True, text=True, ) From b6f3300a74948d8a5288017551d174df04433448 Mon Sep 17 00:00:00 2001 From: raymondcen Date: Thu, 4 Jun 2026 01:13:19 -0700 Subject: [PATCH 4/4] reformat --- src/pipeline/classify_extract.py | 23 ++++++----------------- 1 file changed, 6 insertions(+), 17 deletions(-) diff --git a/src/pipeline/classify_extract.py b/src/pipeline/classify_extract.py index 0ff2118..700cd6f 100644 --- a/src/pipeline/classify_extract.py +++ b/src/pipeline/classify_extract.py @@ -548,9 +548,7 @@ def _run_extraction_on_text( row["extraction_status"] = "success" _populate_metrics_row(row, metrics) print( - f" [OK] species={metrics.get('species_name')} " - f"n={metrics.get('sample_size')} " - f"date={metrics.get('study_date')}", + f" [OK] species={metrics.get('species_name')} " f"n={metrics.get('sample_size')} " f"date={metrics.get('study_date')}", file=sys.stderr, ) return row @@ -631,9 +629,7 @@ def run_pdf_pipeline( row = future.result() except Exception as exc: print(f" [ERROR] Worker failed for {pdf_path.name}: {exc}", file=sys.stderr) - row = _new_summary_row( - pdf_path.name, include_classifier=True, include_preprocess=False - ) + row = _new_summary_row(pdf_path.name, include_classifier=True, include_preprocess=False) row["extraction_status"] = "worker_failed" summary_rows.append(row) else: @@ -732,10 +728,7 @@ def run_txt_pipeline( def main(): parser = argparse.ArgumentParser( - description=( - "Classify PDFs as useful/not-useful and extract structured diet metrics, " - "or extract from pre-classified .txt files with --skip-classifier." - ), + description=("Classify PDFs as useful/not-useful and extract structured diet metrics, " "or extract from pre-classified .txt files with --skip-classifier."), formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: @@ -814,10 +807,7 @@ def main(): "--chunked", action="store_true", default=False, - help=( - "Chunked extraction: split text, score chunks with XGBoost, " - "extract from top-N chunks, merge via majority voting." - ), + help=("Chunked extraction: split text, score chunks with XGBoost, " "extract from top-N chunks, merge via majority voting."), ) parser.add_argument( "--top-chunks", @@ -870,8 +860,7 @@ def main(): else: if input_path.suffix.lower() == ".txt": print( - "[ERROR] .txt input requires --skip-classifier. " - "Use: python src/pipeline/classify_extract.py --skip-classifier", + "[ERROR] .txt input requires --skip-classifier. " "Use: python src/pipeline/classify_extract.py --skip-classifier", file=sys.stderr, ) sys.exit(1) @@ -896,4 +885,4 @@ def main(): if __name__ == "__main__": - main() \ No newline at end of file + main()