Training Module¶
Trainers for every trainable approach in the benchmark: a lexical (LogReg/GBDT) classifier, a Word2Vec retrieval model, a fine-tuned transformer sequence classifier, and a fine-tuned SentenceTransformer embedding retriever. Rule-based and TF-IDF/BM25 retrieval have no learned parameters and therefore no trainer — TF-IDF/BM25 build their index on the fly at query time instead.
Input contract¶
Three of the four trainers (LexicalClassifierTrainer, ClassificationTrainer,
EmbeddingTrainer) share one input type: TrainingData,
which bundles a query Document, a source Document, and a
GroundTruth of labeled pairs. Word2VecTrainer is
unsupervised (it learns word embeddings from raw sentences, not labeled pairs),
so it takes plain Documents instead.
from locisimiles.document import Document
from locisimiles.ground_truth import GroundTruth
from locisimiles.training.data import TrainingData
query_doc = Document("query.csv")
source_doc = Document("source.csv")
positives = GroundTruth("known_positives.csv") # query_id, source_id, label
data = TrainingData(query_doc, source_doc, positives)
TrainingData¶
Iterating a TrainingData yields resolved (query_text, source_text, label)
triples — text is looked up by segment id once here rather than duplicated
across each trainer. TrainingData also carries all four of the paper's
negative-sampling methods as chainable methods, since assembling a training
set and sampling its negatives are naturally the same step.
sample_random_negatives (⟨qry,rnd⟩) is the default — a random
non-positive source segment per query, with no model loading required. The
paper's own ablation found sample_mixed_negatives (⟨qry,mix⟩) the
best-performing method on every reported metric (F1, accuracy, FPR, SMR); it's
offered as an explicit, one-call upgrade from the default rather than the
default itself, since it requires loading a pretrained embedding model to mine
hard negatives:
# Default: no model download needed
data = TrainingData(query_doc, source_doc, positives).sample_random_negatives(
n_per_query=5,
)
# Best-performing per the paper's ablation — random + embedding-mined hard negatives
data = TrainingData(query_doc, source_doc, positives).sample_mixed_negatives(
n_random_per_query=5,
n_hard_per_query=5,
)
sample_random_pairs (⟨rnd,rnd⟩, fully-random pairs — the paper's weakest
method) and sample_hard_negatives (⟨qry,sim⟩ alone) are also available for
completeness/parity with the paper. All four sampling methods return a new
TrainingData and never mutate the original, so they chain freely.
dataclass
¶
TrainingData(
query_doc: Document,
source_doc: Document,
ground_truth: GroundTruth,
)
Bundles a query document, source document, and ground truth for training.
The single input all pair/label trainers (LexicalClassifierTrainer,
ClassificationTrainer, EmbeddingTrainer) take. Iterating a
TrainingData yields resolved (query_text, source_text, label)
triples — text is looked up by segment id once here, rather than
duplicated across each trainer.
Negative-sampling methods are available directly as chainable, immutable methods:
| ATTRIBUTE | DESCRIPTION |
|---|---|
query_doc |
Query corpus.
TYPE:
|
source_doc |
Source corpus.
TYPE:
|
ground_truth |
Labeled query/source pairs.
TYPE:
|
sample_random_pairs(
*,
n_per_query: int = 1,
seed: int = 42,
label: LabelValue = "no_match",
) -> TrainingData
Add ⟨rnd,rnd⟩ negatives: independently-random (query, source) pairs.
Included for completeness/parity with the paper's ablation, which found this the weakest of the four methods.
sample_random_negatives(
*,
n_per_query: int = 1,
seed: int = 42,
label: LabelValue = "no_match",
) -> TrainingData
Add ⟨qry,rnd⟩ negatives — the default method.
For every query segment, samples a random non-positive source segment. Needs no model loading, so this is the package's default, first-reached-for sampling method.
sample_hard_negatives(
*,
n_per_query: int = 1,
embedding_model_name: str = DEFAULT_HARD_NEGATIVE_MODEL_NAME,
device: Optional[Union[str, int]] = None,
label: LabelValue = "no_match",
) -> TrainingData
Add ⟨qry,sim⟩ negatives: nearest non-positive neighbors by embedding similarity.
sample_mixed_negatives(
*,
n_random_per_query: int = 5,
n_hard_per_query: int = 5,
embedding_model_name: str = DEFAULT_HARD_NEGATIVE_MODEL_NAME,
device: Optional[Union[str, int]] = None,
seed: int = 42,
label: LabelValue = "no_match",
) -> TrainingData
Add ⟨qry,mix⟩ negatives — the paper's best-performing method on every
reported metric, combining random and embedding-mined hard negatives
for the same query. Offered as an explicit opt-in upgrade from the
default sample_random_negatives.
Implemented as a thin composition of sample_random_negatives and
sample_hard_negatives — no separate sampling logic.
LexicalClassifierTrainer¶
Trains the TF-IDF/Jaccard/overlap-feature LogReg or GBDT classifier consumed
by LexicalClassifierJudge.
from locisimiles.training.lexical import LexicalClassifierTrainer, LexicalClassifierTrainerConfig
config = LexicalClassifierTrainerConfig(output_dir="models/lexical", classifier="logreg")
trainer = LexicalClassifierTrainer(config)
trainer.fit(data=data)
artifact_path = trainer.save()
dataclass
¶
LexicalClassifierTrainerConfig(
output_dir: Path,
seed: int = 42,
lowercase: bool = True,
normalize_ij_uv: bool = True,
classifier: Literal["logreg", "gbdt"] = "logreg",
lemmatize: bool = True,
output_filename: str = "lexical_classifier.joblib",
logreg_C: float = 1.0,
logreg_max_iter: int = 1000,
gbdt_max_iter: int = 300,
gbdt_learning_rate: float = 0.05,
gbdt_max_depth: int | None = None,
class_weight: str | dict | None = None,
label_names: dict[int, str] | None = None,
)
Configuration for the lexical classifier trainer.
fit() takes a TrainingData whose ground truth label may be an
integer class id or a string class name, e.g. no_match / cit /
cf for the three-class setup, or any two-class scheme for binary
training.
LexicalClassifierTrainer(
config: LexicalClassifierTrainerConfig,
)
Train a TF-IDF/Jaccard/overlap feature-based LogReg or GBDT classifier.
Ensure the output directory exists; fit() validates its TrainingData.
fit(*, data: TrainingData, **kwargs: Any) -> Any
Fit TF-IDF vectorizers and the configured classifier on paired training data.
Persist vectorizers, classifier, and label mapping to one joblib artifact.
Load a previously saved lexical classifier artifact.
Word2VecTrainer¶
Trains the gensim Word2Vec model consumed by Word2VecCandidateGenerator
(the Burns-style retrieval baseline). Unsupervised: takes Documents
directly rather than a TrainingData, since there's no label to learn from.
from locisimiles.training.word2vec import Word2VecTrainer, Word2VecTrainerConfig
config = Word2VecTrainerConfig(output_dir="models/word2vec")
trainer = Word2VecTrainer(config)
trainer.fit(documents=[query_doc, source_doc])
model_path = trainer.save()
dataclass
¶
Word2VecTrainerConfig(
output_dir: Path,
seed: int = 42,
lowercase: bool = True,
normalize_ij_uv: bool = True,
vector_size: int = 300,
window: int = 5,
min_count: int = 1,
sg: int = 1,
workers: int = 1,
epochs: int = 10,
output_filename: str = "latin_w2v.model",
)
Configuration specific to Word2Vec training.
Word2VecTrainer(config: Word2VecTrainerConfig)
Train a local gensim Word2Vec model from one or more Documents.
Word2Vec is unsupervised (it learns word embeddings from raw sentences,
not labeled pairs), so unlike the pair/label trainers it takes plain
Documents rather than a TrainingData.
ClassificationTrainer¶
Fine-tunes the transformer sequence-classification model consumed by
ClassificationJudge — binary (no_match/match) or
multiclass (no_match/cit/cf), inferred from the distinct labels in
data. Mirrors the paper's recipe by default: a fixed epoch count with no
checkpoint selection, optional balanced class-weighting or focal loss, and
the same pair-truncation strategy ClassificationJudge uses at inference
time. model.config.id2label/label2id are set automatically before
saving, so a freshly trained model is immediately usable by
ClassificationJudge — no manual upload/labeling step.
Pass eval_data to fit() to compute eval-set loss after every epoch. By
default this is purely for visibility (the paper's fixed-epoch-then-save
recipe); set select_best_checkpoint=True to instead keep the epoch with
the lowest eval loss, optionally with early_stopping_patience to stop
once that loss stops improving:
trainer.fit(
data=data,
eval_data=eval_data,
) # select_best_checkpoint=False (default): eval_data is logged only
config = ClassificationTrainerConfig(
output_dir="models/classifier",
select_best_checkpoint=True,
early_stopping_patience=2,
)
trainer = ClassificationTrainer(config)
trainer.fit(data=data, eval_data=eval_data) # keeps the best-loss epoch
from locisimiles.training.classification import ClassificationTrainer, ClassificationTrainerConfig
config = ClassificationTrainerConfig(
output_dir="models/classifier-3class",
model_name="xlm-roberta-base",
label_names={0: "no_match", 1: "cit", 2: "cf"},
epochs=4,
class_weight="balanced",
)
trainer = ClassificationTrainer(config)
trainer.fit(data=data)
model_path = trainer.save()
dataclass
¶
ClassificationTrainerConfig(
output_dir: Path,
seed: int = 42,
lowercase: bool = True,
normalize_ij_uv: bool = True,
model_name: str = "xlm-roberta-base",
label_names: Optional[Dict[int, str]] = None,
epochs: int = 4,
batch_size: int = 32,
learning_rate: float = 2e-05,
max_length: int = 512,
class_weight: Optional[Literal["balanced"]] = None,
use_focal_loss: bool = False,
focal_gamma: float = 2.0,
apply_roberta_separator_fix: bool = False,
disable_compile: bool = False,
device: str = "cpu",
output_filename: str = "classifier",
select_best_checkpoint: bool = False,
early_stopping_patience: Optional[int] = None,
)
Configuration for the transformer sequence-classification trainer.
fit() takes a TrainingData whose ground truth label may be
binary (e.g. no_match/match) or multiclass (e.g.
no_match/cit/cf) — the number of classes is inferred from the
distinct labels seen during fit().
ClassificationTrainer(config: ClassificationTrainerConfig)
Fine-tune a transformer sequence-classification model on labeled pairs.
Mirrors the experiment recipe: a hand-rolled AdamW loop, no LR
scheduler, a fixed epoch count by default (the model is saved after the
last epoch), with optional balanced class-weighting or focal loss. Pass
select_best_checkpoint=True (and eval_data) to instead keep the
epoch with the lowest eval loss, optionally with early stopping. Pair
truncation reuses the exact strategy
:class:~locisimiles.pipeline.judge.classification.ClassificationJudge
uses at inference time, so train/inference tokenization matches.
Ensure the output directory exists; fit() validates its TrainingData.
fit(
*,
data: TrainingData,
eval_data: Optional[TrainingData] = None,
**kwargs: Any,
) -> Any
Fine-tune the classifier on resolved (query_text, source_text, label) pairs.
eval_data, if given, has its cross-entropy loss computed after
every epoch — purely for visibility by default
(select_best_checkpoint=False, matching the paper's
fixed-epoch-then-save recipe). Set select_best_checkpoint=True to
instead keep the epoch with the lowest eval_data loss, optionally
with early_stopping_patience to stop early once that loss stops
improving.
tune_threshold(
*,
data: TrainingData,
method: str = "max_f1",
negative_label: str = "no_match",
**kwargs: Any,
) -> ThresholdSet
Tune one-vs-rest decision thresholds on an evaluation TrainingData.
For multiclass label schemes, tunes an independent threshold per positive class; ties between classes that both clear their threshold are broken by comparing probabilities (the only tie-break rule currently supported).
Persist the fine-tuned model, tokenizer, and (if tuned) thresholds.
Sets model.config.id2label/label2id from the resolved label
mapping before saving, so the saved directory is immediately usable
by :class:~locisimiles.pipeline.judge.classification.ClassificationJudge
with no manual step.
Load a previously saved classifier directory.
Threshold tuning and application¶
A trained classifier's raw argmax decision isn't always what you want —
tune_threshold sweeps one-vs-rest decision thresholds per positive class on
an evaluation TrainingData (with a tie-break rule for when two positive
classes both clear their threshold on the same pair), and save() persists
the result as a threshold.json sidecar alongside the model.
Applying a tuned ThresholdSet is a deliberate standalone post-processing
step, not something baked into ClassificationJudge — run judge.judge(...)
as normal, then optionally apply the tuned thresholds:
trainer.tune_threshold(data=eval_data, method="max_f1")
model_path = trainer.save() # writes threshold.json alongside the model
# --- later, at inference time ---
from locisimiles.pipeline.judge import ClassificationJudge
from locisimiles.training.classification.threshold import ThresholdSet, apply_thresholds_to_judgments
judge = ClassificationJudge(classification_name=str(model_path))
results = judge.judge(query=query_doc, candidates=candidates)
thresholds = ThresholdSet.from_json(model_path / "threshold.json")
final = apply_thresholds_to_judgments(results, thresholds, negative_label="no_match")
dataclass
¶
ThresholdSet(
thresholds: Dict[str, float],
method: str = "max_f1",
tie_break: str = "max_probability",
)
Per-class decision thresholds tuned by :meth:ClassificationTrainer.tune_threshold.
| ATTRIBUTE | DESCRIPTION |
|---|---|
thresholds |
Mapping from positive class label to its tuned one-vs-rest decision threshold.
TYPE:
|
method |
Threshold-tuning method used (
TYPE:
|
tie_break |
Rule used to resolve ties when multiple positive classes
clear their threshold for the same pair. Only
TYPE:
|
Persist this threshold set as a JSON sidecar file.
classmethod
¶
from_json(path: Union[str, Path]) -> ThresholdSet
Load a threshold set previously written by :meth:to_json.
apply_thresholds_to_judgments(
judgments: CandidateJudgeOutput,
thresholds: ThresholdSet,
*,
negative_label: str = "no_match",
) -> CandidateJudgeOutput
Re-decide judgment_score/predicted_label using tuned thresholds.
Runs :func:apply_thresholds over every judgment's existing
class_probabilities and returns a new CandidateJudgeOutput
with judgment_score/predicted_label/predicted_class_id
overwritten accordingly. Judgments without class_probabilities
(e.g. from a judge run without emit_class_metadata=True) are passed
through unchanged.
| PARAMETER | DESCRIPTION |
|---|---|
judgments
|
Output of
TYPE:
|
thresholds
|
Tuned thresholds from :func:
TYPE:
|
negative_label
|
Label treated as the non-link class.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
CandidateJudgeOutput
|
A new |
apply_thresholds(
class_probabilities: Dict[str, float],
thresholds: ThresholdSet,
*,
negative_label: str = "no_match",
) -> Tuple[str, Optional[int]]
Apply tuned one-vs-rest thresholds and the tie-break rule to one probability row.
| PARAMETER | DESCRIPTION |
|---|---|
class_probabilities
|
Mapping from class label to probability, e.g.
as emitted by
TYPE:
|
thresholds
|
Tuned thresholds from :func:
TYPE:
|
negative_label
|
Label returned when no positive class clears its threshold.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
str
|
|
Optional[int]
|
always |
Tuple[str, Optional[int]]
|
not a label-to-id mapping. |
EmbeddingTrainer¶
Fine-tunes the SentenceTransformer bi-encoder consumed by
EmbeddingCandidateGenerator. Uses OnlineContrastiveLoss by
default (the paper's production loss); prompts maps the "query"/"match"
sides to asymmetric prefixes (E5-style by default) and is baked into the saved
model, so EmbeddingCandidateGenerator(embedding_model_name=...,
prompt_name="query"/"match") works out of the box after training.
from locisimiles.training.embedding import EmbeddingTrainer, EmbeddingTrainerConfig
config = EmbeddingTrainerConfig(
output_dir="models/embedding",
model_name="intfloat/multilingual-e5-large",
)
trainer = EmbeddingTrainer(config)
trainer.fit(data=data, eval_data=eval_data) # eval_data is optional
model_path = trainer.save()
As with ClassificationTrainer, eval_data alone (via a
BinaryClassificationEvaluator run once per epoch) is purely for
visibility by default; set select_best_checkpoint=True — optionally with
early_stopping_patience — to keep the epoch with the best eval score
instead of the last one:
config = EmbeddingTrainerConfig(
output_dir="models/embedding",
select_best_checkpoint=True,
early_stopping_patience=2,
)
trainer = EmbeddingTrainer(config)
trainer.fit(data=data, eval_data=eval_data)
dataclass
¶
EmbeddingTrainerConfig(
output_dir: Path,
seed: int = 42,
lowercase: bool = True,
normalize_ij_uv: bool = True,
model_name: str = "intfloat/multilingual-e5-small",
loss_type: Literal[
"online_contrastive", "contrastive"
] = "online_contrastive",
epochs: int = 4,
batch_size: int = 32,
learning_rate: float = 2e-05,
weight_decay: float = 0.01,
warmup_ratio: float = 0.1,
prompts: Dict[str, str] = _default_prompts(),
negative_label: str = "no_match",
device: str = "cpu",
output_filename: str = "embedding_model",
select_best_checkpoint: bool = False,
early_stopping_patience: Optional[int] = None,
)
Configuration for the SentenceTransformer bi-encoder trainer.
prompts maps the dataset column names fit() builds internally
("query"/"match") to the asymmetric prefixes prepended at
training/inference time (E5-style by default). This is exactly what
:class:~locisimiles.pipeline.generator.embedding.EmbeddingCandidateGenerator
later calls with prompt_name="query"/"match".
EmbeddingTrainer(config: EmbeddingTrainerConfig)
Fine-tune a SentenceTransformer bi-encoder on labeled query/source pairs.
Uses OnlineContrastiveLoss (the paper's production loss) by default;
triplet-loss variants tried in the experiments were superseded and are
not ported. By default, trains for a fixed epoch count with no
checkpoint selection, matching the experiment recipe; pass
select_best_checkpoint=True (and eval_data) to opt into
best-epoch selection and, optionally, early stopping.
Ensure the output directory exists; fit() validates its TrainingData.
fit(
*,
data: TrainingData,
eval_data: Optional[TrainingData] = None,
**kwargs: Any,
) -> Any
Fine-tune the embedding model on resolved (query_text, source_text, label) pairs.
eval_data, if given, runs a BinaryClassificationEvaluator once
per epoch. By default (select_best_checkpoint=False) this is
purely for visibility — the final epoch's weights are what get
returned/saved, matching the paper's fixed-epoch-then-save recipe. Set
select_best_checkpoint=True to instead keep the epoch with the
best eval_data score (average precision over cosine similarity),
optionally with early_stopping_patience to stop early once that
score stops improving.
Persist the fine-tuned SentenceTransformer directory.
Re-asserts model.prompts from config as a defensive measure (in
case anything reset it after fit()), then saves — unlike the
classification trainer, no id/label bookkeeping is needed here.
Load a previously saved SentenceTransformer directory.
Cross-validation¶
The paper reports mean±std across folds rather than a single train/test
split. cross_validate reproduces that protocol: it splits a
GroundTruth into folds grouped by query_id (via
split_ground_truth_by_query — a query's positives and negatives never
straddle a train/eval boundary), then for each fold calls a train_fn
(trains and returns a model/pipeline from that fold's training data) and an
evaluate_fn (evaluates it on the held-out fold, returning a flat metrics
dict), aggregating the results into a CVResult with per-fold, mean, and
std metrics.
cross_validate is deliberately generic — train_fn/evaluate_fn are
plain callables, so it works the same way across all four trainers rather
than hardcoding one. evaluate_with_pipeline is a convenience wrapper for
the common case of evaluating with IntertextEvaluator:
from locisimiles.training.cross_validation import cross_validate, evaluate_with_pipeline
from locisimiles.training.classification import ClassificationTrainer, ClassificationTrainerConfig
from locisimiles.pipeline import Pipeline
from locisimiles.pipeline.generator import ExhaustiveCandidateGenerator
from locisimiles.pipeline.judge import ClassificationJudge
def train_fn(fold_train_data):
trainer = ClassificationTrainer(ClassificationTrainerConfig(output_dir="models/cv"))
trainer.fit(data=fold_train_data)
return trainer.save()
def evaluate_fn(model_path, fold_eval_data):
judge = ClassificationJudge(classification_name=str(model_path))
pipeline = Pipeline(generator=ExhaustiveCandidateGenerator(), judge=judge)
return evaluate_with_pipeline(pipeline, fold_eval_data)
result = cross_validate(
query_doc=query_doc,
source_doc=source_doc,
ground_truth=ground_truth,
n_folds=5,
train_fn=train_fn,
evaluate_fn=evaluate_fn,
)
print(result.mean, result.std)
print(result.to_dataframe()) # per-fold rows + mean/std summary rows
dataclass
¶
CVFold(
index: int,
train_data: TrainingData,
eval_data: TrainingData,
)
One train/eval split of a K-fold cross-validation run.
| ATTRIBUTE | DESCRIPTION |
|---|---|
index |
Zero-based fold index.
TYPE:
|
train_data |
Training data — the union of every fold except this one.
TYPE:
|
eval_data |
This fold's held-out evaluation data.
TYPE:
|
dataclass
¶
Aggregated results of a :func:cross_validate run.
| ATTRIBUTE | DESCRIPTION |
|---|---|
fold_metrics |
One metrics dict per fold, in fold order.
TYPE:
|
mean |
Per-metric mean across folds.
TYPE:
|
std |
Per-metric standard deviation across folds (
TYPE:
|
Return per-fold metrics as a DataFrame, with mean/std summary rows appended.
cross_validate(
*,
query_doc: Document,
source_doc: Document,
ground_truth: GroundTruth,
n_folds: int,
train_fn: Callable[[TrainingData], Any],
evaluate_fn: Callable[
[Any, TrainingData], Dict[str, float]
],
seed: int = 42,
) -> CVResult
Run K-fold cross-validation, reproducing the paper's mean±std-across-folds protocol.
For each fold, train_fn is called with the fold's training
TrainingData (the union of every other fold) and must return a
trained model/pipeline; evaluate_fn is then called with that return
value and the fold's held-out TrainingData, and must return a flat
metric-name-to-value dict. Metrics are aggregated (mean/std) across
folds. Folds are grouped by query_id (see
:func:split_ground_truth_by_query), so a query's positives and
negatives never straddle a train/eval boundary.
| PARAMETER | DESCRIPTION |
|---|---|
query_doc
|
Query corpus.
TYPE:
|
source_doc
|
Source corpus.
TYPE:
|
ground_truth
|
Full ground truth to cross-validate over.
TYPE:
|
n_folds
|
Number of folds (at least 2).
TYPE:
|
train_fn
|
Trains and returns a model/pipeline from one fold's training data.
TYPE:
|
evaluate_fn
|
Evaluates a trained model/pipeline on one fold's held-out data.
TYPE:
|
seed
|
RNG seed for the fold split.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
A
|
class:
TYPE:
|
Example
from locisimiles.training.cross_validation import cross_validate, evaluate_with_pipeline
from locisimiles.training.classification import ClassificationTrainer, ClassificationTrainerConfig
from locisimiles.pipeline import Pipeline
from locisimiles.pipeline.generator import ExhaustiveCandidateGenerator
from locisimiles.pipeline.judge import ClassificationJudge
def train_fn(data):
trainer = ClassificationTrainer(
ClassificationTrainerConfig(output_dir="models/cv", epochs=4)
)
trainer.fit(data=data)
return trainer.save()
def evaluate_fn(model_path, eval_data):
judge = ClassificationJudge(classification_name=str(model_path))
pipeline = Pipeline(generator=ExhaustiveCandidateGenerator(), judge=judge)
return evaluate_with_pipeline(pipeline, eval_data)
result = cross_validate(
query_doc=query_doc,
source_doc=source_doc,
ground_truth=ground_truth,
n_folds=5,
train_fn=train_fn,
evaluate_fn=evaluate_fn,
)
print(result.mean, result.std)
evaluate_with_pipeline(
pipeline: Pipeline,
eval_data: TrainingData,
*,
top_k: int = 10,
average: str = "macro",
**evaluator_kwargs: Any,
) -> Dict[str, float]
Evaluate a pipeline on one fold's held-out data via IntertextEvaluator.
A convenience wrapper for the common evaluate_fn case in
:func:cross_validate: builds an IntertextEvaluator from
eval_data and returns evaluator.evaluate(average=average) as a
flat dict.
| PARAMETER | DESCRIPTION |
|---|---|
pipeline
|
A trained/configured pipeline to evaluate.
TYPE:
|
eval_data
|
Held-out
TYPE:
|
top_k
|
Candidates per query segment, forwarded to
TYPE:
|
average
|
TYPE:
|
**evaluator_kwargs
|
Additional keyword arguments forwarded to
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
Dict[str, float]
|
A flat metric-name-to-value dict for this fold. |
make_cv_folds(
*,
query_doc: Document,
source_doc: Document,
ground_truth: GroundTruth,
n_folds: int,
seed: int = 42,
) -> List[CVFold]
Build n_folds train/eval TrainingData splits, grouped by query_id.
Fold i's eval_data is the held-out fold produced by
:func:split_ground_truth_by_query; its train_data is the union of
every other fold.
| PARAMETER | DESCRIPTION |
|---|---|
query_doc
|
Query corpus.
TYPE:
|
source_doc
|
Source corpus.
TYPE:
|
ground_truth
|
Ground truth to split.
TYPE:
|
n_folds
|
Number of folds (at least 2).
TYPE:
|
seed
|
RNG seed forwarded to :func:
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
List[CVFold]
|
A list of |
split_ground_truth_by_query(
ground_truth: GroundTruth,
n_folds: int,
*,
seed: int = 42,
) -> List[GroundTruth]
Split a GroundTruth into n_folds folds grouped by query_id.
Every entry for a given query id lands in the same fold. Query ids are
shuffled with seed before being assigned round-robin to folds, so
fold sizes are as balanced as the query distribution allows.
| PARAMETER | DESCRIPTION |
|---|---|
ground_truth
|
Ground truth to split.
TYPE:
|
n_folds
|
Number of folds (at least 2).
TYPE:
|
seed
|
RNG seed for shuffling query ids before assignment.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
List[GroundTruth]
|
A list of |
List[GroundTruth]
|
recovers all of |
Negative sampling functions¶
The plain functions underlying TrainingData's sampling methods, for callers
who want them independent of the TrainingData wrapper.
sample_random_pairs(
*,
query_doc: Document,
source_doc: Document,
positives: GroundTruth,
n_per_query: int = 1,
seed: int = 42,
label: LabelValue = "no_match",
) -> GroundTruth
⟨rnd,rnd⟩ — independently-random (query, source) pairs.
Draws n_per_query * len(query_doc) pairs by picking a query segment
and a source segment independently and uniformly at random (not
conditioned on any specific query, unlike the other sampling functions
here), skipping any draw that happens to coincide with a known positive
in positives. Included for completeness/parity with the paper's
ablation — its own results table shows this is the weakest of the four
methods.
| PARAMETER | DESCRIPTION |
|---|---|
query_doc
|
Query corpus.
TYPE:
|
source_doc
|
Source corpus.
TYPE:
|
positives
|
Known positive pairs to avoid mislabeling as negative.
TYPE:
|
n_per_query
|
Target number of random pairs per query segment.
TYPE:
|
seed
|
RNG seed for reproducibility.
TYPE:
|
label
|
Label assigned to sampled negatives.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
GroundTruth
|
A |
sample_random_negatives(
*,
query_doc: Document,
source_doc: Document,
positives: GroundTruth,
n_per_query: int = 1,
seed: int = 42,
label: LabelValue = "no_match",
) -> GroundTruth
⟨qry,rnd⟩ — for every query segment, n_per_query random non-positive
source segments.
Uses a seed derived deterministically per query id (rather than one shared RNG stream) so results are reproducible regardless of query iteration order.
| PARAMETER | DESCRIPTION |
|---|---|
query_doc
|
Query corpus.
TYPE:
|
source_doc
|
Source corpus.
TYPE:
|
positives
|
Known positive pairs to exclude from sampling.
TYPE:
|
n_per_query
|
Number of negatives to sample per query segment.
TYPE:
|
seed
|
Base RNG seed for reproducibility.
TYPE:
|
label
|
Label assigned to sampled negatives.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
GroundTruth
|
A |
sample_hard_negatives(
*,
query_doc: Document,
source_doc: Document,
positives: GroundTruth,
n_per_query: int = 1,
embedding_model_name: str = DEFAULT_HARD_NEGATIVE_MODEL_NAME,
device: Optional[Union[str, int]] = None,
label: LabelValue = "no_match",
) -> GroundTruth
⟨qry,sim⟩ — for every query segment, the top n_per_query
non-positive source segments ranked by embedding similarity.
Reuses :class:~locisimiles.pipeline.generator.embedding.EmbeddingCandidateGenerator
to mine negatives, matching the paper's "a pretrained embedding model"
nearest-neighbor mining. embedding_model_name should be a
general-purpose pretrained model, not the fine-tuned model being
trained — mining hard negatives with the very model about to be
fine-tuned would be circular.
| PARAMETER | DESCRIPTION |
|---|---|
query_doc
|
Query corpus.
TYPE:
|
source_doc
|
Source corpus.
TYPE:
|
positives
|
Known positive pairs to exclude from the mined candidates.
TYPE:
|
n_per_query
|
Number of hard negatives to keep per query segment.
TYPE:
|
embedding_model_name
|
Pretrained sentence-transformer used to mine negatives.
TYPE:
|
device
|
Torch device string for the mining model.
TYPE:
|
label
|
Label assigned to sampled negatives.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
GroundTruth
|
A |
Migrating from the CSV-based trainer API¶
LexicalClassifierTrainer and Word2VecTrainer previously read a flat,
denormalized CSV directly. As of this training-API unification, all trainers
take Document/TrainingData instead — this is a breaking change:
# Before
config = LexicalClassifierTrainerConfig(train_path="train.csv", output_dir="models/")
trainer = LexicalClassifierTrainer(config)
trainer.fit()
# After
config = LexicalClassifierTrainerConfig(output_dir="models/")
trainer = LexicalClassifierTrainer(config)
trainer.fit(data=TrainingData(query_doc, source_doc, ground_truth))
# Before
config = Word2VecTrainerConfig(train_path="train.csv", output_dir="models/")
trainer = Word2VecTrainer(config)
trainer.fit()
# After
config = Word2VecTrainerConfig(output_dir="models/")
trainer = Word2VecTrainer(config)
trainer.fit(documents=[query_doc, source_doc])
IntertextEvaluator's ground_truth_csv parameter is likewise renamed to
ground_truth (still accepts a path, DataFrame, or now also a
GroundTruth), and load_example_ground_truth() returns a GroundTruth
instead of a list of dicts.