API Reference#
Configuration#
Config#
Dataclass holding all training and model hyperparameters.
Parameter |
Default |
Description |
|---|---|---|
Training |
||
|
|
Random seed for reproducibility |
|
|
Training batch size |
|
|
Evaluation/inference batch size |
|
|
Maximum training epochs |
|
|
AdamW learning rate |
|
|
Epochs without val loss improvement before stopping |
Architecture |
||
|
|
Dropout probability in MLP layers |
|
|
Hidden layer width |
|
|
Number of hidden layers in MLP |
|
|
Max token length for tokenizer |
Optimization |
||
|
|
Max gradient norm for clipping |
|
|
Linear warmup steps |
Embedding |
||
|
|
Encoder type ( |
|
|
HuggingFace model ID (auto-resolved from |
|
|
|
Uncertainty |
||
|
|
Below this probability, flag as uncertain |
|
|
Above this probability, flag as uncertain |
Cross-validation |
||
|
|
Number of stratified CV folds |
|
|
Post-hoc calibration method |
Active learning |
||
|
|
Query strategy ( |
|
|
Records per AL iteration |
|
|
Initial random sample fraction (benchmark simulations fix this at 1%) |
Categorical encoding |
||
|
|
Minimum count for a categorical value to get its own embedding |
Class weighting |
||
|
|
|
SAFE stopping |
||
|
|
Consecutive irrelevant records before stopping |
|
|
Minimum fraction screened before stopping allowed |
|
|
Random sample fraction for recall estimation |
|
|
Switch model during screening phases |
Preset Configurations#
Preset |
Description |
|---|---|
|
Balanced defaults for general use |
|
Fewer epochs, larger batch size for quick experiments |
|
More epochs, lower learning rate for production |
|
Human-in-the-loop screening settings |
|
Domain-specific presets (science, medicine, general, modernbert) |
sentence_transformer_models#
Set of models that use the SentenceTransformer encoder (frozen, no fine-tuning).
Data#
preprocess_dataset#
Tokenize text columns and encode categorical/numeric features into a dataset. Pass fitted_transforms=None to fit from training data, or a FittedTransforms object to reuse on val/test sets. Returns (CustomDataset, FittedTransforms).
column_specifications#
Dictionary specifying which DataFrame columns to use.
Key |
Value |
Example |
|---|---|---|
|
List of text column names |
|
|
List of categorical column names |
|
|
List of numeric column names |
|
|
String (single-label) or list (multi-label) |
|
Numeric Transforms#
Transform |
Description |
|---|---|
|
Subtract minimum value |
|
Divide by maximum value |
|
Subtract mean |
|
Quantile transform to normal distribution |
|
RobustScaler (median + IQR), better for outliers |
|
log1p then quantile transform, for skewed features |
load_data#
Load data from CSV or Excel files.
split_data#
Split a DataFrame into train, validation, and test sets.
FittedTransforms#
Stores fitted parameters from training data for reuse on val/test sets. Serialize with to_dict() and restore with FittedTransforms.from_dict(d).
CustomDataset#
Dataset holding tokenized text, categorical/numeric features, and labels.
CachedEmbeddingDataset#
Dataset of precomputed sentence embeddings + tabular features (cached-embedding fast path).
create_dataloader#
Create DataLoader with custom collate function.
collate_fn#
Custom collate to handle text lists in batches.
Model#
PubMLP#
Multi-layer perceptron that combines transformer embeddings with categorical and numeric features.
Parameter |
Description |
|---|---|
|
List of vocab sizes for nn.Embedding per categorical column |
|
1 for single-label, N for multi-label |
Training & Evaluation#
train_evaluate_model#
Full training loop with validation, early stopping, and test evaluation. Returns (train_losses, val_losses, train_accs, val_accs, test_acc, best_val_loss, best_model_state, best_epoch).
calculate_loss#
Average loss across all batches.
calculate_accuracy#
Accuracy (%) across all batches; multi-label returns average per-label accuracy.
calculate_pos_weight#
Compute pos_weight from label distribution (neg_count / pos_count per label).
calculate_evaluation_metrics#
Compute classification report, confusion matrix, and ROC-AUC. Single label: returns accuracy, precision, recall, specificity, f1_score, roc_auc. Multi-label: returns per_label metrics, macro_f1, hamming_loss.
calculate_wss_at_recall#
WSS@recall (Cohen et al., 2006): fraction of screening effort saved at target recall.
calculate_ndcg#
NDCG via sklearn.metrics.ndcg_score (Järvelin & Kekäläinen, 2002).
calculate_ece#
Expected calibration error with equal-width probability bins (default 10). Multi-label (2-D) input returns the per-label macro average. Pair with calibrate_model to report calibration quality before and after temperature scaling.
calculate_brier#
Brier score (Brier, 1950): mean squared error of predicted probabilities. Multi-label (2-D) input returns the per-label macro average.
cross_validate#
Stratified K-fold cross-validation with per-fold metrics.
plot_results#
Plot training/validation loss and accuracy curves.
plot_al_progress#
Plot active learning learning curve (Macro F1 vs labeled instances), with optional per-criterion overlays.
TemperatureScaling#
Post-hoc temperature scaling for model calibration.
calibrate_model#
Fit temperature scaling on validation data.
collect_logits#
Collect raw logits from a trained model.
Prediction#
predict_model#
Run inference and return predictions and probabilities. Single label: flat lists. Multi-label: list of lists.
get_predictions_and_labels#
Get predictions, probabilities, and true labels from a labeled dataloader.
flag_uncertain#
Flag predictions with probabilities in an uncertain range for human review. Multi-label: returns list of lists of bools.
Screening#
regex_screen#
Screen a dataset using regex patterns with optional semantic similarity scoring.
from pubmlp import regex_screen
results = regex_screen("records.csv", inclusion_patterns=["intervention", "randomized"])
extract_window_evidence#
Extract word windows around regex matches.
extract_sentence_evidence#
Extract complete sentences containing regex matches.
extract_all_evidence#
Extract evidence from specified fields in a DataFrame row.
format_evidence_display#
Format evidence list as ‘field: text; field: text; …’.
calculate_semantic_scores#
Calculate cosine similarity between evidence texts and criterion description.
generate_descriptions#
Draft criterion descriptions from regex pattern terms. Extracts literal terms from each pattern and composes a natural language description. The user should review and refine each description before use.
from pubmlp import generate_descriptions
patterns = {'math': {'pattern': r'\b(algebra|geometry)\w*\b'}}
drafts = generate_descriptions(patterns, domain='K-12 education')
# drafts['math']['description'] → "In K-12 education, the study addresses math..."
# drafts['math']['source'] → 'generated'
confirm_descriptions#
Validate that all criteria have non-empty descriptions and return confirmed patterns. Optionally saves to JSON for reproducibility. Raises ValueError if any description is empty.
from pubmlp import confirm_descriptions
# After user edits drafts['math']['description']
confirmed = confirm_descriptions(drafts, save_path='confirmed.json')
# Returns dict ready for regex_screen() and score_full_text()
score_full_text#
Score full record text (title + abstract) against each criterion description via cosine similarity. Scores all records including those with no regex match. Adds {criterion}_semantic_full column per criterion.
from pubmlp import score_full_text
df = score_full_text(df, confirmed_patterns, fields=['title', 'abstract'])
# df['math_semantic_full'] → cosine similarity for every record
compare_screening_configs#
Run regex_screen with multiple configurations (different descriptions, window sizes, units) and return a summary DataFrame comparing match counts, semantic score distributions, and overlap.
from pubmlp import compare_screening_configs
comparison = compare_screening_configs('data.xlsx', {
'specific': {'inclusion_patterns': patterns_v1},
'broad': {'inclusion_patterns': patterns_v2, 'unit': 'window', 'window_size': 10},
})
# Returns DataFrame: config, criterion, n_matched, match_pct, semantic medians
create_stratified_sample#
Create a stratified random sample with regex pattern highlights for human coding.
save_sample_excel#
Save sample to Excel with conditional formatting for review.
apply_conditional_formatting#
Apply conditional formatting to Excel coding sheet (headers green, pattern counts yellow).
count_pattern_matches#
Count regex matches in text (case-insensitive).
highlight_pattern_matches#
Return up to 3 matched snippets with context for visual inspection.
Active Learning#
safe_stratified_split#
Stratified train/val split with random fallback when rare classes prevent stratification.
from pubmlp import safe_stratified_split
train_idx, val_idx = safe_stratified_split(X, y, test_size=0.2, random_state=42)
select_query_batch#
Select the most uncertain samples for human review.
from pubmlp import select_query_batch
query_indices = select_query_batch(probabilities, strategy='uncertainty', batch_size=30)
For multi-label tasks, pass the full 2-D probability matrix (records × labels): uncertainty averages |p − 0.5| across labels and max_relevance takes the per-record maximum. Collapsing the matrix beforehand (e.g., probabilities.max(axis=1)) treats a record that is certain on one label as certain overall and skips records still uncertain on the remaining labels.
create_review_batch#
Create a review batch DataFrame with model probability and prediction columns.
compare_reviewers#
Compute inter-rater agreement (kappa + agreement rate) between model and human.
merge_human_labels#
Merge human decisions from review batch back into the main DataFrame.
ALState#
Dataclass tracking active learning iteration state.
simulate_al#
Offline AL simulation using ground truth labels; model_fn(train_df, unlabeled_df) returns probabilities.
rank_by_hybrid_max_uncertainty#
95% max-relevance + 5% uncertainty ranking strategy.
rank_by_hybrid_max_random#
95% max-relevance + 5% random ranking strategy.
Benchmarks#
list_benchmarks#
Names available from the Synergy collection. Needs the benchmark extra: pip install pubmlp[benchmark].
load_benchmark#
One Synergy dataset as a normalized frame (title, abstract, year, journal, label_included).
load_manifest_corpus#
Rebuild a private corpus by joining a packaged IDs+labels manifest against a user-supplied database export. Raises when ID coverage falls below 95%.
normalize_benchmark_frame#
Map any labeled frame onto the benchmark column set.
build_column_specs#
Column specifications, numeric transform, and metadata fusion level (full, partial, text_only) for a benchmark frame.
embed_dataset#
One-time sentence embedding per dataset, cached to disk by model name and data hash.
run_simulation#
Per-seed strategy simulation (random, active_learning, tfidf_nb, semantic) with cached-embedding or fine-tune engines; one tidy row per evaluation with F1, ROC-AUC, ECE, Brier, NDCG, WSS@95, and stopping statistics.
Calibration metrics (ece, brier) are meaningful only for the random and active_learning strategies, whose probabilities are temperature-scaled; tfidf_nb reports raw naive-Bayes posteriors and semantic reports min-max-scaled cosine scores.
CLI: python -m pubmlp.benchmark --dataset NAME_OR_PATH --strategies random,active_learning,tfidf_nb --seeds 20 --out runs.csv
summarize_runs#
Mean (SD) per dataset and strategy at final effort, plus F1 curve frames.
LLM Screener#
llm_screen#
Screen records with a language model as an additional independent screener, never a sole screener. Takes a respond callable so the caller owns provider access and credentials. Returns the frame with {criterion}_llm, {criterion}_llm_confidence, {criterion}_llm_rationale, and llm_meets_all_criteria, plus a provenance dict for reporting. Self-reported confidence is not calibrated and is not comparable to predict_model probabilities.
build_prompt#
Compose the screening prompt for one record from the criterion descriptions.
parse_response#
Extract per-criterion decisions from a model reply; missing or unparseable decisions return None rather than a guess.
Evidence#
find_keyword_spans#
Find a keyword in text with the surrounding context. Returns each hit with the match offset inside its own context and the absolute bounds in the source, so two hits sharing context can be told from two that stand apart. The keyword is a literal, not a pattern.
search_document#
Search a parsed document for keywords, anchoring each hit to the printed page and the section it sits in.
highlight_markdown#
Render one span with the matched term in bold.
format_evidence#
Join spans into a single string for an evidence column.
pattern_from_terms#
Build a search pattern for regex_screen from plain terms, taking the truncation already used in a database search. A phrase matches across a line break; everything else is literal.
Full Text#
read_pdf#
Read a PDF into per-page text through pdfplumber.
detect_page_labels#
Read the page number printed in the document, accepted only when the printed-to-file offset is consistent.
detect_sections#
Assign each page the heading it falls under, including headings that share a line with body text.
extract_fulltext_evidence#
Return spans anchored to printed page and section for each criterion.
format_anchor#
Render a span’s location, marking it as a file position when the document carried no printed number.
Retrieval#
chunk_text#
Split text into overlapping chunks, breaking at whitespace.
build_index#
Embed chunks once for reuse across queries.
retrieve#
Rank chunks against a query by cosine similarity, in memory and with no vector database.
extract_with_rag#
Answer questions from retrieved passages. Generation goes through the same respond callable llm_screen takes. Every answer carries its passages and a grounding score; answers below the threshold are flagged for a second reviewer rather than accepted.
Confidence#
score_answer#
Score how well an extracted answer is grounded in its passages. A heuristic for triage, not a probability, and not comparable to predict_model output or to a model’s self-report.
interpret_confidence#
Band label for a grounding score.
needs_escalation#
Whether an extraction should go to a second reviewer.
score_extractions#
Score a list of extractions and flag the ones needing review.
confidence_report#
Counts by band and how many extractions need a second reviewer.
Provenance#
ProvenanceTracker#
Record what a run did: environment and package versions, model checkpoint, random seed, data size, retrieval configuration, calibration, decision threshold, and criterion descriptions verbatim. Prompts are recorded as sent with a running SHA-256 digest.
load_provenance#
Read a saved provenance record.
compare_provenances#
Report which recorded fields differ between two runs, ignoring timestamps.
Stopping Rules#
recall_target_test#
Statistical stopping test for recall-based screening. Returns stop decision, recall lower bound, and maximum missed relevant records.
should_stop#
SAFE Phase 2 stopping test. Criteria 3 and 4 (minimum screened, consecutive irrelevant) come from state and config. Criteria 1 and 2 (all predefined relevant found, screened at least twice the expected relevant count) are evaluated only when all_known_relevant_found and n_expected_relevant are supplied, and skipped otherwise. Omitting both tests two of the four criteria.
expected_relevant#
Project the corpus relevant count from a probability sample. Under active learning that is the prior-knowledge draw, since every later batch is model-selected.
update_stopping_state#
Update stopping state counters after a human screening decision.
estimate_recall#
Wilson score lower bound estimate of recall.
generate_stopping_report#
Generate a summary report of stopping criteria.
calculate_wss#
Calculate Work Saved over Sampling at a given recall level.
transition_phase#
Advance phase based on screening progress.
StoppingState#
Dataclass tracking stopping-rule state across iterations.
Audit#
AuditTrail#
Record and persist screening decisions for reproducibility.
AuditEntry#
Single audit log entry dataclass.
summarize_human_decisions#
Summarize human reviewer decisions from an audit trail.
generate_prisma_report#
Generate a PRISMA-style flow diagram report.
interpret_kappa#
Interpret Cohen’s kappa agreement level.
Utilities#
get_device#
Return the best available PyTorch device (CUDA or CPU).
auto_batch_size#
Suggest a batch size based on available GPU memory.
unpack_batch#
Move batch tensors to device and return unpacked components.
default_forward_fn#
Tokenized forward through encoder + classifier head.
cached_forward_fn#
Cached-embedding forward through head only (skips encoder).
compute_cls_embeddings#
Run the encoder once over a dataloader; returns (N, hidden_size) CPU tensor for the cached-embedding fast path.