Improved PDF Classification Model Using XGBoost

[SeanClay10]

Updated the document classification pipeline to replace the previous model with an XGBoost binary classifier trained on TF-IDF features. The update improves predictive accuracy, probability calibration, and robustness on small and imbalanced datasets. The classification script, training pipeline, and test suite were all updated to load XGBoost models, probability outputs, and confidence reporting. Pytest unit tests were enhanced to validate model loading, training behavior, and expected outcomes.

Key Enhancements

1. Migration to XGBoost (binary:logistic)

• Implemented an XGBoost classifier configured for binary document usefulness prediction.
• The model now outputs a calibrated probability representing confidence that a document is “useful.”
• The new approach handles sparse TF-IDF vectors efficiently and generalizes better on small datasets compared to previous linear models.

2. Improved Training Pipeline
   Reworked train_model.py to:

• Load raw training text and labels.
• Encode labels using LabelEncoder.
• Convert documents into TF-IDF features using up to 10,000 weighted terms across unigrams, bigrams, and trigrams.
• Create XGBoost DMatrices for efficient GPU/CPU training.
• Stratified train/test splitting ensures evaluation remains balanced and representative.

[raymondcen]

Lower iterations. 1000 iterations with max depth of 6 is too much and model will overfit and training will take too long with 600+ files unless we add multi threading. If not you can also add an early stopping parameter to it there's probably something equivalent to that with XGBoost. It finishes the training immediately if it doesn't improve after k rounds.

Add 'subsample' parameter to also mitigate memorization.

Consider adding L1 and L2 parameter or something equivalent. I forget which but one shrinks down confident predictions and another removes weak predictions. They basically help in overfitting. Makes it focus on more data driven patterns.

Small things:
You can add the 'verbose' parameter onto it an it should show/print the progress it's training at every round. That way you can also see which parmeter is working and which is not.

You can also add 'scale_pos_weight' parameter to it when we get a more balanced data set.

Other than those things it looks good. The parameters are mainly for the future when we even out the 'not useful' class but they would still be helpful currently (maybe not the scale_pos_weight though).