codefinetuner 0.2.1__tar.gz → 0.4.0__tar.gz

{codefinetuner-0.2.1 → codefinetuner-0.4.0}/.gitignore

@@ -17,4 +17,5 @@ wheels/
 /data/ 
 backup/
 base_models/
+unsloth_compiled_cache/
 TODO.md

{codefinetuner-0.2.1 → codefinetuner-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: codefinetuner
-Version: 0.2.1
+Version: 0.4.0
 Summary: Add your description here
 Project-URL: Homepage, https://github.com/cuolm/codefinetuner
 Project-URL: Repository, https://github.com/cuolm/codefinetuner.git
@@ -23,12 +23,14 @@ Requires-Dist: nltk>=3.9.2
 Requires-Dist: numpy>=2.3.4
 Requires-Dist: omegaconf>=2.3.0
 Requires-Dist: peft>=0.17.1
+Requires-Dist: rapidfuzz>=3.14.5
 Requires-Dist: scikit-learn>=1.7.2
 Requires-Dist: sentencepiece>=0.2.1
 Requires-Dist: torch>=2.9.0
 Requires-Dist: transformers>=4.57.1
-Requires-Dist: tree-sitter-language-pack>=0.2.0
+Requires-Dist: tree-sitter-language-pack<1.6.3,>=0.2.0
 Requires-Dist: tree-sitter<0.23.0,>=0.22.0
+Requires-Dist: unsloth>=2026.4.8; sys_platform == 'linux' or sys_platform == 'win32'
 Description-Content-Type: text/markdown
 
 # CodeFinetuner
@@ -44,8 +46,8 @@ It trains a Low-Rank Adapter ([LoRA](https://arxiv.org/abs/2106.09685)) on Fill-
 - [Architecture](https://github.com/cuolm/codefinetuner/tree/master#architecture)
 - [Project Structure](https://github.com/cuolm/codefinetuner/tree/master#project-structure)
 - [How Training Examples Are Created](https://github.com/cuolm/codefinetuner/tree/master#how-training-examples-are-created)
-- [Installation](https://github.com/cuolm/codefinetuner/tree/master#installation)
 - [Quick Start](https://github.com/cuolm/codefinetuner/tree/master#quick-start)
+- [Installation](https://github.com/cuolm/codefinetuner/tree/master#installation)
 - [Configuration](https://github.com/cuolm/codefinetuner/tree/master#configuration)
 - [Usage](https://github.com/cuolm/codefinetuner/tree/master#usage)
 - [Finetuned Model Usage](https://github.com/cuolm/codefinetuner/tree/master#finetuned-model-usage)
@@ -69,7 +71,7 @@ Raw Code Files
 [Finetune]    -- LoRA adapter training -> merged safetensors model
      |
      v
-[Evaluate]    -- CodeBLEU, SentenceBLEU, exact match, line match, perplexity
+[Evaluate]    -- CodeBLEU, SentenceBLEU, edit similarity, exact match, line match, perplexity
      |
      v
 [Convert]      -- GGUF conversion -> quantized model for deployment
@@ -120,38 +122,43 @@ Here is an example illustrating how a single FIM example is created:
 <|fim_suffix|>    }\n    return count;
 <|fim_middle|>count = count + (value & 1);\n    value = (value >> 1);
 ```
+## Quick Start
+Install CodeFinetuner globally to run the pipeline anywhere on your system:
+```bash
+uv tool install codefinetuner
+```
+Create a configuration file according to the [Configuration](https://github.com/cuolm/codefinetuner/tree/master#configuration) section and run the pipeline:
+```bash
+codefinetuner --config="codefinetuner_config.yaml"
+```
 
 ## Installation
-
-### From PyPI
+ 
+### As a Global CLI Tool
+```bash
+uv tool install codefinetuner
+```
+ 
+### As a Library Dependency
 ```bash
 uv add codefinetuner
 # or
 pip install codefinetuner
 ```
+ 
 ### From Source
 ```bash
 git clone --recurse-submodules https://github.com/cuolm/codefinetuner
 cd codefinetuner
-
+ 
 # Using uv (Recommended)
 uv sync
-
+ 
 # Using pip
 pip install -r requirements.txt
 pip install -e .
 ```
 
----
-
-## Quick Start
-Create a configuration file according to the [Configuration](https://github.com/cuolm/codefinetuner/tree/master#configuration) section.
-```python
-import codefinetuner
-
-codefinetuner.run_pipeline("codefinetuner_config.yaml")
-```
-
 ## Configuration
 The pipeline uses a single-source-of-truth YAML configuration file. It utilizes YAML anchors (`&globals`) to share core parameters across all stages (`preprocess`, `finetune`, `evaluate`, `convert`), ensuring consistency and reducing redundancy.
 
@@ -170,13 +177,14 @@ globals: &globals
   fim_pad_token: "<|fim_pad|>"
   eos_token: "<|endoftext|>"
   label_pad_token_id: -100
+  max_token_sequence_length: 512
   data_language: "c"
   data_extensions: [".c", ".h"]
+  use_unsloth: True
 
 preprocess:
   <<: *globals  # inherits all global parameters
   split_mode: "manual"
-  max_token_sequence_length: 1024
   # ... (preprocess specific settings)
 
 finetune:
@@ -204,11 +212,16 @@ Place source files in your `raw_data_path` (default: `workspace_path/data`).
 ## Usage
 
 ### CLI Usage
-Run the pipeline using the unified CLI:
+If installed via `uv tool install`:
+```bash
+codefinetuner --config="codefinetuner_config.yaml"
+```
+ 
+If running within the source repository cloned from GitHub:
 ```bash
 uv run codefinetuner --config="config/codefinetuner_config.yaml"
 ```
-
+ 
 **Pipeline flags**
 - `--config`: Use a different config file.
 - `--skip-preprocess`: Skip preprocessing.
@@ -216,16 +229,17 @@ uv run codefinetuner --config="config/codefinetuner_config.yaml"
 - `--skip-evaluate`: Skip evaluation.
 - `--skip-convert`: Skip conversion.
 
+
 ### Python Module Usage
 ```python
 import codefinetuner
 
 # Full pipeline
-codefinetuner.run_pipeline("path/to/codefinetuner_config.yaml")
+codefinetuner.run_pipeline("codefinetuner_config.yaml")
 
 # Skip stages
 codefinetuner.run_pipeline(
-    "path/to/codefinetuner_config.yaml",
+    "codefinetuner_config.yaml",
     skip_preprocess=True,
     skip_convert=True
 )

{codefinetuner-0.2.1 → codefinetuner-0.4.0}/README.md

@@ -11,8 +11,8 @@ It trains a Low-Rank Adapter ([LoRA](https://arxiv.org/abs/2106.09685)) on Fill-
 - [Architecture](#architecture)
 - [Project Structure](#project-structure)
 - [How Training Examples Are Created](#how-training-examples-are-created)
-- [Installation](#installation)
 - [Quick Start](#quick-start)
+- [Installation](#installation)
 - [Configuration](#configuration)
 - [Usage](#usage)
 - [Finetuned Model Usage](#finetuned-model-usage)
@@ -36,7 +36,7 @@ Raw Code Files
 [Finetune]    -- LoRA adapter training -> merged safetensors model
      |
      v
-[Evaluate]    -- CodeBLEU, SentenceBLEU, exact match, line match, perplexity
+[Evaluate]    -- CodeBLEU, SentenceBLEU, edit similarity, exact match, line match, perplexity
      |
      v
 [Convert]      -- GGUF conversion -> quantized model for deployment
@@ -87,38 +87,43 @@ Here is an example illustrating how a single FIM example is created:
 <|fim_suffix|>    }\n    return count;
 <|fim_middle|>count = count + (value & 1);\n    value = (value >> 1);
 ```
+## Quick Start
+Install CodeFinetuner globally to run the pipeline anywhere on your system:
+```bash
+uv tool install codefinetuner
+```
+Create a configuration file according to the [Configuration](#configuration) section and run the pipeline:
+```bash
+codefinetuner --config="codefinetuner_config.yaml"
+```
 
 ## Installation
-
-### From PyPI
+ 
+### As a Global CLI Tool
+```bash
+uv tool install codefinetuner
+```
+ 
+### As a Library Dependency
 ```bash
 uv add codefinetuner
 # or
 pip install codefinetuner
 ```
+ 
 ### From Source
 ```bash
 git clone --recurse-submodules https://github.com/cuolm/codefinetuner
 cd codefinetuner
-
+ 
 # Using uv (Recommended)
 uv sync
-
+ 
 # Using pip
 pip install -r requirements.txt
 pip install -e .
 ```
 
----
-
-## Quick Start
-Create a configuration file according to the [Configuration](#configuration) section.
-```python
-import codefinetuner
-
-codefinetuner.run_pipeline("codefinetuner_config.yaml")
-```
-
 ## Configuration
 The pipeline uses a single-source-of-truth YAML configuration file. It utilizes YAML anchors (`&globals`) to share core parameters across all stages (`preprocess`, `finetune`, `evaluate`, `convert`), ensuring consistency and reducing redundancy.
 
@@ -137,13 +142,14 @@ globals: &globals
   fim_pad_token: "<|fim_pad|>"
   eos_token: "<|endoftext|>"
   label_pad_token_id: -100
+  max_token_sequence_length: 512
   data_language: "c"
   data_extensions: [".c", ".h"]
+  use_unsloth: True
 
 preprocess:
   <<: *globals  # inherits all global parameters
   split_mode: "manual"
-  max_token_sequence_length: 1024
   # ... (preprocess specific settings)
 
 finetune:
@@ -171,11 +177,16 @@ Place source files in your `raw_data_path` (default: `workspace_path/data`).
 ## Usage
 
 ### CLI Usage
-Run the pipeline using the unified CLI:
+If installed via `uv tool install`:
+```bash
+codefinetuner --config="codefinetuner_config.yaml"
+```
+ 
+If running within the source repository cloned from GitHub:
 ```bash
 uv run codefinetuner --config="config/codefinetuner_config.yaml"
 ```
-
+ 
 **Pipeline flags**
 - `--config`: Use a different config file.
 - `--skip-preprocess`: Skip preprocessing.
@@ -183,16 +194,17 @@ uv run codefinetuner --config="config/codefinetuner_config.yaml"
 - `--skip-evaluate`: Skip evaluation.
 - `--skip-convert`: Skip conversion.
 
+
 ### Python Module Usage
 ```python
 import codefinetuner
 
 # Full pipeline
-codefinetuner.run_pipeline("path/to/codefinetuner_config.yaml")
+codefinetuner.run_pipeline("codefinetuner_config.yaml")
 
 # Skip stages
 codefinetuner.run_pipeline(
-    "path/to/codefinetuner_config.yaml",
+    "codefinetuner_config.yaml",
     skip_preprocess=True,
     skip_convert=True
 )

{codefinetuner-0.2.1 → codefinetuner-0.4.0}/config/codefinetuner_config.yaml

@@ -1,32 +1,32 @@
-# config/pipeline.yaml - Single Source of Truth
 # Shared across all stages 
 globals: &globals
   workspace_path: null  # null: defaults to current working directory (CWD)
-  model_name: "Qwen/Qwen2.5-0.5B" 
+  model_name: "unsloth/Qwen2.5-Coder-3B" 
   fim_prefix_token: "<|fim_prefix|>"
   fim_middle_token: "<|fim_middle|>"
   fim_suffix_token: "<|fim_suffix|>"
   fim_pad_token: "<|fim_pad|>"
   eos_token: "<|endoftext|>"
   label_pad_token_id: -100
-  data_language: "c"
-  data_extensions: [".c", ".h"]
+  max_token_sequence_length: 512 # 1024
+  data_language: "iec61131_3_st"
+  data_extensions: [".st"]
+  use_unsloth: False
 
 # Preprocess stage
 preprocess:
   <<: *globals
   split_mode: "manual"
-  train_ratio: 0.8
-  eval_ratio: 0.1
-  test_ratio: 0.1
-  max_token_sequence_length: 1024
-  max_code_blocks_ast_depth: 2
-  min_middle_tokens_length: 20
+  train_ratio: 0.8  # ignored when split mode manual
+  eval_ratio: 0.1  # ignored when split mode manual  
+  test_ratio: 0.1  # ignored when split mode manual
+  max_code_blocks_ast_depth: 20
+  min_middle_tokens_length: 10
   max_middle_tokens_length: 200
-  fim_examples_per_subblock_ratio: 1.0
+  fim_examples_per_subblock_ratio: 1
   tokenizer_batch_size: 32
-  raw_data_path: "data/data_embedded"  # null: defaults to <workspace_path>/data
-  tree_sitter_parser_path: null  # null: uses tree_sitter_language_pack parser
+  raw_data_path: data/PLC-lang-oscat  # null: defaults to <workspace_path>/data
+  tree_sitter_parser_path: third_party/tree-sitter-iec61131-3-st/libtree-sitter-iec61131-3-st.dylib # null: uses tree_sitter_language_pack parser
   tree_sitter_definitions_path: null  # null: uses package internal tree sitter definitions 
   rng_seed: 0
 
@@ -36,26 +36,27 @@ finetune:
   model_attn_implementation: "sdpa"
   lora_r: 32
   lora_alpha: 64
-  lora_dropout: 0.1
+  lora_dropout: 0.0
   lora_bias: "none"
   lora_target_modules: ["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "down_proj", "up_proj"]
+  selected_checkpoint_strategy : "best"  # "best" or "last"
   trainer_resume_from_checkpoint: null  # null->fresh start, "last"->take last checkpoint
   trainer_clear_checkpoint_dir: false
   trainer_num_train_epochs: 1
-  trainer_per_device_train_batch_size: 2
-  trainer_per_device_eval_batch_size: 2
-  trainer_gradient_accumulation_steps: 32
-  trainer_learning_rate: 2e-5
-  trainer_weight_decay: 0.1
+  trainer_per_device_train_batch_size: 1
+  trainer_per_device_eval_batch_size: 1 
+  trainer_gradient_accumulation_steps: 16 
+  trainer_learning_rate: 5e-5
+  trainer_weight_decay: 0.01
   trainer_max_grad_norm: 1.0
   trainer_lr_scheduler_type: "cosine"
-  trainer_warmup_steps: 50
+  trainer_warmup_steps: 10
   trainer_gradient_checkpointing: true
   trainer_logging_steps: 10
   trainer_eval_strategy: "steps"
-  trainer_eval_steps: 100
+  trainer_eval_steps: 10
   trainer_save_strategy: "steps"
-  trainer_save_steps: 100
+  trainer_save_steps: 10
   trainer_logging_strategy: "steps"
   dataset_shuffle_buffer_size: 50000
   dataset_shuffle_seed: 0
@@ -63,25 +64,24 @@ finetune:
 # Evaluate stage
 evaluate:
   <<: *globals
-  benchmark_sample_size: 4
-  benchmark_min_fim_middle_tokens: 0
+  benchmark_sample_size: 500 
   benchmark_shuffle_buffer_size: 10000000
   benchmark_shuffle_seed: 42
+  generation_checkpoint: "pipeline"  # "pipeline" -> selected checkpoint from finetune stage, "checkpoint-name" -> specific checkpoint
+  generation_batch_size: 5 
   generation_max_new_tokens: 128
   generation_do_sample: false
-  generation_temperature: 0.7
-  generation_top_p: 0.95
-  codebleu_language: "c"
-  codebleu_ngram_weight: 0.25
-  codebleu_weighted_ngram_weight: 0.25
-  codebleu_syntax_ast_weight: 0.25
-  codebleu_dataflow_weight: 0.25
+  generation_temperature: 0.1  # only active if generation_do_sample: true
+  generation_top_p: 0.99  # only active if generation_do_sample: true
+  codebleu_ngram_weight: 0.1
+  codebleu_weighted_ngram_weight: 0.1
+  codebleu_syntax_ast_weight: 0.4
+  codebleu_dataflow_weight: 0.4
   sentencebleu_ngram_weight_1: 0.25
   sentencebleu_ngram_weight_2: 0.25
   sentencebleu_ngram_weight_3: 0.25
   sentencebleu_ngram_weight_4: 0.25
   line_match_number_of_lines: 2
-  trainer_checkpoint: "last"
   plot_only: false
   benchmark_use_existing_dataset: false
 

{codefinetuner-0.2.1 → codefinetuner-0.4.0}/docs/config-file.md

@@ -14,8 +14,10 @@ This document defines all mandatory and optional parameters in the YAML config f
 | `fim_pad_token` | str | `"<\|fim_pad\|>"` | FIM padding token. |
 | `eos_token` | str | `"<\|endoftext\|>"` | End-of-sequence token. |
 | `label_pad_token_id` | int | `-100` | Token ID ignored in loss calculation. |
+| `max_token_sequence_length` | int | `1024` | Maximum tokens per training example. |
 | `data_language` | str | `"c"` | Tree-sitter language identifier. |
 | `data_extensions` | list | `[".c", ".h"]` | File extensions to include in data preprocessing. |
+| `use_unsloth` | bool | `false` | Enables Unsloth optimizations for faster training and lower VRAM. Requires CUDA.
 
 ## Preprocess Parameters (Optional)
 *Controls how raw code is converted into FIM training examples.*
@@ -26,7 +28,6 @@ This document defines all mandatory and optional parameters in the YAML config f
 | `train_ratio` | float | `0.8` | Training set ratio (used in `auto` split). |
 | `eval_ratio` | float | `0.1` | Validation set ratio (used in `auto` split). |
 | `test_ratio` | float | `0.1` | Test set ratio (used in `auto` split). |
-| `max_token_sequence_length` | int | `1024` | Maximum tokens per training example. |
 | `max_code_blocks_ast_depth` | int | `2` | Tree-Sitter AST depth limit for block extraction. Depth 1 is root, 2 includes child nodes (e.g. functions). |
 | `min_middle_tokens_length` | int | `20` | Minimum tokens required in the FIM "middle" section of an example. |
 | `max_middle_tokens_length` | int | `200` | Maximum tokens allowed in the FIM "middle" section of an example. |
@@ -48,6 +49,7 @@ This document defines all mandatory and optional parameters in the YAML config f
 | `lora_dropout` | float | `0.1` | Dropout probability for LoRA layers to prevent overfitting on specific code snippets. |
 | `lora_bias` | str | `"none"` | Specifies if bias parameters are trained (`"none"`, `"all"`, `"lora_only"`). |
 | `lora_target_modules` | list | `["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "down_proj", "up_proj"]` | Model layers targeted for adaptation. Increasing the list improves performance but consumes more VRAM. |
+| `selected_checkpoint_strategy` | str |`"best"` |Which checkpoint to select as the pipeline output after training. `"best"` selects the checkpoint with the lowest eval loss, `"last"` selects the final checkpoint regardless of eval performance. The selected checkpoint is saved to selected_checkpoint_path. |
 | `trainer_resume_from_checkpoint` | str/null| `null` | Path to a specific checkpoint or `"last"` to continue a previous run. |
 | `trainer_clear_checkpoint_dir` | bool | `false` | If `true`, deletes the output folder before starting a new training run. |
 | `trainer_num_train_epochs` | int | `1` | Total passes through the training dataset. |
@@ -75,10 +77,11 @@ This document defines all mandatory and optional parameters in the YAML config f
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
 | `benchmark_sample_size` | int | `4` | Total number of FIM examples extracted to form the evaluation suite. |
-| `benchmark_min_fim_middle_tokens` | int | `0` | Examples with fewer tokens than this in the middle section are ignored for benchmark dataset creation. |
 | `benchmark_shuffle_buffer_size` | int | `10000000` | Size of the shuffle buffer used to shuffle benchmark examples. |
 | `benchmark_shuffle_seed` | int | `42` | Random seed ensuring that dataset shuffling is deterministic and reproducible. |
 | `benchmark_use_existing_dataset` | bool | `false` | If `true`, the pipeline reuses a previously generated benchmark dataset file instead of creating a new one. |
+| `generation_checkpoint`| str | `"pipeline"`| The checkpoint used for generation. `"pipeline"` loads the selected checkpoint from the finetune stage. Any other value is interpreted as a specific checkpoint name (e.g. `"checkpoint-250"`) and loaded directly from the finetune checkpoint directory. |
+| `generation_batch_size` | int | `10` | The number of examples processed in parallel. Increasing this value maximizes GPU hardware utilization and significantly reduces total evaluation time.|
 | `generation_max_new_tokens` | int | `128` | Upper limit on the number of tokens the models are permitted to generate for each code completion. |
 | `generation_do_sample` | bool | `false` | Enables probabilistic sampling. If `false`, the model uses greedy decoding (picking only the top token) and `generation_temperature` and `generation_top_p` are ignored.|
 | `generation_temperature` | float | `0.7` | Probability smoothing factor. Values < 1.0 make the model more confident; > 1.0 make it more random. Only active when `generation_do_sample` is `true`.|
@@ -93,7 +96,6 @@ This document defines all mandatory and optional parameters in the YAML config f
 | `sentencebleu_ngram_weight_3` | float | `0.25` | Medium phrase accuracy. Weight of 3-token chains (3-grams) in the final BLEU score calculation. |
 | `sentencebleu_ngram_weight_4` | float | `0.25` | Long phrase accuracy. Weight of 4-token blocks (4-grams) in the final BLEU score calculation. |
 | `line_match_number_of_lines` | int | `2` | Number of identical consecutive lines required between prediction and ground truth to count as a line match. |
-| `trainer_checkpoint` | str | `"last"` | Specifies the checkpoint folder to load. Use `last` for the latest or provide a specific directory name. |
 | `plot_only` | bool | `false` | If `true`, skips the heavy inference and scoring steps to only generate charts from existing results. |
 
 ## Convert Parameters (Optional)

codefinetuner-0.4.0/docs/finetuning_example/st_example/clean_st_code.py

@@ -0,0 +1,152 @@
+import argparse
+import random
+import re
+from collections import defaultdict
+from pathlib import Path
+from typing import Dict, List, Tuple
+
+
+def clean_st_content(content: str) -> str:
+    """Strips all block and inline comments, IDE pragmas, and normalizes empty lines."""
+    content = re.sub(r"\(\*.*?\*\)", "", content, flags=re.DOTALL)
+    content = re.sub(r"//.*", "", content)
+
+    lines = [line.rstrip() for line in content.splitlines()]
+    cleaned_lines = []
+    for line in lines:
+        if line:
+            cleaned_lines.append(line)
+        elif cleaned_lines and cleaned_lines[-1] != "":
+            cleaned_lines.append("")
+
+    return "\n".join(cleaned_lines).strip()
+
+
+def extract_blocks(content: str) -> List[str]:
+    """Extracts all matching structured text blocks from the content, including modern OOP structures."""
+    pattern = re.compile(
+        r"\b(?:"
+        r"FUNCTION_BLOCK\b.*?\bEND_FUNCTION_BLOCK|"
+        r"FUNCTION\b.*?\bEND_FUNCTION|"
+        r"PROGRAM\b.*?\bEND_PROGRAM|"
+        r"TYPE\b.*?\bEND_TYPE|"
+        r"CONFIGURATION\b.*?\bEND_CONFIGURATION|"
+        r"VAR_GLOBAL\b.*?\bEND_VAR|"
+        r"CLASS\b.*?\bEND_CLASS|"
+        r"INTERFACE\b.*?\bEND_INTERFACE|"
+        r"NAMESPACE\b.*?\bEND_NAMESPACE|"
+        r"METHOD\b.*?\bEND_METHOD|"
+        r"PROPERTY\b.*?\bEND_PROPERTY"
+        r")",
+        re.DOTALL | re.IGNORECASE,
+    )
+    return [match.group(0) for match in pattern.finditer(content)]
+
+
+def categorize_blocks(blocks: List[str]) -> Dict[str, List[str]]:
+    """Groups blocks into a dictionary based on their starting keyword."""
+    categorized = defaultdict(list)
+    for block in blocks:
+        first_word = block.split(maxsplit=1)[0].upper()
+        categorized[first_word].append(block)
+    return categorized
+
+
+def stratify_blocks(
+    categorized: Dict[str, List[str]], train_ratio: float, eval_ratio: float
+) -> Tuple[List[str], List[str], List[str]]:
+    """Splits categorized blocks proportionally into train, eval, and test sets."""
+    train, eval_set, test = [], [], []
+    random.seed(42)
+
+    for category, block_list in categorized.items():
+        random.shuffle(block_list)
+        total = len(block_list)
+
+        train_end = int(total * train_ratio)
+        eval_end = train_end + int(total * eval_ratio)
+
+        train.extend(block_list[:train_end])
+        eval_set.extend(block_list[train_end:eval_end])
+        test.extend(block_list[eval_end:])
+
+        print(
+            f"Category {category:15} total: {total:<4} -> "
+            f"Train: {len(block_list[:train_end]):<4} "
+            f"Eval: {len(block_list[train_end:eval_end]):<4} "
+            f"Test: {len(block_list[eval_end:]):<4}"
+        )
+
+    random.shuffle(train)
+    random.shuffle(eval_set)
+    random.shuffle(test)
+    return train, eval_set, test
+
+
+def write_splits(target_dir: Path, splits: Dict[str, List[str]]) -> None:
+    """Saves the distributed block splits into their respective files."""
+    print("-" * 50)
+    for filename, block_list in splits.items():
+        output_path = target_dir / filename
+        output_path.write_text("\n\n".join(block_list) + "\n", encoding="utf-8")
+        print(f"Wrote {len(block_list)} total stratified blocks to {output_path}")
+
+
+def split_st_by_structure(
+    input_file: Path,
+    output_dir: Path,
+    train_ratio: float,
+    eval_ratio: float,
+    test_ratio: float,
+):
+    """Coordinates file reading, processing, stratification, and export operations."""
+    target_dir = output_dir.resolve()
+    target_dir.mkdir(parents=True, exist_ok=True)
+
+    if not input_file.is_file():
+        print(f"Error: Input file not found at {input_file}")
+        return
+
+    raw_content = input_file.read_text(encoding="utf-8")
+    cleaned_content = clean_st_content(raw_content)
+
+    blocks = extract_blocks(cleaned_content)
+    if not blocks:
+        print("Error: No valid ST blocks found after cleaning.")
+        return
+
+    categorized = categorize_blocks(blocks)
+    train, eval_set, test = stratify_blocks(categorized, train_ratio, eval_ratio)
+
+    write_splits(
+        target_dir, {"train.st": train, "eval.st": eval_set, "test.st": test}
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Clean clutter and split Structured Text (ST) files structurally using stratified sampling."
+    )
+    parser.add_argument("input_file", type=Path, help="Path to the input .st source file")
+    parser.add_argument("output_dir", type=Path, help="Target directory for the clean splits")
+    parser.add_argument("--train", type=float, default=0.80, help="Train ratio (default: 0.80)")
+    parser.add_argument("--eval", type=float, default=0.10, help="Evaluation ratio (default: 0.10)")
+    parser.add_argument("--test", type=float, default=0.10, help="Test ratio (default: 0.10)")
+
+    args = parser.parse_args()
+
+    if not abs((args.train + args.eval + args.test) - 1.0) < 1e-4:
+        print(f"Error: Ratios must sum up to 1.0 (Current sum: {args.train + args.eval + args.test})")
+        return
+
+    split_st_by_structure(
+        input_file=args.input_file,
+        output_dir=args.output_dir,
+        train_ratio=args.train,
+        eval_ratio=args.eval,
+        test_ratio=args.test,
+    )
+
+
+if __name__ == "__main__":
+    main()