Skip to content

merantix-momentum/mwl-finetuning

Repository files navigation

More-With-Less (MWL) Fine-Tuning Library

This repository provides a modular Python library for parameter-efficient fine-tuning of Large Language Models (LLMs) using LoRA. It was developed as part of the BMBF-funded research project "More-with-Less: Effiziente Sprachmodelle für KMUs" (grant no. 01IS23013A), which investigates how small and medium-sized enterprises can leverage open LLMs with limited computational resources.

The library is designed around two abstract interfaces (DatasetGenerator for data preparation and ModelEvaluator for task-specific evaluation) so that new use cases can be integrated without modifying the core training infrastructure. As a concrete application, the repository includes a Legal Tech demonstrator that extracts structured metadata (decision type, jurisdiction, law references, case references) from German court decisions, using the publicly available harshildarji/openlegaldata dataset with over 251,000 documents.

Setup

Requires Python 3.11+. Uses uv for dependency management.

# Install dependencies and create virtual environment
uv sync

# Optional: include vLLM for fast inference, Gradio for the demo UI
uv sync --extra all

# Set up pre-commit hooks
uv run pre-commit install

# Configure environment variables
cp .env.example .env
# Edit .env: WANDB_API_KEY, HF_TOKEN, HF_HOME, PROJECT_ROOT

To add a new dependency:

uv add <package>              # runtime dependency
uv add --dev <package>        # development dependency

Quick Start

Training

Run a LoRA fine-tuning experiment. All configuration is managed via hydra-zen, composing presets for the model, dataset, training schedule, and evaluation. To perform fine-tuning for the Legal Tech demonstrator run the following:

uv run python -m run.workflows.exec workflow=default \
    run=openlegaldata \
    cfg=openlegaldata

The default workflow performs fine-tuning followed by evaluation. Trained LoRA weights are saved under artifacts/runs/.

Choosing a Base Model

The framework ships with several model presets defined in run/conf/model.py. To switch the base model, override the cfg/model group:

# Qwen3-8B (default for the legal demonstrator)
uv run python -m run.workflows.exec workflow=default \
    run=openlegaldata cfg=openlegaldata cfg/model=qwen3_8b

# LLaMA 3.1 8B (base)
uv run python -m run.workflows.exec workflow=default \
    run=openlegaldata cfg=openlegaldata cfg/model=llama31_8b_base

# LLaMA 3.1 8B (instruct)
uv run python -m run.workflows.exec workflow=default \
    run=openlegaldata cfg=openlegaldata cfg/model=llama31_8b_instruct

The modular architecture makes it straightforward to add new models. Simply register a new ModelConfig preset in run/conf/model.py.

Evaluation

Evaluate a fine-tuned checkpoint by pointing to the saved LoRA weights. The eval_only workflow loads the base model, applies the adapter, and runs the task-specific evaluator:

uv run python -m run.workflows.exec workflow=eval_only \
    run=openlegaldata \
    cfg=openlegaldata \
    workflow.build.lora_path=artifacts/runs/.../lora_weights

For a zero-shot baseline without any LoRA adapter:

uv run python -m run.workflows.exec workflow=eval_only \
    run=openlegaldata \
    cfg=openlegaldata

Gradio Demo

An interactive web UI built with Gradio allows non-technical users to try the metadata extraction on arbitrary court decision texts:

uv run python -m scripts.gradio_demo --model Qwen/Qwen3-8B --lora-path artifacts/runs/.../lora_weights

Project Structure

mwl/                       Core library
├── data/                  Dataset generators and prompt formatting
│   ├── base.py            Abstract DatasetGenerator / DatasetConfig
│   ├── formatting.py      Task-agnostic PromptFormatter (chat + Alpaca)
│   └── openlegaldata.py   OpenLegalData dataset for German court decisions
├── model/                 Model and tokenizer loading (ModelConfig)
├── training/              HF Trainer + LoRA pipeline, mid-training eval callback
├── eval/                  Evaluation metrics and ModelEvaluator base class
│   ├── evaluator.py       Generic ModelEvaluator base
│   └── legal_metadata.py  Task-specific LegalMetadataEvaluator
├── inference/             vLLM and HF inference engines
└── utils.py               Shared utilities (dtype conversion, JSON parsing)

run/                       Experiment orchestration (hydra-zen)
├── conf/                  Configuration dataclasses (model, data, training, eval)
├── workflows/             Workflow functions (train, evaluate, build)
└── hydra_zen.py           Custom hydra-zen builds and store utilities

scripts/                   Standalone scripts
└── gradio_demo.py         Interactive Gradio demo

tests/                     Test suite (pytest)

Code Quality

uv run pre-commit run --all-files   # All linting + formatting
uv run pytest                       # Run tests

Adding a New Task

  1. Dataset: Subclass DatasetGenerator (see mwl/data/base.py). Implement train_dataset(), val_dataset(), format_example(), and get_system_prompt(). Each example must include input_text, labels, and system_prompt fields.

  2. Evaluator: Subclass ModelEvaluator (see mwl/eval/evaluator.py). Implement __call__(model, prefix) returning a metrics dict. See mwl/eval/legal_metadata.py for a concrete example.

  3. Model config (optional): Add a new model preset in run/conf/model.py following the existing entries for Qwen3 and LLaMA.

  4. Config registration: Register hydra-zen configs in run/conf/ (see run/conf/data.py and run/conf/eval.py for examples).

License

See LICENSE.

Acknowledgement

We kindly acknowledge funding by the German Federal Ministry of Education and Research (BMBF) within the project "More-with-Less: Effiziente Sprachmodelle fuer KMUs" (grant no. 01IS23013A).

About

Fine-tuning library and demonstrator for the More-With-Less Project

Resources

License

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages