Train system and general refactoring (Machine-Learning-for-Medical-Language#230)

* move examples to examples directory

* move models to models module

* create args and data modules

* add train command to cli

* update submodules tests

* wip train system refactor

* continue train_system refactor, other misc refactoring

* data processing, train system, logging

* model saving, logging, display, tests

* update ruff version

* simplify build process with newer uv version

* fix lint/format errors

* fix pre-commit error

* run pre-commit on all files in `make check`

* fix a few type issues

* simplify train system invocation from cli

* add logfile path to display

* move new tests to test directory

* move notebooks to examples

* add todos for evaluation and prediction

* fix some more type issues

* support python 12 and 13

* consistent naming for args classes

* remove todos

* remove redundant set call

* move old code to `cnlpt.legacy`, rework docs

* remove some refactoring comments

* restore legacy code

* fix REST api data preprocessing

* fix `cnlpt train` help message

* set default python version to 3.13 and test newer python versions in CI

* update README.md

* minor version bumps in lockfile

* set min torch version to 2.6

* drop python 3.13 support due to windows issue

* minor args refactoring

* fix pin memory warning

* only disable mps for tests in CI

* refactor metrics, include transformers logging in logfile

* display name of best checkpoint

* update example READMEs

* support averaging multiple selection metrics

* add some data tests

* add prediction and evaluation to CnlpTrainSystem

* close logfile in train system test

* shutdown logging after train system test, cache HF models in CI

* use close instead of shutdown

* file handler is on root logger, not train system logger

* oops

* add tokens to CnlpPredictions

* json serialization for CnlpPredictions

* ensure "None" label has id 0 for relations tasks

* overwrite predictions file if `overwrite_output_dir` is true

* add data.analysis module

* fix metric averaging

* extract relations in analysis dataframe

* fix chemprot preprocessing

* fix broken import

* fix a couple training display issues

* temporarily remove results and error analysis stuff from chemprot readme

* add a --version arg for docker builds

Author: ianbulovic

.github/workflows/build-and-test.yml

@@ -23,7 +23,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.9", "3.10", "3.11"]
+        python-version: ["3.9", "3.10", "3.11", "3.12"]
         os: [ubuntu-latest, macos-latest, windows-latest]
 
     runs-on: ${{ matrix.os }}
@@ -32,9 +32,17 @@ jobs:
       - uses: actions/checkout@v4
       - name: Install uv
         id: setup-uv
-        uses: astral-sh/setup-uv@v3
+        uses: astral-sh/setup-uv@v6
         with:
+          version: "0.7.11"
           enable-cache: true
+      - name: Cache HF models
+        uses: actions/cache@v3
+        with:
+          path: ~/.cache/huggingface
+          key: ${{ runner.os }}-hf-${{ hashFiles('**/*.py') }}
+          restore-keys: |
+            ${{ runner.os }}-hf-
       - name: Test with pytest
         run: |
-          uv run --frozen --group test -p ${{ matrix.python-version }} pytest test/
+          uv run --frozen --group test -p ${{ matrix.python-version }} pytest -vv test/

.github/workflows/lint-and-format.yml

@@ -25,9 +25,9 @@ jobs:
       - uses: actions/checkout@v4
       - uses: astral-sh/ruff-action@v1
         with:
-          version: 0.7.0
+          version: 0.11.8 # same as in pyproject.toml and pre-commit hooks
           args: check
       - uses: astral-sh/ruff-action@v1
         with:
-          version: 0.7.0
+          version: 0.11.8 # same as in pyproject.toml and pre-commit hooks
           args: format --check

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.7.0 # same as in pyproject.toml
+  rev: v0.11.8 # same as in pyproject.toml and CI
   hooks:
     # Run the linter.
     - id: ruff
@@ -15,4 +15,4 @@ repos:
       - id: check-yaml
       - id: debug-statements
       - id: detect-private-key
-      - id: end-of-file-fixer
+      - id: end-of-file-fixer

.python-version

@@ -0,0 +1 @@
+3.12

.readthedocs.yaml

@@ -1,31 +1,21 @@
-# .readthedocs.yaml
 # Read the Docs configuration file
 # See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
 
-# Required
 version: 2
 
-# Set the version of Python and other tools you might need
-build:
-  os: ubuntu-20.04
-  tools:
-    python: "3.8"
-    # You can also specify other tool versions:
-    # nodejs: "16"
-    # rust: "1.55"
-    # golang: "1.17"
-
-# Build documentation in the docs/ directory with Sphinx
 sphinx:
    configuration: docs/conf.py
 
-# If using Sphinx, optionally build your docs in additional formats such as PDF
-# formats:
-#    - pdf
-
-# Optionally declare the Python requirements required to build your docs
-python:
-   install:
-     - method: pip
-       path: .
-     - requirements: dev-requirements.txt
+build:
+   os: ubuntu-24.04
+   tools:
+      python: "3.12"
+   jobs:
+      pre_create_environment:
+         - asdf plugin add uv
+         - asdf install uv latest
+         - asdf global uv latest
+      create_environment:
+         - uv venv "${READTHEDOCS_VIRTUALENV_PATH}"
+      install:
+         - UV_PROJECT_ENVIRONMENT="${READTHEDOCS_VIRTUALENV_PATH}" uv sync --frozen --group docs

CONTRIBUTING.md

@@ -45,7 +45,9 @@ we have instructions for using [uv](https://github.com/astral-sh/uv)
 2. From the project's base directory, run:
 
    ```bash
-   uv sync --python 3.11 # 3.9 and 3.10 are also supported. uv will install dev dependencies by default.
+   # uv will use python 3.12 and install dev dependencies by default.
+   # If you prefer a different python version, use e.g. `uv sync -p 3.9`
+   uv sync 
    source .venv/bin/activate # activate the virtual environment
    ```
 
@@ -56,7 +58,7 @@ we have instructions for using [uv](https://github.com/astral-sh/uv)
 2. Create a new conda environment:
 
    ```bash
-   conda create -n cnlpt python=3.11 # 3.9 and 3.10 are also supported
+   conda create -n cnlpt python=3.12 # 3.9, 3.10, and 3.11 are also supported
    conda activate cnlpt
    ```
 
@@ -111,7 +113,7 @@ The `lint-and-format` workflow should always pass if `make check` reports
 that everything is correct.
 
 The `build-and-test` workflow will run `pytest` on Linux, MacOS, and Windows,
-for each Python version this project supports (currently 3.9, 3.10, and 3.11).
+for each Python version this project supports (currently 3.9, 3.10, 3.11, and 3.12).
 
 You can see the structure of these CI runs in the
 [**Actions**](https://github.com/Machine-Learning-for-Medical-Language/cnlp_transformers/actions)
@@ -168,8 +170,7 @@ the Semantic Versioning guidelines. The key points are as follows:
 
 When the codebase is ready for release, run `python scripts/prepare_release.py <new version number>`.
 This will walk you through the last few changes that need to be made before release,
-including updating the changelog and setting the setuptools_scm fallback version,
-and will also update the lockfile and your venv with the new package version.
+including updating the changelog and the lockfile.
 
 > [!WARNING]
 > `prepare_release.py` requires uv to update the lockfile.

Makefile

@@ -5,7 +5,7 @@ help:
 	@echo '  hooks  - install pre-commit hooks'
 	@echo '  check  - lint and format using ruff'
 	@echo '  test   - run tests with pytest'
-	@echo '  docs   - build the docs'
+	@echo '  docs   - build the docs locally'
 	@echo '  build  - build cnlp-transformers for distribution'
 
 .PHONY: hooks
@@ -16,19 +16,16 @@ hooks:
 check:
 	ruff check --fix
 	ruff format
+	pre-commit run -a
 
 .PHONY: test
 test:
 	pytest test/
 
 .PHONY: docs
 docs:
-# this script is copied from the old build_doc_source.sh script
-	find docs -maxdepth 1 ! -name 'index.rst' -name '*.rst' -type f -exec rm -f {} +
-	rm -f transformer_objects.inv
-	sphobjinv convert zlib docs/transformer_objects.txt --quiet
-	SPHINX_APIDOC_OPTIONS=members,show-inheritance sphinx-apidoc -feT -o docs src/cnlpt
-	echo "   :noindex:" >> docs/cnlpt.rst
+	scripts/build_html_docs.sh docs/build
+	@echo "Point your browser at file://${PWD}/docs/build/html/index.html to view."
 
 .PHONY: build
 build:

README.md

@@ -11,23 +11,28 @@ Primary use cases include
 This library is _not_ intended to serve as a place for clinical NLP applications to live. If you build something cool that uses transformer models that take advantage of our model definitions, the best practice is probably to rely on it as a library rather than treating it as your workspace. This library is also not intended as a deployment-ready tool for _scalable_ clinical NLP. There is a lot of interest in developing methods and tools that are smaller and can process millions of records, and this library can potentially be used for research along those line. But it will probably never be extremely optimized or shrink-wrapped for applications. However, there should be plenty of examples and useful code for people who are interested in that type of deployment.
 
 ## Install
-> [!WARNING]
-macOS support is currently experimental. We recommend using python3.10 for macOS installations.
 
-> [!NOTE] 
-When installing the library's dependencies, `pip` will probably install 
-PyTorch with CUDA 10.2 support by default. If you would like to run the 
-library in CPU-only mode or with a newer version of CUDA, [install PyTorch 
-to your desired specifications](https://pytorch.org/get-started/locally/) 
-in your virtual environment first before installing `cnlp-transformers`.
+> [!IMPORTANT]
+> When installing the library's dependencies, PyTorch will probably be installed
+> with CUDA 12.6 support by default on linux, and without CUDA support on other platforms.
+> If you would like to run the library in CPU-only mode or with a specific version of CUDA,
+> [install PyTorch to your desired specifications](https://pytorch.org/get-started/locally/)
+> in your virtual environment first before installing `cnlp-transformers`.
+> [See here](https://docs.astral.sh/uv/guides/integration/pytorch/#the-uv-pip-interface) if
+> using uv.
 
 ### Static installation
 
 If you are installing just to fine-tune or run the REST APIs,
-you can install without cloning:
+you can install without cloning using [uv](https://docs.astral.sh/uv/):
+
+```sh
+uv pip install cnlp-transformers
+```
+
+Or with pip:
 
 ```sh
-# Note: if needed, install PyTorch first (see above)
 pip install cnlp-transformers
 ```
 
@@ -110,18 +115,18 @@ We provided the following step-by-step examples how to finetune in clinical NLP
 
 ### Fine-tuning options
 
-Run ```python -m cnlpt.train_system -h``` to see all the available options. In addition to inherited Huggingface Transformers options, there are options to do the following:
+Run `cnlpt train -h` to see all the available options. In addition to inherited Huggingface Transformers options, there are options to do the following:
 
-* Select different models: ```--model hier``` uses a hierarchical transformer layer on top of a specified encoder model. We recommend using a very small encoder: ```--encoder microsoft/xtremedistil-l6-h256-uncased``` so that the full model fits into memory.
+* Select different models: `--model hier` uses a hierarchical transformer layer on top of a specified encoder model. We recommend using a very small encoder: `--encoder microsoft/xtremedistil-l6-h256-uncased` so that the full model fits into memory.
 * Run simple baselines (use ``--model cnn|lstm --tokenizer_name roberta-base`` -- since there is no HF model then you must specify the tokenizer explicitly)
-* Use a different layer's CLS token for the classification (e.g., ```--layer 10```)
-* Probabilistically freeze weights of the encoder (leaving classifier weights all unfrozen) (```--freeze``` alone freezes all encoder weights, ```--freeze <float>``` when given a parameter between 0 and 1, freezes that percentage of encoder weights)
-* Classify based on a token embedding instead of the CLS embedding (```--token``` -- applies to the event/entity classification setting only, and requires the input to have xml-style tags (`<e>`, `</e>`) around the tokens representing the event/entity)
-* Use class-weighted loss function (```--class_weights```)
+* Use a different layer's CLS token for the classification (e.g., `--layer 10`)
+* Probabilistically freeze weights of the encoder (leaving classifier weights all unfrozen) (`--freeze` alone freezes all encoder weights, `--freeze <float>` when given a parameter between 0 and 1, freezes that percentage of encoder weights)
+* Classify based on a token embedding instead of the CLS embedding (`--token` -- applies to the event/entity classification setting only, and requires the input to have xml-style tags (`<e>`, `</e>`) around the tokens representing the event/entity)
+* Use class-weighted loss function (`--class_weights`)
 
 ## Running REST APIs
 
-There are existing REST APIs in the ```src/cnlpt/api``` folder for a few important clinical NLP tasks:
+There are existing REST APIs in the `src/cnlpt/api` folder for a few important clinical NLP tasks:
 
 1. Negation detection
 2. Time expression tagging (spans + time classes)
@@ -133,7 +138,7 @@ There are existing REST APIs in the ```src/cnlpt/api``` folder for a few importa
 To demo the negation API:
 
 1. Install the `cnlp-transformers` package.
-2. Run `cnlpt_negation_rest [-p PORT]`.
+2. Run `cnlpt rest --model-type negation [-p PORT]`.
 3. Open a python console and run the following commands:
 
 #### Setup variables for negation
@@ -167,7 +172,7 @@ The model correctly classifies both nausea and anosmia as negated.
 To demo the temporal API:
 
 1. Install the `cnlp-transformers` package.
-2. Run `cnlpt_temporal_rest [-p PORT]`
+2. Run `cnlpt rest --model-type temporal [-p PORT]`
 3. Open a python console and run the following commands to test:
 
 #### Setup variables for temporal
@@ -217,20 +222,14 @@ should return:
 
 This output indicates the token spans of events and timexes, and relations between events and timexes, where the suffixes are indices into the respective arrays (e.g., TIMEX-0 in a relation refers to the 0th time expression found, which begins at token 6 and ends at token 9 -- ["March 3, 2010"])
 
-To run only the time expression or event taggers, change the run command to:
-
-```uvicorn cnlpt.api.timex_rest:app --host 0.0.0.0``` or
-
-```uvicorn cnlpt.api.event_rest:app --host 0.0.0.0```
-
-then run the same process commands as above (including the same URL). You will get similar json output, but only one of the dictionary elements (timexes or events) will be populated.
-
 ## Citing cnlp_transformers
+
 Please use the following bibtex to cite cnlp_transformers if you use it in a publication:
-```
+
+```latex
 @misc{cnlp_transformers,
   author       = {CNLPT},
-  title = {Clinical {NLP} {Transformers} (cnlp\_transformers)},
+  title        = {Clinical {NLP} {Transformers} (cnlp\_transformers)},
   year         = {2021},
   publisher    = {GitHub},
   journal      = {GitHub repository},
@@ -239,14 +238,15 @@ Please use the following bibtex to cite cnlp_transformers if you use it in a pub
 ```
 
 ## Publications using cnlp_transformers
+
 Please send us any citations that used this library!
 
-1.  Chen S, Guevara M, Ramirez N, Murray A, Warner JL, Aerts HJWL, et al. Natural Language Processing to Automatically Extract the Presence and Severity of Esophagitis in Notes of Patients Undergoing Radiotherapy. JCO Clin Cancer Inform. 2023 Jul;(7):e2300048.
-2.  Li Y, Miller T, Bethard S, Savova G. Identifying Task Groupings for Multi-Task Learning Using Pointwise V-Usable Information [Internet]. arXiv.org. 2024 [cited 2025 May 22]. Available from: https://arxiv.org/abs/2410.12774v1
-3.  Wang L, Li Y, Miller T, Bethard S, Savova G. Two-Stage Fine-Tuning for Improved Bias and Variance for Large Pretrained Language Models. In: Rogers A, Boyd-Graber J, Okazaki N, editors. Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) [Internet]. Toronto, Canada: Association for Computational Linguistics; 2023 [cited 2025 May 22]. p. 15746–61. Available from: https://aclanthology.org/2023.acl-long.877/
-4.  Miller T, Bethard S, Dligach D, Savova G. End-to-end clinical temporal information extraction with multi-head attention. Proc Conf Assoc Comput Linguist Meet. 2023 Jul;2023:313–9. 
-5.  Yoon W, Ren B, Thomas S, Kim C, Savova G, Hall MH, et al. Aspect-Oriented Summarization for Psychiatric Short-Term Readmission Prediction [Internet]. arXiv; 2025 [cited 2025 May 22]. Available from: http://arxiv.org/abs/2502.10388
-6.  Wang L, Zipursky AR, Geva A, McMurry AJ, Mandl KD, Miller TA. A computable case definition for patients with SARS-CoV2 testing that occurred outside the hospital. JAMIA Open. 2023 Oct 1;6(3):ooad047. 
-7.  Bitterman DS, Goldner E, Finan S, Harris D, Durbin EB, Hochheiser H, et al. An End-to-End Natural Language Processing System for Automatically Extracting Radiation Therapy Events From Clinical Texts. Int J Radiat Oncol Biol Phys. 2023 Sep 1;117(1):262–73.
-8.  McMurry AJ, Gottlieb DI, Miller TA, Jones JR, Atreja A, Crago J, et al. Cumulus: A federated EHR-based learning system powered by FHIR and AI. medRxiv. 2024 Feb 6;2024.02.02.24301940. 
-9.  LCD benchmark: long clinical document benchmark on mortality prediction for language models | Journal of the American Medical Informatics Association | Oxford Academic [Internet]. [cited 2025 Jan 23]. Available from: https://academic.oup.com/jamia/article-abstract/32/2/285/7909835?redirectedFrom=fulltext
+1. Chen S, Guevara M, Ramirez N, Murray A, Warner JL, Aerts HJWL, et al. Natural Language Processing to Automatically Extract the Presence and Severity of Esophagitis in Notes of Patients Undergoing Radiotherapy. JCO Clin Cancer Inform. 2023 Jul;(7):e2300048.
+2. Li Y, Miller T, Bethard S, Savova G. Identifying Task Groupings for Multi-Task Learning Using Pointwise V-Usable Information [Internet]. arXiv.org. 2024 [cited 2025 May 22]. Available from: <https://arxiv.org/abs/2410.12774v1>
+3. Wang L, Li Y, Miller T, Bethard S, Savova G. Two-Stage Fine-Tuning for Improved Bias and Variance for Large Pretrained Language Models. In: Rogers A, Boyd-Graber J, Okazaki N, editors. Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) [Internet]. Toronto, Canada: Association for Computational Linguistics; 2023 [cited 2025 May 22]. p. 15746–61. Available from: <https://aclanthology.org/2023.acl-long.877/>
+4. Miller T, Bethard S, Dligach D, Savova G. End-to-end clinical temporal information extraction with multi-head attention. Proc Conf Assoc Comput Linguist Meet. 2023 Jul;2023:313–9.
+5. Yoon W, Ren B, Thomas S, Kim C, Savova G, Hall MH, et al. Aspect-Oriented Summarization for Psychiatric Short-Term Readmission Prediction [Internet]. arXiv; 2025 [cited 2025 May 22]. Available from: <http://arxiv.org/abs/2502.10388>
+6. Wang L, Zipursky AR, Geva A, McMurry AJ, Mandl KD, Miller TA. A computable case definition for patients with SARS-CoV2 testing that occurred outside the hospital. JAMIA Open. 2023 Oct 1;6(3):ooad047.
+7. Bitterman DS, Goldner E, Finan S, Harris D, Durbin EB, Hochheiser H, et al. An End-to-End Natural Language Processing System for Automatically Extracting Radiation Therapy Events From Clinical Texts. Int J Radiat Oncol Biol Phys. 2023 Sep 1;117(1):262–73.
+8. McMurry AJ, Gottlieb DI, Miller TA, Jones JR, Atreja A, Crago J, et al. Cumulus: A federated EHR-based learning system powered by FHIR and AI. medRxiv. 2024 Feb 6;2024.02.02.24301940.
+9. LCD benchmark: long clinical document benchmark on mortality prediction for language models | Journal of the American Medical Informatics Association | Oxford Academic [Internet]. [cited 2025 Jan 23]. Available from: <https://academic.oup.com/jamia/article-abstract/32/2/285/7909835?redirectedFrom=fulltext>

docker/build.py

@@ -16,9 +16,10 @@
 ]
 
 parser = argparse.ArgumentParser()
-parser.add_argument("--model", action="append", choices=["all"] + MODELS)
+parser.add_argument("--model", action="append", choices=["all", *MODELS])
 parser.add_argument("--processor", choices=["all", "cpu", "gpu"], default="all")
 parser.add_argument("--push", action="store_true", default=False)
+parser.add_argument("--cnlpt_version", default=None)
 args = parser.parse_args()
 
 
@@ -69,7 +70,7 @@ def build_one(model: str, processor: str, *, version: str, push: bool = False) -
     else:
         build_args.append("--load")  # to load into docker locally
 
-    subprocess.run(["docker", "buildx", "build"] + build_args, check=True)
+    subprocess.run(["docker", "buildx", "build", *build_args], check=True)
 
 
 if __name__ == "__main__":
@@ -86,7 +87,7 @@ def build_one(model: str, processor: str, *, version: str, push: bool = False) -
     # Our Dockerfiles pull directly from pip, so we want to be setting the same version as we'll install.
     # We don't want to pull the version from our sibling code in this repo, because it might not be released yet,
     # but we still want to be able to push new builds of the existing releases.
-    version = get_latest_pip_version("cnlp-transformers")
+    version = args.cnlpt_version or get_latest_pip_version("cnlp-transformers")
 
     for model in models:
         for processor in processors:

docker/compose.yaml

@@ -19,4 +19,3 @@ services:
           devices:
             - capabilities: [gpu]
     entrypoint: cnlpt_negation_rest -p 8000
-