ML-USAS

This repository contains various DataTrove pipelines, filters, formatters, and helper functions for processing the HuggingFaceFW finewiki dataset, in various languages listed in the languages section, to create a synthetic (silver labelled) training dataset for USAS semantic tags and Multi Word Expression (MWE) identification for some languages.

For more information on the filtering and processing, see the filtering and processing section below and for more information about the data we use see the data section below.

Setup

You can either use the dev container with your favourite editor, e.g. VSCode. Or you can create your setup locally below we demonstrate both.

In both cases they share the same tools, of which these tools are:

• uv for Python packaging and development
• make (OPTIONAL) for automation of tasks, not strictly required but makes life easier.

Dev Container

A dev container uses a docker container to create the required development environment, the Dockerfile we use for this dev container can be found at ./.devcontainer/Dockerfile. To run it locally it requires docker to be installed, you can also run it in a cloud based code editor, for a list of supported editors/cloud editors see the following webpage.

To run for the first time on a local VSCode editor (a slightly more detailed and better guide on the VSCode website):

1. Ensure docker is running.
2. Ensure the VSCode Dev Containers extension is installed in your VSCode editor.
3. Open the command pallete CMD + SHIFT + P and then select Dev Containers: Rebuild and Reopen in Container

You should now have everything you need to develop, uv, make, for VSCode various extensions like Pylance, etc.

If you have any trouble see the VSCode website..

Local

To run locally first ensure you have the following tools installted locally:

• uv for Python packaging and development. (version 0.9.6)
• make (OPTIONAL) for automation of tasks, not strictly required but makes life easier.
  • Ubuntu: apt-get install make
  • Mac: Xcode command line tools includes make else you can use brew.
  • Windows: Various solutions proposed in this blog post on how to install on Windows, inclduing Cygwin, and Windows Subsystem for Linux.

When developing on the project you will want to install the Python package locally in editable format with all the extra requirements, this can be done like so:

uv sync --all-extras

Linting

Linting and formatting with ruff it is a replacement for tools like Flake8, isort, Black etc, and we us ty for type checking.

To run the linting:

make lint

Tests

To run the tests (uses pytest and coverage) and generate a coverage report:

make test

Example of using Flash Attention

from sentence_transformers import SentenceTransformer

model = SentenceTransformer(
    "sentence-transformers/all-MiniLM-L6-v2",
    model_kwargs={"attn_implementation": "flash_attention_2", "torch_dtype": "bfloat16"},
)

sentences = ["This is an example sentence", "Each sentence is converted"]
embeddings = model.encode(sentences)

Offline data download of the USAS silver labelled dataset

uv run processing_scripts/download_dataset.py ./data/usas_silver_data/

HuggingFace Authentication

Before processing or uploading to the HuggingFace hub please authenticate using a token from huggingface.co/settings/tokens;

hf auth login

or by using a token that is set within ./.env, read using dotenv, e.g.

HF_TOKEN="HUGGINGFACE_TOKEN_KEY_VALUE"

Set the relevant permissions, the minimum for this is repository is "read" only permission, if you want to upload the created synthetic silver labelled dataset to HuggingFace please ensure that you have allowed write permission to the namespace/repository you are going to upload too on HuggingFace.

Data

The data will be coming from HuggingFaceFW finewiki dataset and will be filtered so that each Wikipedia article is either rated as a "Good Articles" (GA) or "Featured Articles" (FA) by an editor, we hope that this will remove articles that might be incomplete or require additional editing. This filtering is inspired by Conia et al. 2024 whereby they found training on data from only "featured" and "good" articles performed similarly to training on the far larger Wikipedia articles that contained non-good and non-featured articles thus showing that training on smaller amounts of data is as affective and more efficient. The "Featured" and "Good" article can be defined differently for each Wikipedia language site as stated in the English site definition within the following article. The list of GA and FA can be found at the HuggingFace dataset ucrelnlp/wikipedia-ga-fa-ids.

Software

Train with HuggingFace through sentence-transformers, experimental tracking with trackio, and carbon emission tracking using carboncode.

Notes

• Data can be sampled so that we see a maximum number of tokens, $N$ from a given tag class, if more than $N$ occur for a given tag class then tokens assigned with that class are sub-sampled by $M/N$ where $M$ are the number of token samples for that given class.

Metrics

Accuracy@k — binary hit: counts the query correct if any relevant doc appears in the top-k, regardless of how many positives exist or how many are retrieved. Precision@k — num_correct / k_val, where num_correct counts every retrieved doc in the top-k that's in query_relevant_docs — so multiple positives in the top-k all get counted. Recall@k — num_correct / len(query_relevant_docs), i.e. normalized by the total number of positives for that query, not just 1. MRR@k — only looks at the rank of the first relevant hit and breaks; additional positives don't affect it. NDCG@k — true_relevances = [1] * len(query_relevant_docs), so the ideal DCG (denominator) scales with the number of positives. This means with multiple positives, achieving NDCG=1.0 at a given k requires retrieving all of them in the top-k, not just one. MAP@k — accumulates precision at each rank where a hit is relevant, then divides by min(k_val, len(query_relevant_docs)), i.e. standard average precision generalized to multiple relevant docs. So: the top-k search never special-cases multiple positives — it's pure score ranking. Multiple positives only change results in compute_metrics, where recall, NDCG, and MAP explicitly weight by len(query_relevant_docs), while accuracy and MRR are indifferent to a query having more than one positive (they only care about the first/any hit).

Which would you suggest is the best metric to capture as many of the true positive values in the top k? Recall@k is the direct fit — num_correct / len(query_relevant_docs) literally measures what fraction of a query's true positives you captured within the top-k, and it's the only one of the six that's normalized by the actual positive count rather than by k or by "at least one."

Contrast with the others:

Precision@k counts the same num_correct but divides by k_val, so it measures result purity, not positive coverage — it penalizes you for a large k even if you caught every positive. Accuracy@k and MRR@k only care about the first hit, so they're blind to whether you caught 1 of 5 positives or 5 of 5. NDCG@k and MAP@k are good complements if rank order among the positives also matters to you (e.g., you want the positives near the top of the k, not just present) — but if you strictly want "how many of the true positives did I retrieve," Recall@k is the metric that answers that question directly. If you care about both coverage and rank quality, I'd report Recall@k as the primary number and NDCG@k alongside it.

Future improvements

• Adding un-pad token efficiency for flash-attention: currently blocked because build_model in src/ml_usas/wsd/model.py sets unpad_inputs=False. This is needed because TargetSpanPooling (src/ml_usas/wsd/pooling.py) assumes the padded (batch, seq_len, hidden) layout — it broadcasts word_ids/target_index against attention_mask, which requires a shared seq_len dimension across the batch. Un-padding instead produces a flat (total_tokens, hidden) tensor with cu_seqlens sample boundaries, so TargetSpanPooling would need rewriting to index off that layout instead of the padded mask.

• Adding hard negative examples from a teacher model that finds difficult examples.

• Turn WSD training dataset construction into a streamed/iterable (datasets.IterableDataset) pipeline where negative sampling happens on-the-fly at train time (in the dataset generator or collator) instead of being baked in once at build_wsd_dataset.py build time. Motivation:

  • Now that class-balance loss weighting (dataset.compute_tag_weights) reweights rows by inverse effective sample count, max_examples_per_tag's row-count cap matters less for correcting label imbalance, making a cap-free streaming approach more attractive — it would let training use more of the corpus instead of downsampling per tag.
  • "Epoch" could then be defined by a number of training steps rather than one pass over a fixed, pre-materialized dataset.
  • The negative-sampling strategy could evolve during training — e.g. start with the current rule-based confusable-tag sampling (negatives.sample_hard_negatives), then switch to embedding-based hard negative mining using the model checkpoint actually being trained. This is cheap to do: negatives are always drawn from the small, fixed pool of ~211-222 USAS tag definitions (not other anchors), so periodically re-embedding just that pool with the current checkpoint and ranking by cosine similarity is enough — no full-corpus index needed.
  • The positive_only_no_duplicates_batch_sampler batch-uniqueness constraint that max_examples_per_tag currently helps keep feasible may no longer be a hard requirement — MaskedMultipleNegativesRankingLoss's false-negative masking (build_false_negative_mask) already excludes any in-batch doc whose tag is one of an anchor's true tags, so a duplicate positive tag in a batch shouldn't wrongly penalize a correct match.
  • Open questions for whoever scopes this properly: whether it applies to both dataset shapes (per_tag/masked-mnrl and per_token/local-multi-positive) or just per_tag (the shape max_examples_per_tag actually affects), and the refresh cadence for re-embedding the tag-definition pool during mining.

• I think token_row_list = list(token_rows) should be an iterable of some description within to_ir_eval_data which is in src/ml_usas/wsd/dataset.py

• Create an evaluation metric for top_k accuracy whereby the model has to predict all of the valid tokens, when it is a multi tag token the bi-encoder models will never get those samples correct.

• Logging with trackio

• The loss function is not ideal as it performs in-batch negative sampling which can contain samples that are correct for certain samples as we are performing multi-label classification. Thus we might want to enhance the loss function so that it only uses hard negatives, ignores in batch samples per sample that are positives, or have a loss function that takes into account multiple positive labels per token when it does occur.

• We could actually make the MaskedMultipleNegativesRankingLoss loss more interesting by including multiple positive examples rather than one.

• It would be good to see if we can turn the model into a pure token based model that can perform what is being done at the sentence level but with multiple tokens at once.

• Change the dataset generation so that it sub-samples both the training and evaluation datasets to a max number of samples based on a maximum number of samples per class.

• Switch the model's similarity function from cosine to dot product for faster inference/retrieval (dot product skips the per-comparison normalization that cosine similarity requires, which matters at ANN-index scale, e.g. FAISS IndexFlatIP). This requires adding a Normalize module to the end of build_model's module pipeline in src/ml_usas/wsd/model.py, so every embedding the model produces — at train and inference time alike — is unit-length. dot_score on unit-length embeddings is numerically identical to cos_sim on raw embeddings (cos_sim normalizes internally anyway), so this is free during training: same loss, same gradients, no retuning of MaskedMultipleNegativesRankingLoss/LocalMultiPositiveLoss's scale. Skipping Normalize and switching only the loss's similarity_fct to util.dot_score would not be equivalent — without baked-in normalization, dot product is sensitive to embedding magnitude, which both breaks train/inference consistency and gives the loss a degenerate shortcut (inflating embedding norms lowers the softmax loss without improving actual alignment).

Scripts ran so far;

#uv run processing_scripts/build_wsd_dataset.py data/usas_silver_data/en/ ./data/training_data/en --num-negatives 4 --num-confusable-negatives 2 --max-examples-per-tag 10000 --max-eval-queries 2000 --seed 12 --overwrite


uv run processing_scripts/build_wsd_dataset.py data/usas_silver_data/en/ ./data/training_data/en_test --num-negatives 4 --num-confusable-negatives 2 --max-train-examples-per-tag 5000 --max-eval-examples-per-tag 50 --seed 12 --overwrite --dataset-shape per_tag --class-balance-beta 0.99

Some starter code that will change:

uv run processing_scripts/train_wsd_model.py --no-push-to-hub --attention-implementation "kernels-community/flash-attn2@v3" --data-dir data/training_data/en

uv run processing_scripts/train_wsd_model.py --no-push-to-hub --attention-implementation "kernels-community/flash-attn2@v3" --data-dir data/training_data/en_test --model-max-seq-length 512 --loss-type masked-mnrl --project-name usas-wsd-trial --run-name usas_wsd_trial_ettin_encoder_68m

Languages

The languages that this repository covers and supports, of which this table is also available in machine readable format at ./data/languages.yaml (languages that have the value of True for the key training). These languages have been selected based on semantic tagging support for the given language whereby in some cases setting up the semantic tagger for a given language can be difficult within a large scale tagging pipeline in addition some languages have very few to none GA or FA articles.

Language	ISO 639-3
English	eng
Dutch	nld
Spanish	spa
Danish	dan
Italian	ita
Portuguese	por
Chinese	zho
Finnish	fin

License

The code is licensed under Apache License Version 2.0.

Claude settings

For those that use Anthropic's Claude we have shared some suggested settings, see ./.claude folder that are enforced within this project but can be easily adjusted or removed if you prefer to use your own settings or the default settings of Claude. The project level settings for Claude, can be found at ./.claude/settings.json are auto generated by running the following script;

cd .claude/hooks && uv run generate_settings.py > ../settings.json

This script creates a settings file with;

• Numerous Deny permissions that have come from the list of files, stated in ./.claude/hooks/sensitive_patterns.py, that you do not want Claude to write/edit/read.
• A pre-hook, ./.claude/hooks/block_sensitive_files.py, that catches any write/edit/read to the list of files that the Deny permissions might miss, e.g. a call to Python using Bash.

To note this pre-hook and Deny permissions would not stop Claude from write/edit/read if Claude requests the file through an unusual regex pattern like e*v to get the .env file, but this is a best effort try to reduce Claude's access to these more sensitive files. Generally speaking if you are using API keys reduce the scope as much as possible and limit the time and resource access while developing.