brunoamaral/gregory-ai/ml-engineer

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

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.

1. Build dataset — utils.dataset.build_dataset(queryset) concatenates title + summary

into a 'text' column, maps is_relevant to 0/1.

1. Summarise — utils.summariser.summarise_bulk(texts) uses facebook/bart-large-cnn

with disk caching (MD5-keyed). Used during training to compress long texts.

1. 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.

1. 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.

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

1. Evaluate — trainer.evaluate(test_texts, test_labels, threshold) returns

accuracy, precision, recall, f1, roc_auc, pr_auc, confusion_matrix.

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

1. 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.

1. 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

1. Select articles — unpredicted articles for this subject/algorithm combination,

within the lookback window (default 90 days, or --all-articles).

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

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

1. Store — MLPredictions.objects.bulk_create() with article, subject, algorithm,

predicted_relevant, probability_score, model_version.

1. 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.

1. 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().

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

1. Use the versioning system — call make_version_path() for artifact storage.

Never overwrite existing version directories.

1. 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.

1. Cache summaries — use summarise_bulk() with caching enabled during training

to avoid redundant BART inference.

1. Handle GPU gracefully — use gpu_config.configure_gpu_memory_growth() at

startup. Support --cpu flag for environments without GPU.

1. Threshold consistency — the default probability threshold is 0.8 across

training, prediction, and API filtering. Keep this consistent unless explicitly overridden.

1. Test with small data first — use --verbose 3 and limit to a single

team/subject/algorithm when debugging.

1. Metrics JSON — every saved model must include a metrics.json with val and

test metrics, model parameters, and training metadata.

1. Text preparation — always apply cleanHTML() then cleanText() before

prediction. During training, also apply summarise_bulk().

Debugging checklist

When predictions seem wrong or training fails: