Metadata-Version: 2.4
Name: ML4SE-ToolKit
Version: 0.1.0
Summary: A package aimed at implementing core functionality for ML4SE to simplify and standardize research
Author-email: Jonathan Katzy <J.B.Katzy@TUDelft.nl>
License-Expression: MIT
Project-URL: Homepage, https://github.com/jkatzy/ML4SE-toolkit
Project-URL: Documentation, https://jkatzy.github.io/ML4SE-toolkit/
Project-URL: Issues, https://github.com/jkatzy/ML4SE-toolkit/issues
Keywords: machine learning,software engineering,code comments,fill-in-the-middle
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Operating System :: OS Independent
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: regex>=2026.1.15
Provides-Extra: treesitter
Requires-Dist: tree-sitter<0.22,>=0.21.3; extra == "treesitter"
Requires-Dist: tree-sitter-languages<2,>=1.10.2; extra == "treesitter"
Provides-Extra: datasets
Requires-Dist: huggingface_hub>=0.30; extra == "datasets"
Requires-Dist: pyarrow>=17; extra == "datasets"
Provides-Extra: torch
Requires-Dist: torch>=2.8.0; extra == "torch"
Dynamic: license-file

# 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:

```bash
pip install ML4SE-ToolKit
```

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

Install optional extras only when you need them:

```bash
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:

```bash
uv venv .venv
uv sync --group dev
```

Install optional extras only when you need them:

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

## Build

Build source and wheel distributions locally with:

```bash
uv build
```

Or via the repo command:

```bash
make build
```

Artifacts are written to `dist/`.

To validate the built distributions before uploading:

```bash
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:

```bash
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

```python
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](https://github.com/jkatzy/ML4SE-toolkit/blob/main/docs/input_formatting.md)
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](https://github.com/jkatzy/ML4SE-toolkit/blob/main/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](https://github.com/jkatzy/ML4SE-toolkit/blob/main/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](https://github.com/jkatzy/ML4SE-toolkit/tree/dev/docs/comment_testing):
  syntax research, adversarial tests, fuzzing, LLM judges, and regression intake
