ML4SE-toolkit

ML4SE-toolkit provides small, composable primitives for turning source code into reproducible machine-learning examples for software engineering research.

What is in scope

• Comment-oriented parsing utilities that return a normalized QueryMatch(prefix, suffix, match) contract.
• Comment sanitization utilities that strip comment syntax after extraction.
• Input formatters for fill-in-the-middle, autoregressive, masked diffusion, masked LM, and span corruption models, with sentinel and mask tokens picked from a model name or Hugging Face object.
• A restartable, multi-worker iterator that applies a function to every row of a dataset split into many files, such as The Stack, one file at a time.
• Optional integrations for Tree-sitter queries and PyTorch-style iterable datasets.
• Existing EBNF and grammar-analysis tooling remains available in the repository under src/ml4setk/EBNF; it is not part of the PyPI package.

Installation

Install the published package from PyPI:

pip install ML4SE-ToolKit

The import name is ml4setk, for example from ml4setk import CommentQuery.

Install optional extras only when you need them:

pip install "ML4SE-ToolKit[treesitter]"
pip install "ML4SE-ToolKit[torch]"
pip install "ML4SE-ToolKit[datasets]"

Development setup

Use uv so package state stays locked to the repository:

uv venv .venv
uv sync --group dev

Install optional extras only when you need them:

uv sync --group dev --extra treesitter
uv sync --group dev --extra torch

Build

Build source and wheel distributions locally with:

uv build

Or via the repo command:

make build

Artifacts are written to dist/.

To validate the built distributions before uploading:

uvx twine check dist/*

Publishing

The repository includes a dedicated GitHub Actions workflow at .github/workflows/publish.yml.

• workflow_dispatch: build, validate, and publish the current branch to TestPyPI
• push of a tag matching v*: build, validate, and publish to PyPI

The workflow uses PyPI Trusted Publishing rather than a long-lived API token. Before the first automated release, configure trusted publishers for this project in both PyPI and TestPyPI with:

• owner: jkatzy
• repository: ML4SE-toolkit
• workflow name: publish.yml
• environment: pypi for PyPI, testpypi for TestPyPI

Release flow:

git switch main
python scripts/check_release_version.py
uv run pytest -m "not optional_dependency"
uv run pytest -m "optional_dependency" --no-cov
uv run ruff check src tests examples
git tag v0.1.0
git push origin main v0.1.0

The publish workflow checks that the package version in pyproject.toml, the package __version__, and the release tag all agree before uploading.

Minimal example

from ml4setk import CommentQuery, FIMInput

sample = "prefix\n// explain this branch\nsuffix"
match = CommentQuery("java").parse(sample)[0]

model_input, ground_truth = FIMInput(
    "<fim_prefix>",
    "<fim_suffix>",
    "<fim_middle>",
).generate(match)

print(model_input)
print(ground_truth)

Pass FIMInput.from_model("bigcode/starcoder2-3b") instead of the token strings to use a model's own sentinel tokens. The input formatting guide covers every input format and the models each one supports.

The parsing contract is stable across the core query implementations:

• prefix: text before the match
• suffix: text after the match
• match: the extracted region itself

For the comment extraction and cleaning guide, with a runnable quick start, an explanation of how extraction and cleaning work, and the API reference, see docs/comment_extractor.md.

At the moment the comment extractor covers 776 registry keys, including template and document syntaxes such as astro, coldfusion, jsp, marko, plantuml, slim, smarty, and restructuredtext, plus nested-comment languages such as coldfusion, dafny, frege, grammatical_framework, rexx, tla, and v.

Format-aware coverage includes GNU checksums, ecere_projects, figlet_font, microsoft_visual_studio_solution, AMPL nl, omgrofl, and pogoscript. These entries use narrowly scoped structural or syntax-aware guards where marker-only scanning would misclassify strings, counted records, or binary data.

Development commands

• make setup: create and sync the local development environment with uv
• make test: run the default unit and integration suite with coverage
• make test-optional: run optional-dependency tests without coverage gating
• make comment-fuzz: run deterministic Unicode-heavy parser invariants across every registered language
• make lint: run ruff on maintained source, tests, and examples; the legacy EBNF utility tree keeps its existing separate style baseline
• make smoke: run the minimal end-to-end smoke tests
• make build: create source and wheel distributions in dist/
• make docs: build the MkDocs site into site/ and fail on broken links; make docs-serve previews it at http://127.0.0.1:8000. The Docs workflow publishes the site from main to GitHub Pages.
• make check-main-branch: fail if development-only artifacts are still tracked and the branch is not ready to merge to main
• make check-release-version: fail if package version metadata is inconsistent

Git workflow

main is the release branch and dev is the single long-lived integration branch. Agent instructions, comment-syntax research, adversarial testing guides, and LLM-judge operations live on dev; release code, deterministic regressions, and stable user documentation live on main. Promote tested code from dev to main, then remove any short-lived topic branch used for review. The permanent rules and the enforced main-branch policy are documented in docs/git_workflow.md.

Repository map

• src/ml4setk/Parsing: query primitives, the comment registry, the sanitizer, and optional Tree-sitter support
• src/ml4setk/Generation: input formatters and the model token registries
• src/ml4setk/Datasets: the chunked, restartable dataset iterator
• src/ml4setk/EBNF: grammar-processing and railroad-diagram utilities
• tests: regression, unit, integration, and smoke coverage
• docs/architecture.md: concise architecture and extension notes
• docs/comment_extractor.md: how comment extraction and cleaning work and how to use them
• docs/comment_cleaning_policy.md: what cleaning removes and what it keeps
• docs/input_formatting.md: every input format and its supported models
• docs/chunked_datasets.md: iterating over large file-chunked datasets
• docs/git_workflow.md: branch policy, development-only artifacts, and the main-branch guard
• dev comment workflow: syntax research, adversarial tests, fuzzing, LLM judges, and regression intake