ml-engineer

name

ml-engineer

description

Train, evaluate, debug, and extend GregoryAI's machine learning pipeline for article relevance prediction. Use this skill whenever someone asks about training models, running predictions, debugging ML issues, adding new algorithms, understanding prediction scores, pseudo-labeling, model versioning, or the ML consensus system. Triggers include questions like "train models for team X", "why is recall low", "add a new algorithm", "how does pseudo-labeling work", "run predictions", "check model metrics", "what does the ML pipeline do", or any work involving the gregory.ml module, train_models / predict_articles management commands, or the files_repo_PBL_nsbe research notebooks.

ML Engineer for GregoryAI

Purpose

This skill covers the full ML lifecycle in GregoryAI: data collection, text preprocessing, model training (with optional pseudo-labeling), evaluation, model versioning, prediction, and the consensus system that determines article relevance.

Step 0 — Always read the source first

Before writing or modifying any ML code, read the relevant source files:

# Core ML module (wrappers, factory, pseudo-labeling, GPU config)
django/gregory/ml/__init__.py
django/gregory/ml/trainer.py
django/gregory/ml/bert_wrapper.py
django/gregory/ml/lgbm_wrapper.py
django/gregory/ml/lstm_wrapper.py
django/gregory/ml/pseudo.py
django/gregory/ml/gpu_config.py

# Utilities (dataset building, text cleaning, summarisation, versioning, metrics)
django/gregory/utils/dataset.py
django/gregory/utils/text_utils.py
django/gregory/utils/summariser.py
django/gregory/utils/versioning.py
django/gregory/utils/metrics.py
django/gregory/utils/verboser.py
django/gregory/utils/model_utils.py

# Management commands (orchestration)
django/gregory/management/commands/train_models.py
django/gregory/management/commands/predict_articles.py
django/gregory/management/commands/pipeline.py

# Models (MLPredictions, PredictionRunLog, ArticleSubjectRelevance, Subject consensus fields)
django/gregory/models.py

# Research notebooks and original algorithm implementations
files_repo_PBL_nsbe/

Read the files that are relevant to the task before making changes.

Architecture overview

Three algorithms

Alias	Class	Module	Key artifacts
pubmed_bert	BertTrainer	gregory.ml.bert_wrapper	bert_weights.h5
lgbm_tfidf	LGBMTfidfTrainer	gregory.ml.lgbm_wrapper	tfidf_vectorizer.joblib, lgbm_classifier.joblib
lstm	LSTMTrainer	gregory.ml.lstm_wrapper	lstm.weights.h5, tokenizer.json

All three share a common interface:

trainer = get_trainer(algo, **kwargs)          # factory in gregory.ml.trainer
history = trainer.train(train_texts, train_labels, val_texts, val_labels, epochs, batch_size)
metrics = trainer.evaluate(test_texts, test_labels, threshold)
preds, probs = trainer.predict(texts, threshold)
trainer.save(model_dir)
trainer.load(model_dir)   # or load_weights() for BERT
trainer.export_metrics_json(val_metrics, test_metrics, output_path)

Model storage

Artifacts live at:

BASE_DIR/models/{team_slug}/{subject_slug}/{algorithm}/{version}/

Version format: YYYYMMDD, auto-incrementing _2, _3 etc. if a version already exists for the same day. Generated by gregory.utils.versioning.make_version_path().

Each version directory contains model weights/files plus metrics.json.

Key models in the database

• MLPredictions — stores per-article, per-subject, per-algorithm prediction results (predicted_relevant, probability_score, algorithm, subject, model_version)
• PredictionRunLog — audit trail for every training and prediction run (run_type='train'|'predict', team, subject, algorithm, model_version, metrics JSON, articles_processed)
• ArticleSubjectRelevance — manual human labels (is_relevant boolean per article per subject)
• Subject.ml_consensus_type — determines how many algorithms must agree: 'any' (1), 'majority' (2), 'all' (3)
• Subject.auto_predict — whether automated predictions are enabled for this subject

Training pipeline

Managed by python manage.py train_models. The steps:

1. Collect articles — utils.dataset.collect_articles(team_slug, subject_slug) filters Articles by team and subject, returns only those with human relevance labels.

2. Build dataset — utils.dataset.build_dataset(queryset) concatenates title + summary into a 'text' column, maps is_relevant to 0/1.

3. Summarise — utils.summariser.summarise_bulk(texts) uses facebook/bart-large-cnn with disk caching (MD5-keyed). Used during training to compress long texts.

4. Split — utils.dataset.train_val_test_split(df) does stratified 70/15/15 split. Falls back to non-stratified if fewer than 3 samples per class.

5. Pseudo-labeling (optional) — gregory.ml.pseudo.generate_pseudo_labels() runs iterative self-training: train on labeled data, predict unlabeled, add confident predictions (default threshold 0.9) back to training set, repeat up to max_iter=7.

6. Train — get_trainer(algo).train(...) with algorithm-specific defaults.

7. Evaluate — trainer.evaluate(test_texts, test_labels, threshold) returns accuracy, precision, recall, f1, roc_auc, pr_auc, confusion_matrix.

8. Save — trainer.save(version_dir) writes model artifacts + metrics.json.

9. Log — Creates a PredictionRunLog entry with run_type='train'.

Command flags

python manage.py train_models \
  --team <slug>           # or --all-teams
  --subject <slug>        # required unless --all-teams
  --algo pubmed_bert      # pubmed_bert | lgbm_tfidf | lstm (default: all three)
  --pseudo-label          # enable pseudo-labeling
  --prob-threshold 0.8    # prediction threshold (default 0.8)
  --model-version v1      # override auto-version
  --verbose 2             # 0=quiet, 1=progress, 2=warnings, 3=summary
  --debug                 # extra debug output
  --cpu                   # force CPU mode

Prediction pipeline

Managed by python manage.py predict_articles.

1. Resolve model — finds the latest version directory for the given team/subject/algorithm.

2. Load model — algorithm-specific loading logic:

   • BERT: rebuild architecture then load_weights(bert_weights.h5)
   • LGBM: joblib.load() for vectorizer and classifier
   • LSTM: rebuild architecture, load weights + tokenizer JSON

3. Select articles — unpredicted articles for this subject/algorithm combination, within the lookback window (default 90 days, or --all-articles).

4. Prepare text — cleanHTML(title + ' ' + summary) then cleanText().

5. Predict — trainer.predict(texts, threshold) returns labels and probabilities.

6. Store — MLPredictions.objects.bulk_create() with article, subject, algorithm, predicted_relevant, probability_score, model_version.

7. Log — Creates a PredictionRunLog entry with run_type='predict'.

Command flags

python manage.py predict_articles \
  --team <slug>           # or --all-teams
  --subject <slug>
  --algo pubmed_bert
  --prob-threshold 0.8
  --dry-run               # don't write predictions to DB
  --lookback-days 90      # how far back to look for unpredicted articles
  --all-articles           # ignore lookback, predict everything
  --verbose 2
  --cpu

Consensus system

An article is considered "ML-relevant" for a subject when enough algorithms agree. The logic lives in api/filters.py → ArticleFilter._get_ml_relevant_articles_query():

1. For each Subject with auto_predict=True, read its ml_consensus_type.
2. Map consensus type to minimum algorithm count: any→1, majority→2, all→3.
3. Find articles where at least that many distinct algorithms predicted relevant with probability_score >= threshold.
4. Union with manually-relevant articles (ArticleSubjectRelevance.is_relevant=True).

The relevant=true filter parameter on the articles API triggers this logic. The ml_threshold parameter overrides the default 0.8 probability threshold.

Text preprocessing

Two layers of cleaning in gregory.utils.text_utils:

• cleanHTML(text) — strips HTML tags via BeautifulSoup
• cleanText(text, min_words=10) — lowercases, removes DOI/PMID patterns, removes stopwords, filters texts shorter than min_words

During training, summarise_bulk() compresses texts using BART-large-CNN. Summaries are cached to disk (MD5 hash key) so repeated training runs skip already-summarised articles.

GPU configuration

gregory.ml.gpu_config handles platform detection:

• configure_gpu_memory_growth() — enables TF memory growth (Apple Silicon + NVIDIA)
• disable_gpu() — force CPU via environment variables
• get_device_info() — returns dict with platform, GPU count, device list

The --cpu flag on management commands calls disable_gpu().

Algorithm-specific details

PubMed BERT (pubmed_bert)

• Base model: microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext
• Architecture: BERT → CLS token → Dense(48, ReLU, L2=0.01) → Dropout(0.3) → Dense(2, softmax)
• Default max_len: 400, learning_rate: 2e-5, batch_size: 16
• Early stopping: patience=3 on val_loss
• Supports pseudo-labeling via perform_pseudo_labeling() method
• Saves as HDF5 weights (bert_weights.h5)

LightGBM + TF-IDF (lgbm_tfidf)

• TF-IDF: max_features=10000, min_df=2, max_df=0.95, ngram_range=(1,2)
• LightGBM: binary objective, gbdt boosting, 31 leaves, lr=0.05
• Early stopping: 10 rounds on validation loss
• Saves via joblib (vectorizer + classifier as separate files)

LSTM (lstm)

• TextVectorization: max_tokens=10000, sequence_length=100
• Architecture: Embedding(128) → Bidirectional LSTM(64) → Dropout(0.3) → Dense(1, sigmoid)
• Early stopping: patience=3 on val_loss
• Saves weights as .weights.h5 (Keras 3.x), tokenizer config as JSON

Research notebooks

The files_repo_PBL_nsbe/ directory contains the original research from the NovaSBE collaboration. Key resources:

• model_documentation/ — detailed docs for each algorithm class
• notebooks_to_run/training_model.ipynb — end-to-end training notebook
• notebooks_to_run/classification_of_articles.ipynb — inference notebook
• notebooks_to_run/RNN_models_training_history/ — LSTM experiment notebooks
• code_utils/ — standalone implementations of all algorithms, pseudo-labeling, text processing

The Django wrappers (gregory.ml.*) are the production versions of these notebook implementations. When adding new algorithms or modifying existing ones, check both the wrapper and the original notebook/utility to understand the design intent.

Rules when modifying ML code

1. Preserve the common trainer interface — any new algorithm must implement: train(), evaluate(), predict(), save(), load(), export_metrics_json(). Register it in gregory.ml.trainer.get_trainer().

2. Always log runs — create PredictionRunLog entries for both training and prediction.

3. Use the versioning system — call make_version_path() for artifact storage. Never overwrite existing version directories.

4. Respect the consensus system — new algorithms participate in the consensus count via MLPredictions records. Update consensus thresholds in Subject model if changing the algorithm count.

5. Cache summaries — use summarise_bulk() with caching enabled during training to avoid redundant BART inference.

6. Handle GPU gracefully — use gpu_config.configure_gpu_memory_growth() at startup. Support --cpu flag for environments without GPU.

7. Threshold consistency — the default probability threshold is 0.8 across training, prediction, and API filtering. Keep this consistent unless explicitly overridden.

8. Test with small data first — use --verbose 3 and limit to a single team/subject/algorithm when debugging.

9. Metrics JSON — every saved model must include a metrics.json with val and test metrics, model parameters, and training metadata.

10. Text preparation — always apply cleanHTML() then cleanText() before prediction. During training, also apply summarise_bulk().

Debugging checklist

When predictions seem wrong or training fails:

• Check PredictionRunLog for recent runs and their metrics
• Verify the model version directory exists and contains all expected artifacts
• Check Subject.auto_predict is True and ml_consensus_type is set
• Verify enough human labels exist (ArticleSubjectRelevance records)
• Check class balance in the training dataset (needs both relevant and not-relevant)
• Look at summarise_bulk() cache stats — stale cache can cause issues
• Verify GPU config with get_device_info()
• Check the threshold: training, prediction, and API filter must use the same value