Improved PDF Classification Model Using XGBoost

requirements.txt

@@ -13,3 +13,4 @@ google-auth-oauthlib
 google-auth-httplib2
 numpy
 google-api-python-client
+xgboost

src/model/pdf_classifier.py

@@ -1,6 +1,7 @@
 import argparse
 import joblib
 from pathlib import Path
+import xgboost as xgb
 import sys
 
 sys.path.append(str(Path(__file__).resolve().parents[2]))
@@ -9,32 +10,45 @@
 
 # Classify a single PDF as useful or not useful based on its text content.
 def classify_pdf(pdf_path, model_dir="src/model/models"):
-    model_path = Path(model_dir) / "pdf_classifier_model.pkl"
+    model_path = Path(model_dir) / "pdf_classifier.json"
     vectorizer_path = Path(model_dir) / "tfidf_vectorizer.pkl"
+    encoder_path = Path(model_dir) / "label_encoder.pkl"
 
-    # Load model and TF-IDF vectorizer
-    if not model_path.exists() or not vectorizer_path.exists():
-        print(f"[ERROR] Missing model or vectorizer in {model_dir}")
+    if not model_path.exists() or not vectorizer_path.exists() or not encoder_path.exists():
+        print(f"[ERROR] Missing model, encoder, or vectorizer in {model_dir}")
         return
 
-    model = joblib.load(model_path)
+    # Load model, encoder, and TF-IDF vectorizer
+    model = xgb.Booster()
+    model.load_model(str(model_path))
     vectorizer = joblib.load(vectorizer_path)
+    encoder = joblib.load(encoder_path)
 
     # Extract text from PDF
     text = extract_text_from_pdf(pdf_path)
     if not text.strip():
         print(f"[ERROR] No text extracted from {pdf_path}. Skipping classification.")
         return
 
-    # Transform and classify
+    # Transform text into vectorized TF-IDF format
     X_vec = vectorizer.transform([text])
-    prediction = model.predict(X_vec)[0]
 
-    prob = model.predict_proba(X_vec)[0]
-    confidence = max(prob)
+    # Wrap in DMatrix for XGBoost prediction
+    dtest = xgb.DMatrix(X_vec)
+    pred_prob = model.predict(dtest)[0]
+    pred_class = 1 if pred_prob >= 0.70 else 0
+
+    # Convert numeric class back into original label name
+    pred_label = encoder.inverse_transform([pred_class])[0]
+
+    if pred_class == 0:
+        confidence = 1 - pred_prob
+    else:
+        confidence = pred_prob
+
     print("\n=== PDF Classification Result ===")
-    print(f" File:        {Path(pdf_path).name}")
-    print(f" Prediction:  {prediction.capitalize()} ({confidence:.2%} confidence)")
+    print(f" File: {Path(pdf_path).name}")
+    print(f" Prediction: {pred_label} ({confidence:.2%} confidence)")
     print("=================================\n")
 
 

src/model/train_model.py

@@ -1,106 +1,128 @@
 """
-PDF FracFeedExtractor Classifier
+PDF FracFeedExtractor Classifier Model Training
 -------------------------
 
-This module uses scikit-learn to train and evaluate a text classification model
-that predicts whether a PDF is useful for predator diet data analysis.
-
-The model uses TF-IDF vectorization + Logistic Regression.
+This module trains a PDF classifier using TF-IDF vectorization + XGBoost
+to determine whether a PDF is useful for predator diet data analysis.
 """
 
 from pathlib import Path
 from sklearn.model_selection import train_test_split
 from sklearn.feature_extraction.text import TfidfVectorizer
-from sklearn.linear_model import LogisticRegression
+from sklearn.preprocessing import LabelEncoder
 from sklearn.metrics import classification_report, accuracy_score
+
+import xgboost as xgb
 import joblib
 import json
 import sys
 from collections import Counter
-import os
 
 
-# Load processed text files and their useful vs not useful labels.
+# Load training texts and their labels
 def load_labeled_data(data_dir="data/processed-text", labels_file="data/labels.json"):
-    # Load label mappings from JSON file
     data_dir = Path(data_dir)
+
+    # Load label dictionary
     with open(labels_file, "r", encoding="utf-8") as f:
         labels_map = json.load(f)
 
-    # Iterate through all .txt files in the processed-text directory
+    # Iterate through processed text files
     texts, labels, filenames = [], [], []
     for txt_file in data_dir.glob("*.txt"):
         fname = txt_file.name
-        # Only include files that have a corresponding label
         if fname in labels_map:
             with open(txt_file, "r", encoding="utf-8") as f:
                 texts.append(f.read())
                 labels.append(labels_map[fname])
                 filenames.append(fname)
         else:
             print(f"[WARN] No label found for {fname}, skipping.")
+
     return texts, labels, filenames
 
 
-# Trains and save a scikit-learn classification model.
+# Train an XGBoost text classifier with TF-IDF features.
 def train_pdf_classifier(texts, labels, output_dir="src/model/models"):
-    # Basic validations
+
+    # Ensure dataset is not empty
     if not texts or not labels:
-        print("[ERROR] No training samples found. Ensure PDFs were extracted to text before training.")
+        print("[ERROR] No training samples found.")
         return None
 
+    # Need at least 2 classes to learn a classifier
     class_counts = Counter(labels)
     if len(class_counts) < 2:
-        print(f"[ERROR] Need at least 2 classes to train classifier. Found: {class_counts}")
+        print("[ERROR] Need at least two classes.")
         return None
 
     if any(c < 2 for c in class_counts.values()):
-        print(f"[ERROR] Each class needs at least 2 samples for stratified split. Counts: {class_counts}")
+        print("[ERROR] Each class needs at least 2 samples.")
         return None
 
-    # Decide split strategy accommodating very small CI samples.
-    X_train = X_test = y_train = y_test = None
-    perform_eval = True
-    min_class = min(class_counts.values())
-    n = len(labels)
-    test_size = 0.2
-    tiny_context = n < 10 or min_class < 3 or os.environ.get("CI_TRAIN", "").lower() in {"1", "true", "yes"}
-    if tiny_context:
-        # For tiny datasets, attempt a 50/50 split; if that fails, fall back to train-only.
-        test_size = 0.5
     try:
-        X_train, X_test, y_train, y_test = train_test_split(texts, labels, test_size=test_size, random_state=42, stratify=labels)
+        # Stratified splitter ensures class balance in train/test sets
+        X_train, X_test, y_train, y_test = train_test_split(texts, labels, test_size=0.2, random_state=42, stratify=labels)
     except ValueError as e:
-        print(f"[WARN] Stratified split failed on tiny dataset: {e} — training on all data without evaluation.")
-        perform_eval = False
-        X_train = texts
-        y_train = labels
+        print(f"[ERROR] train_test_split failed: {e}")
+        return None
+
+    # Convert target labels into numeric form
+    enc = LabelEncoder()
+    y_train = enc.fit_transform(y_train)
+    y_test = enc.transform(y_test)
+
+    # Calculate scale_pos_weight for imbalanced classes
+    num_pos = sum(y_train)
+    num_neg = len(y_train) - num_pos
+    scale_pos_weight = num_neg / max(num_pos, 1)
+
+    # Build TF-IDF vectorizer for extracting text features
+    vectorizer = TfidfVectorizer(
+        max_features=10000,
+        stop_words="english",
+        ngram_range=(1, 3),
+    )
 
-    # Convert text data into TF-IDF features
-    vectorizer = TfidfVectorizer(max_features=10000, stop_words="english", ngram_range=(1, 3))
+    # Fit TF-IDF transformer on training data and apply to test data
     X_train_vec = vectorizer.fit_transform(X_train)
-    X_test_vec = vectorizer.transform(X_test) if perform_eval else None
-
-    # Initialize a logistic regression classifier
-    model = LogisticRegression(max_iter=100000, solver="liblinear")
-    model.fit(X_train_vec, y_train)  # Train the model
-
-    # Evaluate model performance
-    if perform_eval and X_test_vec is not None:
-        y_pred = model.predict(X_test_vec)
-        acc = accuracy_score(y_test, y_pred)
-        print("\n=== Model Evaluation ===")
-        print(f"Accuracy: {acc:.2f}")
-        print(classification_report(y_test, y_pred, digits=2))
-        print("========================\n")
-    else:
-        acc = None
-        print("[INFO] Trained on full tiny dataset (no evaluation split).")
-
-    # Save model and vectorizer for future use
+    X_test_vec = vectorizer.transform(X_test)
+
+    # Convert to DMatrix for XGBoost
+    dtrain = xgb.DMatrix(X_train_vec, label=y_train)
+    dtest = xgb.DMatrix(X_test_vec, label=y_test)
+
+    # XGBoost parameters
+    params = {
+        "objective": "binary:logistic",  # binary classification
+        "eval_metric": "logloss",  # log loss metric
+        "eta": 0.05,  # learning rate
+        "max_depth": 6,
+        "subsample": 0.8,  # use 80% of data per boosting round
+        "alpha": 1.0,  # L1 regularization
+        "lambda": 1.0,  # L2 regularization
+        "scale_pos_weight": scale_pos_weight,
+    }
+
+    # Train the model
+    model = xgb.train(params, dtrain, num_boost_round=500, evals=[(dtrain, "train"), (dtest, "eval")], early_stopping_rounds=20, verbose_eval=True)  # stop if no improvement for 20 rounds
+
+    # Predict on test set and convert probabilities to labels
+    y_pred_prob = model.predict(dtest)
+    y_pred = [1 if p >= 0.5 else 0 for p in y_pred_prob]
+
+    # Evaluate accuracy
+    acc = accuracy_score(y_test, y_pred)
+    print("\n=== Model Evaluation ===")
+    print(f"Accuracy: {acc:.2f}")
+    print(classification_report(y_test, y_pred, digits=2))
+    print("========================\n")
+
+    # Save artifacts
     Path(output_dir).mkdir(parents=True, exist_ok=True)
-    joblib.dump(model, Path(output_dir) / "pdf_classifier_model.pkl")
+    model.save_model(str(Path(output_dir) / "pdf_classifier.json"))
     joblib.dump(vectorizer, Path(output_dir) / "tfidf_vectorizer.pkl")
+    joblib.dump(enc, Path(output_dir) / "label_encoder.pkl")
 
     return {"accuracy": acc}
 

tests/test_pdf_classifier.py

@@ -1,28 +1,41 @@
 import pytest
 import joblib
 from pathlib import Path
-from sklearn.linear_model import LogisticRegression
 import numpy as np
+import xgboost as xgb
 from sklearn.feature_extraction.text import TfidfVectorizer
-from sklearn.feature_extraction.text import CountVectorizer
-from unittest.mock import patch, MagicMock
+from sklearn.preprocessing import LabelEncoder
+from unittest.mock import patch
 from src.model.pdf_classifier import classify_pdf
 
 
 @pytest.fixture
 def model_dir_with_mock_model(tmp_path):
+    """Prepare a minimal XGBoost model, vectorizer, and label encoder on disk."""
     model_dir = tmp_path / "models"
     model_dir.mkdir(parents=True)
 
+    # Fake small training dataset
     texts = ["predator stomach content", "fish prey analysis", "rock study geology", "mineral content paper"]
     labels = ["useful", "useful", "not useful", "not useful"]
 
-    vectorizer = TfidfVectorizer(max_features=4)
+    # Encode labels
+    enc = LabelEncoder()
+    y = enc.fit_transform(labels)
+
+    # Vectorize text
+    vectorizer = TfidfVectorizer(max_features=6)
     X = vectorizer.fit_transform(texts)
-    model = LogisticRegression().fit(X, labels)
 
-    joblib.dump(model, model_dir / "pdf_classifier_model.pkl")
+    # Train tiny XGBoost model
+    dtrain = xgb.DMatrix(X, label=y)
+    params = {"objective": "binary:logistic", "eval_metric": "logloss"}
+    model = xgb.train(params, dtrain, num_boost_round=3)
+
+    # Save artifacts
+    model.save_model(str(model_dir / "pdf_classifier.json"))
     joblib.dump(vectorizer, model_dir / "tfidf_vectorizer.pkl")
+    joblib.dump(enc, model_dir / "label_encoder.pkl")
 
     return model_dir
 
@@ -39,11 +52,11 @@ def test_classify_pdf_valid_case(mock_extract, model_dir_with_mock_model, capsys
 
 
 def test_classify_pdf_missing_model(capsys, tmp_path):
-    model_dir = tmp_path / "empty_models"
+    """Verifier missing model directory gracefully errors."""
+    model_dir = tmp_path / "empty"
     model_dir.mkdir()
 
-    test_pdf = Path("tests/test.pdf")
-    classify_pdf(test_pdf, model_dir)
+    classify_pdf(Path("tests/test.pdf"), model_dir)
 
     output = capsys.readouterr().out
     assert "[ERROR]" in output
@@ -52,30 +65,24 @@ def test_classify_pdf_missing_model(capsys, tmp_path):
 
 @patch("src.model.pdf_classifier.extract_text_from_pdf", return_value="")
 def test_classify_pdf_no_text(mock_extract, model_dir_with_mock_model, capsys):
-    test_pdf = Path("tests/empty.pdf")
-    classify_pdf(test_pdf, model_dir_with_mock_model)
+    classify_pdf(Path("tests/empty.pdf"), model_dir_with_mock_model)
 
     output = capsys.readouterr().out
     assert "[ERROR]" in output
     assert "No text extracted" in output
 
 
 @patch("src.model.pdf_classifier.extract_text_from_pdf", return_value="diet prey stomach analysis")
-def test_classify_pdf_confidence_and_prediction(mock_extract, model_dir_with_mock_model, capsys):
-    test_pdf = Path("tests/valid.pdf")
-    classify_pdf(test_pdf, model_dir_with_mock_model)
+def test_classify_pdf_prediction_output(mock_extract, model_dir_with_mock_model, capsys):
+    classify_pdf(Path("tests/valid.pdf"), model_dir_with_mock_model)
 
     output = capsys.readouterr().out
     assert "Prediction:" in output
-    assert "Confidence" not in output
-    assert "useful" in output.lower()
-
+    assert "%" in output
 
-def test_classify_pdf_handles_nonexistent_model_dir(tmp_path, capsys):
-    model_dir = tmp_path / "nonexistent_dir"
-    test_pdf = Path("tests/missing_model.pdf")
 
-    classify_pdf(test_pdf, model_dir)
+def test_classify_pdf_handles_nonexistent_dir(tmp_path, capsys):
+    classify_pdf(Path("tests/missing.pdf"), tmp_path / "nope")
 
     output = capsys.readouterr().out
     assert "[ERROR]" in output