initial upload
This commit is contained in:
sam-paech
9
.devnotes
Normal file
9
.devnotes
Normal file
@@ -0,0 +1,9 @@
|
||||
# updating submodules:
|
||||
|
||||
git submodule update --remote --merge # add the submodule path to update just one
|
||||
git add antislop-vllm # 2. Record the new SHA(s) in the parent repo
|
||||
git commit -m "Update submodule antislop-vllm to latest"
|
||||
|
||||
# 2. Record the new SHA(s) in the parent repo
|
||||
git add slop-forensics # or `git add path/to/others`
|
||||
git commit -m "Update submodule slop-forensics to latest"
|
||||
7
.gitignore
vendored
Normal file
7
.gitignore
vendored
Normal file
@@ -0,0 +1,7 @@
|
||||
config.yaml
|
||||
__pycache__
|
||||
experiments
|
||||
*.pyc
|
||||
unsloth_compiled_cache
|
||||
results
|
||||
bct.txt
|
||||
349
README.md
Normal file
349
README.md
Normal file
@@ -0,0 +1,349 @@
|
||||
# Auto-Antislop
|
||||
|
||||
Auto-Antislop is an automated pipeline which takes a model and does the following:
|
||||
|
||||
1. Generates a baseline dataset on a set of prompts that you specify
|
||||
2. Identifies the model's slop (over-represented words, phrases & n-grams)
|
||||
3. Using [antislop-vllm](https://github.com/sam-paech/antislop-vllm) generates a preference dataset for fine-tuning
|
||||
4. Fine tunes the model on the generated preference dataset using a novel trainer (FTPO: final token preference optimisation)
|
||||
|
||||
|
||||
<strong>📚 Table of Contents</strong>
|
||||
|
||||
- [🚀 Installation](#-installation)
|
||||
- [Prerequisites](#prerequisites)
|
||||
- [Clone the Repository (with submodules)](#clone-the-repository-with-submodules)
|
||||
- [Install Dependencies](#install-dependencies)
|
||||
- [NLTK Data](#nltk-data)
|
||||
|
||||
- [⚙️ Configuration](#️-configuration)
|
||||
- [VLLM & Antislop Generation Parameters](#vllm--antislop-generation-parameters)
|
||||
- [Fine-Tuning Parameters](#fine-tuning-parameters)
|
||||
- [FTPO Parameters](#ftpo-parameters)
|
||||
|
||||
- [🛠️ Usage](#️-usage)
|
||||
- [Quickstart](#quickstart)
|
||||
- [Commandline Args](#commandline-args)
|
||||
- [Resuming Runs](#resuming-runs)
|
||||
|
||||
- [⚙️ How It Works (Pipeline Flow)](#️-how-it-works-pipeline-flow)
|
||||
|
||||
- [📖 FTPO Explained](#-ftpo-explained)
|
||||
- [How a training example is created in the Auto-Antislop pipeline](#how-a-training-example-is-created-in-the-auto-antislop-pipeline)
|
||||
- [Loss formulation](#loss-formulation)
|
||||
- [Which to choose: mse_target_tokenwise or mse_target_aggregate](#which-to-choose-mse_target_tokenwise-or-mse_target_aggregate)
|
||||
- [FTPO Tunable hyper-parameters](#ftpo-tunable-hyper-parameters)
|
||||
|
||||
- [📂 Output Structure](#-output-structure)
|
||||
|
||||
- [🧪 Post-Finetuning: Testing the Model](#-post-finetuning-testing-the-model)
|
||||
|
||||
- [🧩 Submodules](#-submodules)
|
||||
|
||||
- [💡 Notes & Troubleshooting](#-notes--troubleshooting)
|
||||
|
||||
- [📜 Citation](#-citation)
|
||||
|
||||
|
||||
|
||||
|
||||
## 🚀 Installation
|
||||
|
||||
1. **Prerequisites:**
|
||||
* Python 3.9+
|
||||
* NVIDIA GPU with CUDA installed (for vLLM and finetuning).
|
||||
* Git.
|
||||
|
||||
2. **Clone the Repository (with submodules):**
|
||||
```bash
|
||||
git clone --recurse-submodules https://github.com/your-username/auto-antislop.git
|
||||
cd auto-antislop
|
||||
```
|
||||
If you've already cloned without submodules, run:
|
||||
```bash
|
||||
git submodule update --init --recursive
|
||||
```
|
||||
|
||||
3. **Install Dependencies:**
|
||||
|
||||
*Preferably do this in a venv! Unsloth likes to install its own required dep versions, including torch.*
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
* **Note on vLLM & Torch:** `vllm` and `torch` are listed in `requirements.txt` but commented out. It's often best to install versions compatible with your specific CUDA setup.
|
||||
* If you plan to use the `--manage-vllm` feature, install `vllm` (e.g., `pip install vllm`).
|
||||
* Ensure PyTorch is installed with CUDA support (see [pytorch.org](https://pytorch.org/)).
|
||||
|
||||
4. **NLTK Data:**
|
||||
The script will attempt to download necessary NLTK resources (`punkt`, `punkt_tab`, `stopwords`) on first run. If this fails due to network issues, you might need to download them manually:
|
||||
```python
|
||||
import nltk
|
||||
nltk.download('punkt')
|
||||
nltk.download('punkt_tab') # For NLTK 3.9+
|
||||
nltk.download('stopwords')
|
||||
```
|
||||
|
||||
5. **Troubleshooting:**
|
||||
If you are getting import errors during the training step (after installing the dependencies in requirements.txt), one thing to try is building flash-attn from source. The current prebuilt wheels (flash-attn==2.8.0.post2) installed by pip are non-functional with torch 2.7. Install it from source like:
|
||||
|
||||
```
|
||||
pip uninstall -y flash-attn
|
||||
pip install -U wheel ninja packaging cmake
|
||||
MAX_JOBS=12 pip install git+https://github.com/Dao-AILab/flash-attention.git@v2.8.0.post2#egg=flash_attn --no-build-isolation
|
||||
```
|
||||
|
||||
## ⚙️ Configuration
|
||||
|
||||
The primary configuration is done through a YAML file. Refer to the examples provided in `configs/`.
|
||||
|
||||
While the pipeline is ostensibly automatic end to end, there are a lot of options you can tweak in the config file, some of which are important for a successful result. Let's walk through some of them:
|
||||
|
||||
#### VLLM & Antislop Generation Parameters
|
||||
|
||||
* **`model_id`**: The model you want to unslop. Can be a huggingface id or a local dir.
|
||||
* **`vllm management`:**: The pipeline can launch vllm automatically with the settings you provide; alternatively you can provide an openai compatible endpoint.
|
||||
* **`generation_threads`:** How many parallel threads to generate with (set according to your system specs; if in doubt try 30).
|
||||
* **`generation_max_prompts`:** The number of prompts to generate a response for. Set to 50 for a quick test. Recommmended is 1000-2000 for good coverage.
|
||||
* **`generation_hf_dataset_name`:** The huggingface dataset used to source prompts (expects sharegpt format). The kinds of prompts you use determines the slop that will be removed from the model in fine tuning.
|
||||
* **`extra_slop_phrases_to_ban`:** Set your own slop list to ban. The strings you add here will be trained out of the model.
|
||||
|
||||
#### Fine-Tuning Parameters
|
||||
|
||||
* **`finetune_use_unsloth`:** Supports unsloth or transformers/trl. Unsloth has lower vram usage but it doesn't work with all modesl in this pipeline.
|
||||
* **`finetune_mode`:** Set to "ftpo" or "dpo". FTPO is our trainer implemented specifically for the preference dataset we generate in this pipeline. It's more surgical in training out the slop words without impacting the weights otherwise, compared to DPO.
|
||||
* **`finetune_early_stopping_wins`:** This stops training when "chosen" tokens are preferred more than "rejected" tokens by this fraction. Early stopping is important to avoid overtraining. We find a good number is 0.8-0.85. Some models will degrade more easily than others, in which case you can try a lower stopping threshold.
|
||||
* **`finetune_lora_r`:** FTPO works best with a high lora rank (128-256), likely much higher than you are used to. This is because we are trying to do surgical updates of weights, and a high rank means less collateral damage on unrelated weights. Feel free to experiment; ymmv.
|
||||
* **`finetune_target_modules`:** The modules we target in fine tuning seems to be model dependent in terms of what works best. Check out the training recipes in `configs/` for working examples, or try just `["lm_head"]` for a minimally invasive fine-tune.
|
||||
* **`finetune_learning_rate`:** Set the learning rate manually.
|
||||
* **`finetune_auto_learning_rate`:** OR use an automatic learning rate that adjusts to dataset size, batch size & lora rank. Adjustable via `finetune_auto_learning_rate_adjustment_scaling`.
|
||||
* **`finetune_max_train_examples`:** The number of training examples. Suggest 8000-12000.
|
||||
|
||||
#### FTPO Parameters
|
||||
|
||||
These params control the final token preference optimisation trainer. The defaults are probably fine for a first pass. See below for a breakdown of FTPO and the configurable parameters.
|
||||
|
||||
## 🛠️ Usage
|
||||
|
||||
### Quickstart:
|
||||
|
||||
Once your .yaml is set up:
|
||||
|
||||
```
|
||||
python main.py --config myconfig.yaml
|
||||
```
|
||||
|
||||
### Commandline Args
|
||||
|
||||
**Key Command-Line Arguments (override config settings):**
|
||||
* `--config-file`: Path to the main YAML configuration file (default: `auto_antislop_config.yaml`).
|
||||
* `--resume-from-dir`: Path to an existing experiment run directory to resume.
|
||||
* `--log-level`: Set logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL).
|
||||
* `--manage-vllm [true/false]`: Override vLLM management.
|
||||
* `--vllm-port <port>`: Override vLLM port.
|
||||
* `--vllm-model-id <model_id>`: Override vLLM model.
|
||||
* `--num-iterations <N>`: Override number of iterations.
|
||||
* `--generation-max-prompts <N>`: Override max prompts for generation.
|
||||
* `--run-finetune [true/false]`: Override finetuning step.
|
||||
* `--finetune-base-model-id <model_id>`: Override base model for DPO.
|
||||
* `--finetune-mode [dpo/ftpo]`: Override finetuning mode.
|
||||
* `--finetune-cuda-visible-devices "0,1"`: Set specific GPUs for finetuning.
|
||||
|
||||
**Resuming Runs:**
|
||||
If the pipeline is interrupted, you can resume it by providing the path to the experiment directory:
|
||||
```bash
|
||||
python main.py --config-file path/to/your_config.yaml --resume-from-dir results/auto_antislop_runs/run_YYYYMMDD_HHMMSS
|
||||
```
|
||||
The script will attempt to pick up from the last successfully completed part of the iteration. Generation for an iteration is considered complete if the output JSONL file exists and contains the expected number of `prompt_id`s.
|
||||
|
||||
|
||||
## ⚙️ How It Works (Pipeline Flow)
|
||||
|
||||
1. **vLLM Server Management:** If `manage_vllm` is true, the script starts a vLLM server. If a server is already running on the configured port, or if `manage_vllm` is false, the script assumes an external vLLM server.
|
||||
2. **Antislop Iteration Loop (`num_iterations` times):**
|
||||
* **Generation (Iter 0 - Baseline):**
|
||||
* The `antislop-vllm` script generates text from a source dataset (e.g., writing prompts) *without* any ban lists active. This forms the baseline dataset.
|
||||
* **Analysis (All Iterations):**
|
||||
* The generated text is analyzed using tools from `slop-forensics` and custom analysis scripts:
|
||||
* N-gram frequencies are compared against a human writing profile.
|
||||
* Over-represented words are identified.
|
||||
* Common "slop phrases" are extracted.
|
||||
* **Ban List Update (All Iterations):**
|
||||
* Based on the analysis, n-gram and slop phrase ban lists are created or updated. User-supplied `extra_ngrams_to_ban`, `extra_slop_phrases_to_ban`, and `extra_regex_patterns` from the config are merged.
|
||||
* **Generation (Iter 1+ - Anti-Slop):**
|
||||
* `antislop-vllm` generates text again on the same prompts, but this time it uses the accumulated ban lists (n-grams, slop phrases, regex) to avoid slop.
|
||||
* While generating, preference pairs are produced each time a ban occurs, containing a *rejected* token and a number of coherent alternative *chosen* tokens, as well as the preceding context.
|
||||
3. **FTPO Dataset Creation:**
|
||||
* Pairs are created:
|
||||
* `prompt`: The original input prompt + generated context so far.
|
||||
* `rejected`: The first token in a banned sequence, at the time a ban occurred during antislop generation.
|
||||
* `chosen`: A number of coherent alternative tokens at that position, constrained by min_p.
|
||||
4. **FTPO Finetuning:**
|
||||
* If `finetune_enabled` is true, the script runs FTPO finetuning using the preference dataset.
|
||||
* Supports LoRA and optional 4-bit quantization
|
||||
* Supports unsloth or transformers/trl training paths (though some models may not work with both)
|
||||
* Saves the LoRA adapters and optionally a merged 16-bit model.
|
||||
|
||||
### 📖 FTPO Explained
|
||||
|
||||
FTPO (Final-Token Preference Optimisation) is a surgical preference optimisation training algorithm that constrains gradient updates to just a rejected/chosen *continuation token*, and avoids training on the preceding context. The intent is to push probability mass **away from the first token of a banned phrase (the *rejected* token)** and **toward one or more viable alternatives (the *chosen* tokens)** while leaving the rest of the model distribution largely intact.
|
||||
|
||||
The loss function operates entirely in logit-space (in contrast to most preference optimisation trainers that use softmax), resulting in minimalist targeted weight updates.
|
||||
|
||||
---
|
||||
|
||||
#### How a training example is created in the Auto-Antislop pipeline
|
||||
|
||||
1. **Generation runs with antislop active.**
|
||||
In the auto-antislop pipeline, the FTPO dataset is generated in iterations > 0, when antislop is actively banning slop that it surfaced during the first iteration. Whenever the sampler encounters a banned n-gram / phrase / regex, it halts and constructs a training example before resuming inference with a non-banned continuation.
|
||||
|
||||
2. **Rejected token.**
|
||||
The would-have-been next token (the first token of the banned phrase) is stored as the `rejected` continuation token.
|
||||
|
||||
3. **Chosen tokens.**
|
||||
The sampler then draws further candidates for that same position, applying a *min-p* filter ¹ to keep only continuations whose tail probability mass is above a given threshold, to ensure they are coherent. These candidates are further filtered per the banned phrases list. The remaining tokens are stored as the `chosen` continuation tokens in the sample.
|
||||
|
||||
* If no alternative passes the filter the event is discarded and no FTPO sample is written.
|
||||
|
||||
4. **Context.**
|
||||
The full prompt (and any chat template markers) up to but **not including** the banned token is stored.
|
||||
This means the model receives identical context for both the rejected and chosen tokens.
|
||||
|
||||
Result: a single JSONL line contains the shared context plus one rejected token and a small set of chosen tokens — exactly the information FTPO needs.
|
||||
|
||||
---
|
||||
|
||||
#### Loss formulation
|
||||
|
||||
**Preference term.**
|
||||
For each example the trainer computes
|
||||
`Δ = logit(chosen) − logit(rejected)`
|
||||
for all chosen+rejected pairs. In a given sample there is 1 rejected and typically 4+ chosen tokens.
|
||||
The loss is a function of the amount that all the chosen logits are beating the rejected, averaged across chosen tokens for that sample.
|
||||
This encourages the model to rank *every* chosen token above the rejected one.
|
||||
Once a chosen logit is beating rejected by a given margin, it no longer contributes to the loss. This helps to avoid unnecessarily moving the weights when chosen is already winning.
|
||||
|
||||
**Three-term MSE regulariser.**
|
||||
A key part of the loss function is the MSE loss which is split into three terms:
|
||||
|
||||
* lambda_mse \* mse_non_target +
|
||||
* lambda_mse_target \* mse_target
|
||||
|
||||
Where "target" refers to the chosen & rejected logits for a given training example, and non-target refers to the remaining vocab.
|
||||
|
||||
We use MSE loss in logit space rather than KL loss, because applying softmax as part of the loss function (like KL does) creates learning pressure on the whole vocab, when we are instead trying to do minimal targeted gradient updates.
|
||||
|
||||
* **mse_target**: This term applies tokenwise loss pressure on just the target tokens (rejected & chosen) to keep them close to the original weights. We apply this separately so that we can apply a weaker MSE loss to the target tokens, allowing them to move more freely than the remainder of the vocab. This is because the target tokens need to move significantly relative to one another, since the "rejected" token is typically highest prob by a large margin (that's why it's slop!). Generally you should enable *either* mse_target_tokenwise or mse_target_aggregate.
|
||||
* **lambda_mse_target**: The scaling strength applied to the mse_target_tokenwise term. Set to 0 in the config to disable this loss term.
|
||||
* **mse_non_target**: This loss term represents tokenwise loss for the remaining vocab (other than the target chosen + rejected tokens).
|
||||
* **lambda_mse**: The scaling strength applied to the mse_non_target term. Set to 0 in the config to disable this loss term.
|
||||
|
||||
**Tau parameters**
|
||||
There is also a `tau` parameter for each of the two *target* loss terms, which acts as penalty free range (in logits) within which logits can move relative to baseline without incurring loss. Setting tau to > 0 can be helpful when using mse_target_tokenwise, to allow the model to learn more easily and reach higher preference accuracies when training. A reasonable range is 0-1.5. Higher values may lead to degradation with some models.
|
||||
|
||||
**Which to choose: mse_target_tokenwise or mse_target_aggregate**
|
||||
You can use both loss terms together, but mse_target_tokenwise will tend dominate as it's a applied per-token, where the aggregate term allows logits to move significantly before it kicks in.
|
||||
Some models can tolerate the more permissive mse_target_aggregate term without degrading; other models are more sensitive and need mse_target_tokenwise to keep logits closer to baseline.
|
||||
Check in the `configs/` dir for training recipes for specific models that have worked.
|
||||
|
||||
#### FTPO Tunable hyper-parameters
|
||||
|
||||
** These are settable in the config file: **
|
||||
```
|
||||
# ── FTPO-specific hyper-parameters ─────────────────────────────────────────
|
||||
# Leave any of these out (or set to null) to fall back to FTPOTrainer defaults.
|
||||
ftpo_beta: 0.1 # Global scale on pref loss (higher = steeper sigmoid).
|
||||
|
||||
# MSE loss term 1: light mse loss applied tokenwise on target tokens
|
||||
ftpo_lambda_mse_target_tokenwise: 0.05 # Strength of MSE loss tether on the individual logits in the
|
||||
# chosen+rejected set vs reference.
|
||||
ftpo_tau_mse_target_tokenwise: 0.5 # Grace bandwidth (logits) before the above MSE loss kicks in.
|
||||
|
||||
# MSE loss term 2: stronger mse term applied to remaining (non-target) vocab
|
||||
ftpo_lambda_mse: 0.4
|
||||
|
||||
ftpo_clip_epsilon_logits: 2 # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off"
|
||||
```
|
||||
|
||||
## 📂 Output Structure
|
||||
|
||||
Outputs are saved in `experiment_base_dir` (e.g., `results/auto_antislop_runs/`), under a timestamped directory for each run (e.g., `run_YYYYMMDD_HHMMSS/`):
|
||||
|
||||
* `run_config_YYYYMMDD_HHMMSS.yaml`: The exact configuration used for this run.
|
||||
* `iter_N_creative_writing_generations.jsonl`: Raw generated text for iteration `N`.
|
||||
* `iter_N_ftpo_pairs.jsonl`: (If FTPO mode is active) Fine-grained preference pairs for iteration `N`.
|
||||
* `iter_N_analysis_results/`: Directory containing:
|
||||
* `bigrams__dictionary_sorted.csv`, `trigrams__non_dictionary_sorted.csv`, etc.: N-gram analysis results.
|
||||
* `overrepresented_words.csv`: Analysis of over-represented words.
|
||||
* `slop_list_phrases.jsonl` (inside `phrase_tmp/`): Candidate slop phrases from `slop-forensics`.
|
||||
* `banned_ngrams_used.json`, `banned_slop_phrases_used.json`: Copies of ban lists *used* for this iteration's generation (for iter > 0).
|
||||
* `banned_ngrams_new_this_iter.json`, `banned_slop_phrases_new_this_iter.json`: Ban list entries *added* after this iteration's analysis.
|
||||
* `orchestration.log`: Log specific to analysis and ban list updates for this iteration.
|
||||
* `banned_ngrams.json`: Aggregated list of banned n-grams across iterations.
|
||||
* `banned_slop_phrases.json` (or custom name): Aggregated list of banned slop phrases.
|
||||
* `user_defined_regex_blocklist.json`: Copy of user-defined regex patterns.
|
||||
* `dpo_pairs_dataset.jsonl`: The final preference dataset for DPO/FTPO.
|
||||
* `final_iteration_statistics.csv`: Summary metrics for each iteration.
|
||||
* `finetuned_model_SUFFIX/`: (If finetuning is run)
|
||||
* `lora_adapters/`: Saved LoRA adapter weights and tokenizer config.
|
||||
* `merged_16bit/`: (If `finetune_save_merged_16bit: true`) Full model with LoRA weights merged, in 16-bit precision.
|
||||
* `gguf_q8_0.gguf`: (If `finetune_save_gguf_q8_0: true`) GGUF quantized model.
|
||||
* `logprob_gap_analysis/`: (If FTPO mode) JSONL files with pre/post training logprob gap statistics.
|
||||
|
||||
## 🧪 Post-Finetuning: Testing the Model
|
||||
|
||||
A simple script `test_inference.py` is provided to load the latest finetuned model and run a test generation.
|
||||
|
||||
```bash
|
||||
python test_inference.py
|
||||
```
|
||||
This script automatically searches for the most recent `merged_16bit` model in the standard output directories. You can modify the prompt within the script.
|
||||
|
||||
##🧩 Submodules
|
||||
|
||||
* **`antislop-vllm`**: (Path: `antislop-vllm/`)
|
||||
* Handles the core text generation using vLLM.
|
||||
* Crucially, it implements the logic for dynamic banning of n-grams, phrases, and regex patterns during the generation process.
|
||||
* **`slop-forensics`**: (Path: `slop-forensics/`)
|
||||
* Provides tools and algorithms for analyzing text to identify various types of "slop," including over-represented n-grams and common undesirable phrases.
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
---
|
||||
|
||||
¹ *min-p sampler: see* *Nguyen et al., “Turning Up the Heat: Min-p Sampling for Creative and Coherent LLM Outputs” (arXiv:2407.01082).*
|
||||
|
||||
|
||||
|
||||
## 💡 Notes & Troubleshooting
|
||||
|
||||
* **GPU Memory:** Running vLLM and finetuning (especially with larger models) requires significant GPU VRAM. Adjust `vllm_gpu_memory_utilization` and finetuning batch sizes/quantization accordingly. If running both on the same GPU, the script attempts to stop vLLM before finetuning to free up VRAM.
|
||||
* **Submodule Issues:** If you encounter errors related to `antislop-vllm` or `slop-forensics`, ensure the submodules are correctly initialized (`git submodule update --init --recursive`).
|
||||
* **NLTK Data:** If `ensure_core_nltk_resources()` fails, download the resources manually as described in Installation.
|
||||
* **Unsloth Cache:** Unsloth might create a `unsloth_compiled_cache` directory. This is ignored by git.
|
||||
* **Gemma-3 Checkpoints:** The `utils/model_helpers.py` contains a `fix_gemma3_checkpoint` function to handle potential inconsistencies in Gemma-3 model key naming, and `detie_lm_head` to ensure proper saving of merged models.
|
||||
* **FTPO Mode:** The "ftpo" (Final Token Preference Optimization) mode uses `FTPOTrainer` which focuses on the preference for a single next token, given a context. This is useful for correcting specific token choices rather than entire continuations.
|
||||
|
||||
|
||||
## Citation
|
||||
|
||||
If you use Auto-Antislop or the concepts from the original `antislop-sampler` in your research, please consider citing:
|
||||
|
||||
```bibtex
|
||||
@misc{paech2024antislop,
|
||||
title={AntiSlop Sampler},
|
||||
author={Samuel J. Paech},
|
||||
year={2024},
|
||||
howpublished={\url{https://github.com/sam-paech/antislop-sampler}}
|
||||
}
|
||||
|
||||
@misc{paech2024antislop,
|
||||
title={Auto-Antislop},
|
||||
author={Samuel J. Paech},
|
||||
year={2025},
|
||||
howpublished={\url{https://github.com/sam-paech/auto-antislop}}
|
||||
}
|
||||
```
|
||||
And/or link to this repository: `https://github.com/sam-paech/auto-antislop`
|
||||
279
auto_antislop_config.yaml
Normal file
279
auto_antislop_config.yaml
Normal file
@@ -0,0 +1,279 @@
|
||||
################################################################################
|
||||
# MAIN AUTO-ANTISLOP CONFIGURATION
|
||||
################################################################################
|
||||
|
||||
################################################################################
|
||||
# RUN SETUP
|
||||
################################################################################
|
||||
experiment_base_dir: "results/auto_antislop_runs" # Base for timestamped run directories
|
||||
human_profile_path: "data/human_writing_profile.json"
|
||||
log_level: "INFO"
|
||||
# Iteration 0: Generates the baseline dataset & computes slop strings/ngrams to ban
|
||||
# Iteration 1: Generates a dataset using antislop, banning those strings & ngrams. Recomputes the slop strings/ngrams at the end & adds any new slop to the ban lists
|
||||
# Iteration 2+: Extra iterations catch slop that emerges after the initial set is banned
|
||||
num_iterations: 2 # Minimum 2 iterations (this is enough to catch most slop)
|
||||
model_id: "unsloth/gemma-3-4b-it" # Global model id for the pipeline. Can be overridden on individual steps.
|
||||
|
||||
################################################################################
|
||||
# VLLM SERVER MANAGEMENT (Conditional: if --manage-vllm is True)
|
||||
################################################################################
|
||||
manage_vllm: true
|
||||
vllm_model_id: null # Model served by vLLM (if unset, will use model_id)
|
||||
vllm_port: 8000
|
||||
vllm_hf_token: null # Optional: Your Hugging Face token if model is gated
|
||||
vllm_cuda_visible_devices: "0" # set to e.g. "0,1,2,3" for multiple gpus
|
||||
vllm_gpu_memory_utilization: 0.85 # leave some room for the refusal classifier if you are using it (about 3gb)
|
||||
vllm_max_model_len: 2500
|
||||
vllm_dtype: "bfloat16"
|
||||
# Additional raw CLI arguments for vLLM server, e.g., ["--tensor-parallel-size", "4"] for multiple gpus
|
||||
vllm_extra_args: [] # each param as a separate string, e.g. ["--quantization", "bitsandbytes"]
|
||||
vllm_env: # env vars for the vLLM process
|
||||
# VLLM_USE_V1: "1" # may be needed for amd gpus
|
||||
|
||||
|
||||
################################################################################
|
||||
# GENERATION PARAMETERS (using antislop-vllm)
|
||||
################################################################################
|
||||
generation_step_enabled: true
|
||||
|
||||
# --- API & Model Configuration ---
|
||||
# If you set manage_vllm=true, leave the base url unset
|
||||
#generation_api_base_url: "http://localhost:8000/v1"
|
||||
#generation_api_base_url: "https://apjmbtwbrb8t61-8888.proxy.runpod.net/v1"
|
||||
generation_model_id: null # Model id for generation requests (if unset, uses model_id)
|
||||
generation_api_key: "xxx" # API key for the vLLM server
|
||||
|
||||
# --- Core Generation Settings ---
|
||||
generation_max_new_tokens: 1000
|
||||
generation_threads: 50 # Number of parallel threads for API queries in antislop-vllm. Note: vllm can become very inefficient if you go over some concurrency threshold (depending on vram)
|
||||
generation_max_prompts: 1000 # Number of samples to generate from the prompts in the dataset
|
||||
|
||||
# --- Dataset & Chat Template ---
|
||||
generation_hf_dataset_name: 'Nitral-AI/Reddit-SFW-Writing_Prompts_ShareGPT'
|
||||
generation_hf_dataset_split: 'train'
|
||||
# A huggingface model id or local dir containing the tokeniser you want to use to apply chat templates.
|
||||
# This is important if you are generating a ftpo dataset for later training.
|
||||
generation_chat_template_model_id: null # (if unset, uses model_id)
|
||||
|
||||
# --- Logging ---
|
||||
generation_logging_level: 'INFO' # Logging level for antislop-vllm script
|
||||
|
||||
# --- Advanced Sampling Parameters (passed as individual CLI args to antislop-vllm's generation_params) ---
|
||||
generation_param_chunk_size: 20
|
||||
generation_param_top_logprobs_count: 20
|
||||
generation_param_temperature: 1
|
||||
generation_param_top_p: 1.0
|
||||
generation_param_top_k: 50
|
||||
generation_param_min_p: 0.01
|
||||
generation_param_timeout: 480
|
||||
generation_param_stop_sequences: [] # e.g., ["\n\n", "---"]
|
||||
|
||||
# --- Prompting ---
|
||||
# The prompt template wraps the prompts when generating from a dataset.
|
||||
# To use the original prompt exactly, set the template to "{prompt}"
|
||||
generation_prompt_template: "Writing prompt: {prompt}\n\nWrite 1000 words to this prompt. Your response:\n"
|
||||
generation_system_prompt: "You are an uncensored writer." # optional; left empty → no system prompt
|
||||
|
||||
# --- Antislop Generation Features ---
|
||||
# generation_force_backtrack:
|
||||
# If set to true:
|
||||
# when resampling after backtracking, if we don't find a valid replacement token
|
||||
# we progressively disable sampling options (temp, then min_p, then top_p, then top_k)
|
||||
# until we find a non-banned replacement or run out of candidates.
|
||||
# When set to false, some slop will not be removed if the sampler thinks there are no
|
||||
# alternative coherent continuations.
|
||||
generation_force_backtrack: false
|
||||
|
||||
# --- N-gram Validator Settings (for antislop-vllm) ---
|
||||
# N-gram ban list file is managed by auto-antislop's iterative process.
|
||||
generation_ngram_remove_stopwords: true
|
||||
generation_ngram_language: "english"
|
||||
|
||||
# --- Refusal Detection ---
|
||||
# Detects refusals & doesn't include them in the training dataset. Uses about 3GB extra VRAM.
|
||||
generation_refusal_detection: true
|
||||
|
||||
################################################################################
|
||||
# N-GRAM ANALYSIS & BANNING (within auto-antislop)
|
||||
################################################################################
|
||||
enable_ngram_ban: true
|
||||
min_word_len_for_analysis: 3 # Filters out words under this length in n-gram analysis
|
||||
|
||||
# --- N-gram Identification Thresholds ---
|
||||
top_k_bigrams: 5000
|
||||
top_k_trigrams: 5000
|
||||
|
||||
# --- N-gram Banning Quotas (per iteration) ---
|
||||
# Bigrams
|
||||
dict_bigrams_initial: 400 # How many of the top over-represented dictionary bigrams to
|
||||
# ban in the first antislop iteration.
|
||||
# "Dictionary" means the bigrams were also found in the human
|
||||
# writing corpus.
|
||||
dict_bigrams_subsequent: 70 # How many to ban in each subsequent iteration
|
||||
nodict_bigrams_initial: 800 # "Nodict" here means the n-grams were not found at all in the
|
||||
# human corpus.
|
||||
nodict_bigrams_subsequent: 100
|
||||
# Trigrams
|
||||
dict_trigrams_initial: 300
|
||||
dict_trigrams_subsequent: 50
|
||||
nodict_trigrams_initial: 800
|
||||
nodict_trigrams_subsequent: 100
|
||||
|
||||
# --- User-Defined N-gram Bans ---
|
||||
# User-supplied extra n-grams to always ban (processed by auto-antislop)
|
||||
extra_ngrams_to_ban: [
|
||||
# "voice barely whisper",
|
||||
]
|
||||
|
||||
################################################################################
|
||||
# OVER-REPRESENTED WORD ANALYSIS & BANNING
|
||||
################################################################################
|
||||
compute_overrep_words: true
|
||||
top_k_words_for_overrep_analysis: 200000
|
||||
|
||||
# --- Quotas for Adding Over-represented Words to Slop Phrase Ban List ---
|
||||
dict_overrep_initial: 800 # How many of the top over-represented dictionary words to
|
||||
# ban in the first antislop iteration.
|
||||
# "Dictionary" means the words were also found in the human
|
||||
# writing corpus.
|
||||
dict_overrep_subsequent: 200 # How many to ban in each subsequent iteration
|
||||
nodict_overrep_initial: 80 # "Nodict" here means the n-grams were not found at all in the
|
||||
# human corpus.
|
||||
nodict_overrep_subsequent: 20
|
||||
|
||||
################################################################################
|
||||
# SLOP PHRASE BANNING
|
||||
################################################################################
|
||||
|
||||
# Slop phrases are over-represented whole phrases extracted from the generated texts.
|
||||
enable_slop_phrase_ban: true
|
||||
min_phrase_freq_to_keep: 2 # Min frequency for a new phrase from slop-forensics to be considered
|
||||
top_n_initial_slop_ban: 600 # New slop phrases from slop-forensics to ban in iter 0
|
||||
top_n_subsequent_slop_ban: 100 # New slop phrases from slop-forensics to ban in later iters
|
||||
|
||||
# --- User-Defined Slop Phrase Bans ---
|
||||
# User supplied list of strings to always ban
|
||||
# - case insensitive
|
||||
# To trigger a ban, the sequence must not have a word-like character
|
||||
# (not punctuation or whitespace) directly on either side. That is to say, we
|
||||
# are not banning disallowed sequences that occur as substrings in longer
|
||||
# words. The exception is if the banned string is already bookended by
|
||||
# a non-word character.
|
||||
#
|
||||
# Examples:
|
||||
# banned string "cat"
|
||||
# - won't trigger a ban for "cation"
|
||||
# - will trigger a ban on "cat[morecat]"
|
||||
# banned string "cat["
|
||||
# - *will* trigger a ban on "cat[morecat]", because the banned string
|
||||
# ends with a non-word character.
|
||||
extra_slop_phrases_to_ban: [
|
||||
# "testament to",
|
||||
#"…", "*", " –", "–", "#",
|
||||
]
|
||||
|
||||
# --- Whitelisted Strings ---
|
||||
# These will be excluded from the list of slop strings that the pipeline finds.
|
||||
# Note: special tokens in the tokenizer and parts of the chat template are
|
||||
# automatically whitelisted.
|
||||
whitelist_strings: [
|
||||
# "think", "thinking"
|
||||
]
|
||||
|
||||
################################################################################
|
||||
# REGEX BANNING
|
||||
################################################################################
|
||||
# User-supplied regex patterns to ban
|
||||
# Note: unoptimised regex patterns can slow down antislop generation, as they will be called often on large texts.
|
||||
extra_regex_patterns: [
|
||||
# These ones ban "it's not x, it's y" type patterns:
|
||||
|
||||
#"\\b(?:\\w+n(?:['’]t)|not\\s+(?:just|only|merely|because))\\s+(?:(?![.;:?!…]).){1,100}?[.;:?!…]\\s*(?:it|they|you)(?:['’](?:s|re|m))?\\b(?!\\s+(?:was|were|is|are|wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t)\\b)(?:\\s*[*…]?\\s*)?(?!when\\b|then\\b|but\\b|and\\b|yet\\b)(?!right\\b)(?!normal\\b)(?!true\\b)(?!sure\\b)(?!only\\b)(?!still\\b)(?!rarely\\b)(?!already\\b)(?!wrong\\b)(?!want\\b)(?!just\\b)(?!couldn\\b)(?!could\\b)(?!saw\\b)(?!started\\b)(?!remember\\b)(?!struggled\\b)(?!watched\\b)(?!goal\\b)(?!took\\b)(?!kept\\b)(?!reminded\\b)(?!time\\b)(?!have\\b)(?!acted\\b)(?!smiled\\b)(?!think\\b)(?!give\\b)(?!grab\\b)(?!gave\\b)(?!turn\\b)(?!justify\\b)(?!\\w+ly\\b)(?=[a-z]{4,}\\b)[a-z]+\\w*",
|
||||
|
||||
#"\\b(?:\\w+n(?:['’]t)|not)\\s+(?:just|only|merely)?\\s*(?:(?![-–—]|[.?!…]).){1,80}?[-–—]{1,2}\\s*\\w+(?:['’]\\w+)?\\s+",
|
||||
|
||||
#"\\b(?:wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t|not)\\s+(?!\\b(?:minute|minutes|hour|hours|day|days|year|years|second|seconds)\\b)(?!with\\b)(?!even\\b)(?:(?![.;:?!…]).){2,120}?[.;:?!…]\\s*(?:it|they|you|that)(?:\\s+(?:was|were|is|are)\\b(?:\\s+[*_~]?\\w+[*_~]?)?|(?:['’](?:s|re|m))\\b(?:\\s+[*_~]?\\w+[*_~]?)?)",
|
||||
|
||||
#"\\bno\\s+longer\\s+(?:just|only|merely)?\\s+[^.;:?!…]{1,120}[.;:?!…]\\s*(?:it|they|you)\\s+(?:is|are|was|were)\\b(?:\\s+[*_~]?\\w+[*_~]?)?",
|
||||
|
||||
#"\\b(?:wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t|not)\\s+(?:just|only|merely)?\\s*(?:(?!\\bbut\\b|[.?!…]).){1,80}?[,;:\\-–—]\\s*but\\s+(?!I\\b)(?:also\\s+)?"
|
||||
|
||||
]
|
||||
|
||||
################################################################################
|
||||
# FINETUNING
|
||||
################################################################################
|
||||
finetune_enabled: true
|
||||
|
||||
# --- General Finetuning Setup ---
|
||||
finetune_use_unsloth: false
|
||||
finetune_mode: "ftpo" # dpo / ftpo (final token preference optimisation)
|
||||
finetune_ftpo_dataset: "" # you can specify an existing ftpo dataset, or leave unset to let the
|
||||
# pipeline use the one produced in the generation step
|
||||
finetune_base_model_id: null # Base model for DPO (if unset, uses model_id)
|
||||
finetune_max_seq_length: 3500 # this may truncate some outputs
|
||||
finetune_load_in_4bit: false # qlora
|
||||
|
||||
# --- Early Stopping ---
|
||||
finetune_early_stopping_wins: 0.85 # Early stopping threshold for fraction of *chosen* completions that are selected over *rejected*.
|
||||
# More than 0.85 may be overtrained. Set to > 1.0 to disable early stopping.
|
||||
finetune_early_stopping_loss: null # Loss threshold for early stopping. Set to null to disable.
|
||||
|
||||
# --- LoRA Configuration ---
|
||||
finetune_lora_r: 128 # the ftpo trainer works best with a high lora rank
|
||||
finetune_lora_alpha: 128
|
||||
finetune_lora_dropout: 0.05
|
||||
finetune_weight_decay: 0.01
|
||||
finetune_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "lm_head"]
|
||||
|
||||
# --- Layer Freezing ---
|
||||
finetune_freeze_early_layers: true
|
||||
finetune_n_layers_unfrozen: 10
|
||||
|
||||
# --- Training Process ---
|
||||
finetune_gradient_checkpointing: "unsloth"
|
||||
finetune_chat_template: "" # e.g. "gemma-3" -- get the chat template from unsloth's helper if required, otherwise leave the string blank to use the tokeniser's chat template
|
||||
finetune_batch_size: 1
|
||||
finetune_gradient_accumulation_steps: 16
|
||||
finetune_warmup_ratio: 0.1
|
||||
finetune_num_epochs: 1
|
||||
|
||||
# --- Learning Rate ---
|
||||
finetune_learning_rate: 0.000001
|
||||
finetune_auto_learning_rate: true # true: automatically determine learning rate based on dataset size, effective batch size & lora rank
|
||||
finetune_auto_learning_rate_adjustment_scaling: 0.15 # scale the auto-lr by this factor
|
||||
|
||||
# --- DPO/FTPO Specific ---
|
||||
finetune_beta: 0.1 # DPO beta
|
||||
|
||||
# --- Output & Saving ---
|
||||
finetune_output_dir_suffix: "_ftpo_exp01" # Appended to experiment run dir
|
||||
finetune_save_merged_16bit: true
|
||||
finetune_save_gguf_q8_0: false
|
||||
|
||||
# --- Dataset Handling for Finetuning ---
|
||||
finetune_max_train_examples: 1000 # adjust as needed
|
||||
finetune_shuffle_seed: 666
|
||||
|
||||
# --- FTPO Sample Regularization ---
|
||||
# 0 = off; 0.9 strongly downsamples overrepresented rule violations
|
||||
# (this is useful because the raw generated dataset is typically very skewed)
|
||||
ftpo_sample_rejected_regularisation_strength: 0.8
|
||||
ftpo_sample_chosen_regularisation_strength: 0.2
|
||||
ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list
|
||||
|
||||
# ── FTPO-specific hyper-parameters ─────────────────────────────────────────
|
||||
# Leave any of these out (or set to null) to fall back to FTPOTrainer defaults.
|
||||
|
||||
# Loss terms are computed separately for the target (chosen + rejected) tokens vs the remainder of the vocab.
|
||||
# This is because we want to allow more freedom of movement for the target tokens.
|
||||
|
||||
# MSE loss term 1: light mse loss applied tokenwise on target tokens
|
||||
ftpo_lambda_mse_target: 0.05 # Strength of MSE loss tether on the individual logits in the
|
||||
# chosen+rejected set vs reference.
|
||||
ftpo_tau_mse_target: 0.5 # Grace bandwidth (logits) before the above MSE loss kicks in.
|
||||
|
||||
# MSE loss term 2: stronger mse term applied to remaining (non-target) vocab
|
||||
ftpo_lambda_mse: 0.4
|
||||
|
||||
ftpo_clip_epsilon_logits: 2 # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off"
|
||||
286
configs/GLM-4-32B-0414.yaml
Normal file
286
configs/GLM-4-32B-0414.yaml
Normal file
@@ -0,0 +1,286 @@
|
||||
################################################################################
|
||||
# MAIN AUTO-ANTISLOP CONFIGURATION
|
||||
################################################################################
|
||||
|
||||
################################################################################
|
||||
# RUN SETUP
|
||||
################################################################################
|
||||
experiment_base_dir: "results/auto_antislop_runs" # Base for timestamped run directories
|
||||
human_profile_path: "data/human_writing_profile.json"
|
||||
log_level: "INFO"
|
||||
# Iteration 0: Generates the baseline dataset & computes slop strings/ngrams to ban
|
||||
# Iteration 1: Generates a dataset using antislop, banning those strings & ngrams. Recomputes the slop strings/ngrams at the end & adds any new slop to the ban lists
|
||||
# Iteration 2+: Extra iterations catch slop that emerges after the initial set is banned
|
||||
num_iterations: 2 # Minimum 2 iterations (this is enough to catch most slop)
|
||||
model_id: "THUDM/GLM-4-32B-0414" # Global model id for the pipeline. Can be overridden on individual steps.
|
||||
|
||||
# !! NEEDED TO SET ATTN TO EAGER
|
||||
# add to finetuning.py after model load:
|
||||
# os.environ["FLASH_ATTENTION_FORCE_EAGER"] = "1"
|
||||
# model.config._attn_implementation = "eager"
|
||||
|
||||
|
||||
################################################################################
|
||||
# VLLM SERVER MANAGEMENT (Conditional: if --manage-vllm is True)
|
||||
################################################################################
|
||||
manage_vllm: true
|
||||
vllm_model_id: null # Model served by vLLM (if unset, will use model_id)
|
||||
vllm_port: 8000
|
||||
vllm_hf_token: null # Optional: Your Hugging Face token if model is gated
|
||||
vllm_cuda_visible_devices: "0" # set to e.g. "0,1,2,3" for multiple gpus
|
||||
vllm_gpu_memory_utilization: 0.85 # leave some room for the refusal classifier if you are using it (about 3gb)
|
||||
vllm_max_model_len: 4500
|
||||
vllm_dtype: "bfloat16"
|
||||
# Additional raw CLI arguments for vLLM server, e.g., ["--tensor-parallel-size", "4"] for multiple gpus
|
||||
vllm_extra_args: ["--quantization", "bitsandbytes"]
|
||||
vllm_env: # env vars for the vLLM process
|
||||
# VLLM_USE_V1: "1" # may be needed for amd gpus
|
||||
|
||||
|
||||
################################################################################
|
||||
# GENERATION PARAMETERS (using antislop-vllm)
|
||||
################################################################################
|
||||
generation_step_enabled: true
|
||||
|
||||
# --- API & Model Configuration ---
|
||||
# If you set manage_vllm=true, leave the base url unset
|
||||
#generation_api_base_url: "http://localhost:8000/v1"
|
||||
#generation_api_base_url: "https://apjmbtwbrb8t61-8888.proxy.runpod.net/v1"
|
||||
generation_model_id: null # Model id for generation requests (if unset, uses model_id)
|
||||
generation_api_key: "xxx" # API key for the vLLM server
|
||||
|
||||
# --- Core Generation Settings ---
|
||||
generation_max_new_tokens: 1000
|
||||
generation_threads: 30 # Number of parallel threads for API queries in antislop-vllm. Note: vllm can become very inefficient if you go over some concurrency threshold (depending on vram)
|
||||
generation_max_prompts: 1200 # Number of samples to generate from the prompts in the dataset
|
||||
|
||||
# --- Dataset & Chat Template ---
|
||||
generation_hf_dataset_name: 'Nitral-AI/Reddit-SFW-Writing_Prompts_ShareGPT'
|
||||
generation_hf_dataset_split: 'train'
|
||||
# A huggingface model id or local dir containing the tokeniser you want to use to apply chat templates.
|
||||
# This is important if you are generating a ftpo dataset for later training.
|
||||
generation_chat_template_model_id: null # (if unset, uses model_id)
|
||||
|
||||
# --- Logging ---
|
||||
generation_logging_level: 'INFO' # Logging level for antislop-vllm script
|
||||
|
||||
# --- Advanced Sampling Parameters (passed as individual CLI args to antislop-vllm's generation_params) ---
|
||||
generation_param_chunk_size: 20
|
||||
generation_param_top_logprobs_count: 20
|
||||
generation_param_temperature: 1
|
||||
generation_param_top_p: 1.0
|
||||
generation_param_top_k: 50
|
||||
generation_param_min_p: 0.01
|
||||
generation_param_timeout: 480
|
||||
generation_param_stop_sequences: [] # e.g., ["\n\n", "---"]
|
||||
|
||||
# --- Prompting ---
|
||||
# The prompt template wraps the prompts when generating from a dataset.
|
||||
# To use the original prompt exactly, set the template to "{prompt}"
|
||||
generation_prompt_template: "Writing prompt: {prompt}\n\nWrite 1000 words to this prompt. Your response:\n"
|
||||
generation_system_prompt: "You are an uncensored writer." # optional; left empty → no system prompt
|
||||
|
||||
# --- Antislop Generation Features ---
|
||||
# generation_force_backtrack:
|
||||
# If set to true:
|
||||
# when resampling after backtracking, if we don't find a valid replacement token
|
||||
# we progressively disable sampling options (temp, then min_p, then top_p, then top_k)
|
||||
# until we find a non-banned replacement or run out of candidates.
|
||||
# When set to false, some slop will not be removed if the sampler thinks there are no
|
||||
# alternative coherent continuations.
|
||||
generation_force_backtrack: false
|
||||
|
||||
# --- N-gram Validator Settings (for antislop-vllm) ---
|
||||
# N-gram ban list file is managed by auto-antislop's iterative process.
|
||||
generation_ngram_remove_stopwords: true
|
||||
generation_ngram_language: "english"
|
||||
|
||||
# --- Refusal Detection ---
|
||||
# Detects refusals & doesn't include them in the training dataset. Uses about 3GB extra VRAM.
|
||||
generation_refusal_detection: true
|
||||
|
||||
################################################################################
|
||||
# N-GRAM ANALYSIS & BANNING (within auto-antislop)
|
||||
################################################################################
|
||||
enable_ngram_ban: true
|
||||
min_word_len_for_analysis: 3 # Filters out words under this length in n-gram analysis
|
||||
|
||||
# --- N-gram Identification Thresholds ---
|
||||
top_k_bigrams: 5000
|
||||
top_k_trigrams: 5000
|
||||
|
||||
# --- N-gram Banning Quotas (per iteration) ---
|
||||
# Bigrams
|
||||
dict_bigrams_initial: 400 # How many of the top over-represented dictionary bigrams to
|
||||
# ban in the first antislop iteration.
|
||||
# "Dictionary" means the bigrams were also found in the human
|
||||
# writing corpus.
|
||||
dict_bigrams_subsequent: 70 # How many to ban in each subsequent iteration
|
||||
nodict_bigrams_initial: 800 # "Nodict" here means the n-grams were not found at all in the
|
||||
# human corpus.
|
||||
nodict_bigrams_subsequent: 100
|
||||
# Trigrams
|
||||
dict_trigrams_initial: 300
|
||||
dict_trigrams_subsequent: 50
|
||||
nodict_trigrams_initial: 800
|
||||
nodict_trigrams_subsequent: 100
|
||||
|
||||
# --- User-Defined N-gram Bans ---
|
||||
# User-supplied extra n-grams to always ban (processed by auto-antislop)
|
||||
extra_ngrams_to_ban: [
|
||||
# "voice barely whisper",
|
||||
]
|
||||
|
||||
################################################################################
|
||||
# OVER-REPRESENTED WORD ANALYSIS & BANNING
|
||||
################################################################################
|
||||
compute_overrep_words: true
|
||||
top_k_words_for_overrep_analysis: 200000
|
||||
|
||||
# --- Quotas for Adding Over-represented Words to Slop Phrase Ban List ---
|
||||
dict_overrep_initial: 800 # How many of the top over-represented dictionary words to
|
||||
# ban in the first antislop iteration.
|
||||
# "Dictionary" means the words were also found in the human
|
||||
# writing corpus.
|
||||
dict_overrep_subsequent: 200 # How many to ban in each subsequent iteration
|
||||
nodict_overrep_initial: 80 # "Nodict" here means the n-grams were not found at all in the
|
||||
# human corpus.
|
||||
nodict_overrep_subsequent: 20
|
||||
|
||||
################################################################################
|
||||
# SLOP PHRASE BANNING
|
||||
################################################################################
|
||||
|
||||
# Slop phrases are over-represented whole phrases extracted from the generated texts.
|
||||
enable_slop_phrase_ban: true
|
||||
min_phrase_freq_to_keep: 2 # Min frequency for a new phrase from slop-forensics to be considered
|
||||
top_n_initial_slop_ban: 600 # New slop phrases from slop-forensics to ban in iter 0
|
||||
top_n_subsequent_slop_ban: 100 # New slop phrases from slop-forensics to ban in later iters
|
||||
|
||||
# --- User-Defined Slop Phrase Bans ---
|
||||
# User supplied list of strings to always ban
|
||||
# - case insensitive
|
||||
# To trigger a ban, the sequence must not have a word-like character
|
||||
# (not punctuation or whitespace) directly on either side. That is to say, we
|
||||
# are not banning disallowed sequences that occur as substrings in longer
|
||||
# words. The exception is if the banned string is already bookended by
|
||||
# a non-word character.
|
||||
#
|
||||
# Examples:
|
||||
# banned string "cat"
|
||||
# - won't trigger a ban for "cation"
|
||||
# - will trigger a ban on "cat[morecat]"
|
||||
# banned string "cat["
|
||||
# - *will* trigger a ban on "cat[morecat]", because the banned string
|
||||
# ends with a non-word character.
|
||||
extra_slop_phrases_to_ban: [
|
||||
# "testament to",
|
||||
#"…", "*", " –", "–", "#",
|
||||
]
|
||||
|
||||
# --- Whitelisted Strings ---
|
||||
# These will be excluded from the list of slop strings that the pipeline finds.
|
||||
# Note: special tokens in the tokenizer and parts of the chat template are
|
||||
# automatically whitelisted.
|
||||
whitelist_strings: [
|
||||
# "think", "thinking"
|
||||
]
|
||||
|
||||
################################################################################
|
||||
# REGEX BANNING
|
||||
################################################################################
|
||||
# User-supplied regex patterns to ban
|
||||
# Note: unoptimised regex patterns can slow down antislop generation, as they will be called often on large texts.
|
||||
extra_regex_patterns: [
|
||||
# These ones ban "it's not x, it's y" type patterns:
|
||||
|
||||
#"\\b(?:\\w+n(?:['’]t)|not\\s+(?:just|only|merely|because))\\s+(?:(?![.;:?!…]).){1,100}?[.;:?!…]\\s*(?:it|they|you)(?:['’](?:s|re|m))?\\b(?!\\s+(?:was|were|is|are|wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t)\\b)(?:\\s*[*…]?\\s*)?(?!when\\b|then\\b|but\\b|and\\b|yet\\b)(?!right\\b)(?!normal\\b)(?!true\\b)(?!sure\\b)(?!only\\b)(?!still\\b)(?!rarely\\b)(?!already\\b)(?!wrong\\b)(?!want\\b)(?!just\\b)(?!couldn\\b)(?!could\\b)(?!saw\\b)(?!started\\b)(?!remember\\b)(?!struggled\\b)(?!watched\\b)(?!goal\\b)(?!took\\b)(?!kept\\b)(?!reminded\\b)(?!time\\b)(?!have\\b)(?!acted\\b)(?!smiled\\b)(?!think\\b)(?!give\\b)(?!grab\\b)(?!gave\\b)(?!turn\\b)(?!justify\\b)(?!\\w+ly\\b)(?=[a-z]{4,}\\b)[a-z]+\\w*",
|
||||
|
||||
#"\\b(?:\\w+n(?:['’]t)|not)\\s+(?:just|only|merely)?\\s*(?:(?![-–—]|[.?!…]).){1,80}?[-–—]{1,2}\\s*\\w+(?:['’]\\w+)?\\s+",
|
||||
|
||||
#"\\b(?:wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t|not)\\s+(?!\\b(?:minute|minutes|hour|hours|day|days|year|years|second|seconds)\\b)(?!with\\b)(?!even\\b)(?:(?![.;:?!…]).){2,120}?[.;:?!…]\\s*(?:it|they|you|that)(?:\\s+(?:was|were|is|are)\\b(?:\\s+[*_~]?\\w+[*_~]?)?|(?:['’](?:s|re|m))\\b(?:\\s+[*_~]?\\w+[*_~]?)?)",
|
||||
|
||||
#"\\bno\\s+longer\\s+(?:just|only|merely)?\\s+[^.;:?!…]{1,120}[.;:?!…]\\s*(?:it|they|you)\\s+(?:is|are|was|were)\\b(?:\\s+[*_~]?\\w+[*_~]?)?",
|
||||
|
||||
#"\\b(?:wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t|not)\\s+(?:just|only|merely)?\\s*(?:(?!\\bbut\\b|[.?!…]).){1,80}?[,;:\\-–—]\\s*but\\s+(?!I\\b)(?:also\\s+)?"
|
||||
|
||||
]
|
||||
|
||||
################################################################################
|
||||
# FINETUNING
|
||||
################################################################################
|
||||
finetune_enabled: true
|
||||
|
||||
# --- General Finetuning Setup ---
|
||||
finetune_use_unsloth: false
|
||||
finetune_mode: "ftpo" # dpo / ftpo (final token preference optimisation)
|
||||
finetune_ftpo_dataset: "" # you can specify an existing ftpo dataset, or leave unset to let the
|
||||
# pipeline use the one produced in the generation step
|
||||
finetune_base_model_id: null # Base model for DPO (if unset, uses model_id)
|
||||
finetune_max_seq_length: 3500 # this may truncate some outputs
|
||||
finetune_load_in_4bit: true # qlora
|
||||
|
||||
# --- Early Stopping ---
|
||||
finetune_early_stopping_wins: 0.85 # Early stopping threshold for fraction of *chosen* completions that are selected over *rejected*.
|
||||
# More than 0.85 may be overtrained. Set to > 1.0 to disable early stopping.
|
||||
finetune_early_stopping_loss: null # Loss threshold for early stopping. Set to null to disable.
|
||||
|
||||
# --- LoRA Configuration ---
|
||||
finetune_lora_r: 128 # the ftpo trainer works best with a high lora rank
|
||||
finetune_lora_alpha: 128
|
||||
finetune_lora_dropout: 0.05
|
||||
finetune_weight_decay: 0.01
|
||||
finetune_target_modules: ["lm_head"] #["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "lm_head"]
|
||||
|
||||
# --- Layer Freezing ---
|
||||
finetune_freeze_early_layers: false
|
||||
finetune_n_layers_unfrozen: 10
|
||||
|
||||
# --- Training Process ---
|
||||
finetune_gradient_checkpointing: "unsloth"
|
||||
finetune_chat_template: "" # e.g. "gemma-3" -- get the chat template from unsloth's helper if required, otherwise leave the string blank to use the tokeniser's chat template
|
||||
finetune_batch_size: 1
|
||||
finetune_gradient_accumulation_steps: 16
|
||||
finetune_warmup_ratio: 0.1
|
||||
finetune_num_epochs: 1
|
||||
|
||||
# --- Learning Rate ---
|
||||
finetune_learning_rate: 0.000001
|
||||
finetune_auto_learning_rate: true # true: automatically determine learning rate based on dataset size, effective batch size & lora rank
|
||||
finetune_auto_learning_rate_adjustment_scaling: 0.15 # scale the auto-lr by this factor
|
||||
|
||||
# --- DPO/FTPO Specific ---
|
||||
finetune_beta: 0.1 # DPO beta
|
||||
|
||||
# --- Output & Saving ---
|
||||
finetune_output_dir_suffix: "_ftpo_exp01" # Appended to experiment run dir
|
||||
finetune_save_merged_16bit: true
|
||||
finetune_save_gguf_q8_0: false
|
||||
|
||||
# --- Dataset Handling for Finetuning ---
|
||||
finetune_max_train_examples: 8000 # adjust as needed
|
||||
finetune_shuffle_seed: 666
|
||||
|
||||
# --- FTPO Sample Regularization ---
|
||||
# 0 = off; 0.9 strongly downsamples overrepresented rule violations
|
||||
# (this is useful because the raw generated dataset is typically very skewed)
|
||||
ftpo_sample_rejected_regularisation_strength: 0.8
|
||||
ftpo_sample_chosen_regularisation_strength: 0.2
|
||||
ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list
|
||||
|
||||
|
||||
# ── FTPO-specific hyper-parameters ─────────────────────────────────────────
|
||||
# Leave any of these out (or set to null) to fall back to FTPOTrainer defaults.
|
||||
|
||||
# Loss terms are computed separately for the target (chosen + rejected) tokens vs the remainder of the vocab.
|
||||
# This is because we want to allow more freedom of movement for the target tokens.
|
||||
|
||||
# MSE loss term 1: light mse loss applied tokenwise on target tokens
|
||||
ftpo_lambda_mse_target: 0.05 # Strength of MSE loss tether on the individual logits in the
|
||||
# chosen+rejected set vs reference.
|
||||
ftpo_tau_mse_target: 0.5 # Grace bandwidth (logits) before the above MSE loss kicks in.
|
||||
|
||||
# MSE loss term 2: stronger mse term applied to remaining (non-target) vocab
|
||||
ftpo_lambda_mse: 0.4
|
||||
|
||||
ftpo_clip_epsilon_logits: 2 # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off"
|
||||
286
configs/Mistral-Small-3.2-24B-Instruct-2506.yaml
Normal file
286
configs/Mistral-Small-3.2-24B-Instruct-2506.yaml
Normal file
@@ -0,0 +1,286 @@
|
||||
################################################################################
|
||||
# MAIN AUTO-ANTISLOP CONFIGURATION
|
||||
################################################################################
|
||||
|
||||
################################################################################
|
||||
# RUN SETUP
|
||||
################################################################################
|
||||
experiment_base_dir: "results/auto_antislop_runs" # Base for timestamped run directories
|
||||
human_profile_path: "data/human_writing_profile.json"
|
||||
log_level: "INFO"
|
||||
# Iteration 0: Generates the baseline dataset & computes slop strings/ngrams to ban
|
||||
# Iteration 1: Generates a dataset using antislop, banning those strings & ngrams. Recomputes the slop strings/ngrams at the end & adds any new slop to the ban lists
|
||||
# Iteration 2+: Extra iterations catch slop that emerges after the initial set is banned
|
||||
num_iterations: 2 # Minimum 2 iterations (this is enough to catch most slop)
|
||||
model_id: "unsloth/Mistral-Small-3.2-24B-Instruct-2506" #"mistralai/Mistral-Small-3.2-24B-Instruct-2506" # Global model id for the pipeline. Can be overridden on individual steps.
|
||||
|
||||
# !! NEEDED TO SET ATTN TO EAGER
|
||||
# add to finetuning.py after model load:
|
||||
# os.environ["FLASH_ATTENTION_FORCE_EAGER"] = "1"
|
||||
# model.config._attn_implementation = "eager"
|
||||
|
||||
|
||||
################################################################################
|
||||
# VLLM SERVER MANAGEMENT (Conditional: if --manage-vllm is True)
|
||||
################################################################################
|
||||
manage_vllm: true
|
||||
vllm_model_id: null # Model served by vLLM (if unset, will use model_id)
|
||||
vllm_port: 8000
|
||||
vllm_hf_token: null # Optional: Your Hugging Face token if model is gated
|
||||
vllm_cuda_visible_devices: "0" # set to e.g. "0,1,2,3" for multiple gpus
|
||||
vllm_gpu_memory_utilization: 0.85 # leave some room for the refusal classifier if you are using it (about 3gb)
|
||||
vllm_max_model_len: 4500
|
||||
vllm_dtype: "bfloat16"
|
||||
# Additional raw CLI arguments for vLLM server, e.g., ["--tensor-parallel-size", "4"] for multiple gpus
|
||||
vllm_extra_args: ["--tokenizer_mode", "mistral", "--load_format", "mistral", "--config_format", "mistral"]
|
||||
vllm_env: # env vars for the vLLM process
|
||||
VLLM_USE_V1: "0" # may be needed for amd gpus
|
||||
|
||||
|
||||
################################################################################
|
||||
# GENERATION PARAMETERS (using antislop-vllm)
|
||||
################################################################################
|
||||
generation_step_enabled: true
|
||||
|
||||
# --- API & Model Configuration ---
|
||||
# If you set manage_vllm=true, leave the base url unset
|
||||
#generation_api_base_url: "http://localhost:8000/v1"
|
||||
#generation_api_base_url: "https://apjmbtwbrb8t61-8888.proxy.runpod.net/v1"
|
||||
generation_model_id: null # Model id for generation requests (if unset, uses model_id)
|
||||
generation_api_key: "xxx" # API key for the vLLM server
|
||||
|
||||
# --- Core Generation Settings ---
|
||||
generation_max_new_tokens: 1000
|
||||
generation_threads: 50 # Number of parallel threads for API queries in antislop-vllm. Note: vllm can become very inefficient if you go over some concurrency threshold (depending on vram)
|
||||
generation_max_prompts: 1500 # Number of samples to generate from the prompts in the dataset
|
||||
|
||||
# --- Dataset & Chat Template ---
|
||||
generation_hf_dataset_name: 'ganjaninja/writing-prompts-sfw-nsfw-interleaved' #'Nitral-AI/Reddit-SFW-Writing_Prompts_ShareGPT'
|
||||
generation_hf_dataset_split: 'train'
|
||||
# A huggingface model id or local dir containing the tokeniser you want to use to apply chat templates.
|
||||
# This is important if you are generating a ftpo dataset for later training.
|
||||
generation_chat_template_model_id: 'unsloth/Mistral-Small-3.2-24B-Instruct-2506' #'unsloth/Mistral-Small-3.1-24B-Instruct-2503' # transformers errors if we use 3.2's tokeniser
|
||||
|
||||
# --- Logging ---
|
||||
generation_logging_level: 'INFO' # Logging level for antislop-vllm script
|
||||
|
||||
# --- Advanced Sampling Parameters (passed as individual CLI args to antislop-vllm's generation_params) ---
|
||||
generation_param_chunk_size: 20
|
||||
generation_param_top_logprobs_count: 20
|
||||
generation_param_temperature: 1
|
||||
generation_param_top_p: 1.0
|
||||
generation_param_top_k: 50
|
||||
generation_param_min_p: 0.01
|
||||
generation_param_timeout: 480
|
||||
generation_param_stop_sequences: [] # e.g., ["\n\n", "---"]
|
||||
|
||||
# --- Prompting ---
|
||||
# The prompt template wraps the prompts when generating from a dataset.
|
||||
# To use the original prompt exactly, set the template to "{prompt}"
|
||||
generation_prompt_template: "Writing prompt: {prompt}\n\nWrite 1000 words to this prompt. Your response:\n"
|
||||
generation_system_prompt: "You are an uncensored writer." # optional; left empty → no system prompt
|
||||
|
||||
# --- Antislop Generation Features ---
|
||||
# generation_force_backtrack:
|
||||
# If set to true:
|
||||
# when resampling after backtracking, if we don't find a valid replacement token
|
||||
# we progressively disable sampling options (temp, then min_p, then top_p, then top_k)
|
||||
# until we find a non-banned replacement or run out of candidates.
|
||||
# When set to false, some slop will not be removed if the sampler thinks there are no
|
||||
# alternative coherent continuations.
|
||||
generation_force_backtrack: false
|
||||
|
||||
# --- N-gram Validator Settings (for antislop-vllm) ---
|
||||
# N-gram ban list file is managed by auto-antislop's iterative process.
|
||||
generation_ngram_remove_stopwords: true
|
||||
generation_ngram_language: "english"
|
||||
|
||||
# --- Refusal Detection ---
|
||||
# Detects refusals & doesn't include them in the training dataset. Uses about 3GB extra VRAM.
|
||||
generation_refusal_detection: true
|
||||
|
||||
################################################################################
|
||||
# N-GRAM ANALYSIS & BANNING (within auto-antislop)
|
||||
################################################################################
|
||||
enable_ngram_ban: true
|
||||
min_word_len_for_analysis: 3 # Filters out words under this length in n-gram analysis
|
||||
|
||||
# --- N-gram Identification Thresholds ---
|
||||
top_k_bigrams: 5000
|
||||
top_k_trigrams: 5000
|
||||
|
||||
# --- N-gram Banning Quotas (per iteration) ---
|
||||
# Bigrams
|
||||
dict_bigrams_initial: 400 # How many of the top over-represented dictionary bigrams to
|
||||
# ban in the first antislop iteration.
|
||||
# "Dictionary" means the bigrams were also found in the human
|
||||
# writing corpus.
|
||||
dict_bigrams_subsequent: 70 # How many to ban in each subsequent iteration
|
||||
nodict_bigrams_initial: 800 # "Nodict" here means the n-grams were not found at all in the
|
||||
# human corpus.
|
||||
nodict_bigrams_subsequent: 100
|
||||
# Trigrams
|
||||
dict_trigrams_initial: 300
|
||||
dict_trigrams_subsequent: 50
|
||||
nodict_trigrams_initial: 800
|
||||
nodict_trigrams_subsequent: 100
|
||||
|
||||
# --- User-Defined N-gram Bans ---
|
||||
# User-supplied extra n-grams to always ban (processed by auto-antislop)
|
||||
extra_ngrams_to_ban: [
|
||||
# "voice barely whisper",
|
||||
]
|
||||
|
||||
################################################################################
|
||||
# OVER-REPRESENTED WORD ANALYSIS & BANNING
|
||||
################################################################################
|
||||
compute_overrep_words: true
|
||||
top_k_words_for_overrep_analysis: 200000
|
||||
|
||||
# --- Quotas for Adding Over-represented Words to Slop Phrase Ban List ---
|
||||
dict_overrep_initial: 800 # How many of the top over-represented dictionary words to
|
||||
# ban in the first antislop iteration.
|
||||
# "Dictionary" means the words were also found in the human
|
||||
# writing corpus.
|
||||
dict_overrep_subsequent: 200 # How many to ban in each subsequent iteration
|
||||
nodict_overrep_initial: 80 # "Nodict" here means the n-grams were not found at all in the
|
||||
# human corpus.
|
||||
nodict_overrep_subsequent: 20
|
||||
|
||||
################################################################################
|
||||
# SLOP PHRASE BANNING
|
||||
################################################################################
|
||||
|
||||
# Slop phrases are over-represented whole phrases extracted from the generated texts.
|
||||
enable_slop_phrase_ban: true
|
||||
min_phrase_freq_to_keep: 2 # Min frequency for a new phrase from slop-forensics to be considered
|
||||
top_n_initial_slop_ban: 600 # New slop phrases from slop-forensics to ban in iter 0
|
||||
top_n_subsequent_slop_ban: 100 # New slop phrases from slop-forensics to ban in later iters
|
||||
|
||||
# --- User-Defined Slop Phrase Bans ---
|
||||
# User supplied list of strings to always ban
|
||||
# - case insensitive
|
||||
# To trigger a ban, the sequence must not have a word-like character
|
||||
# (not punctuation or whitespace) directly on either side. That is to say, we
|
||||
# are not banning disallowed sequences that occur as substrings in longer
|
||||
# words. The exception is if the banned string is already bookended by
|
||||
# a non-word character.
|
||||
#
|
||||
# Examples:
|
||||
# banned string "cat"
|
||||
# - won't trigger a ban for "cation"
|
||||
# - will trigger a ban on "cat[morecat]"
|
||||
# banned string "cat["
|
||||
# - *will* trigger a ban on "cat[morecat]", because the banned string
|
||||
# ends with a non-word character.
|
||||
extra_slop_phrases_to_ban: [
|
||||
# "testament to",
|
||||
#"…", "*", " –", "–", "#",
|
||||
]
|
||||
|
||||
# --- Whitelisted Strings ---
|
||||
# These will be excluded from the list of slop strings that the pipeline finds.
|
||||
# Note: special tokens in the tokenizer and parts of the chat template are
|
||||
# automatically whitelisted.
|
||||
whitelist_strings: [
|
||||
# "think", "thinking"
|
||||
]
|
||||
|
||||
################################################################################
|
||||
# REGEX BANNING
|
||||
################################################################################
|
||||
# User-supplied regex patterns to ban
|
||||
# Note: unoptimised regex patterns can slow down antislop generation, as they will be called often on large texts.
|
||||
extra_regex_patterns: [
|
||||
# These ones ban "it's not x, it's y" type patterns:
|
||||
|
||||
#"\\b(?:\\w+n(?:['’]t)|not\\s+(?:just|only|merely|because))\\s+(?:(?![.;:?!…]).){1,100}?[.;:?!…]\\s*(?:it|they|you)(?:['’](?:s|re|m))?\\b(?!\\s+(?:was|were|is|are|wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t)\\b)(?:\\s*[*…]?\\s*)?(?!when\\b|then\\b|but\\b|and\\b|yet\\b)(?!right\\b)(?!normal\\b)(?!true\\b)(?!sure\\b)(?!only\\b)(?!still\\b)(?!rarely\\b)(?!already\\b)(?!wrong\\b)(?!want\\b)(?!just\\b)(?!couldn\\b)(?!could\\b)(?!saw\\b)(?!started\\b)(?!remember\\b)(?!struggled\\b)(?!watched\\b)(?!goal\\b)(?!took\\b)(?!kept\\b)(?!reminded\\b)(?!time\\b)(?!have\\b)(?!acted\\b)(?!smiled\\b)(?!think\\b)(?!give\\b)(?!grab\\b)(?!gave\\b)(?!turn\\b)(?!justify\\b)(?!\\w+ly\\b)(?=[a-z]{4,}\\b)[a-z]+\\w*",
|
||||
|
||||
#"\\b(?:\\w+n(?:['’]t)|not)\\s+(?:just|only|merely)?\\s*(?:(?![-–—]|[.?!…]).){1,80}?[-–—]{1,2}\\s*\\w+(?:['’]\\w+)?\\s+",
|
||||
|
||||
#"\\b(?:wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t|not)\\s+(?!\\b(?:minute|minutes|hour|hours|day|days|year|years|second|seconds)\\b)(?!with\\b)(?!even\\b)(?:(?![.;:?!…]).){2,120}?[.;:?!…]\\s*(?:it|they|you|that)(?:\\s+(?:was|were|is|are)\\b(?:\\s+[*_~]?\\w+[*_~]?)?|(?:['’](?:s|re|m))\\b(?:\\s+[*_~]?\\w+[*_~]?)?)",
|
||||
|
||||
#"\\bno\\s+longer\\s+(?:just|only|merely)?\\s+[^.;:?!…]{1,120}[.;:?!…]\\s*(?:it|they|you)\\s+(?:is|are|was|were)\\b(?:\\s+[*_~]?\\w+[*_~]?)?",
|
||||
|
||||
#"\\b(?:wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t|not)\\s+(?:just|only|merely)?\\s*(?:(?!\\bbut\\b|[.?!…]).){1,80}?[,;:\\-–—]\\s*but\\s+(?!I\\b)(?:also\\s+)?"
|
||||
|
||||
]
|
||||
|
||||
################################################################################
|
||||
# FINETUNING
|
||||
################################################################################
|
||||
finetune_enabled: true
|
||||
|
||||
# --- General Finetuning Setup ---
|
||||
finetune_use_unsloth: false
|
||||
finetune_mode: "ftpo" # dpo / ftpo (final token preference optimisation)
|
||||
finetune_ftpo_dataset: "" # you can specify an existing ftpo dataset, or leave unset to let the
|
||||
# pipeline use the one produced in the generation step
|
||||
finetune_base_model_id: null # Base model for DPO (if unset, uses model_id)
|
||||
finetune_max_seq_length: 2500 # this may truncate some outputs
|
||||
finetune_load_in_4bit: true # qlora
|
||||
|
||||
# --- Early Stopping ---
|
||||
finetune_early_stopping_wins: 0.85 # Early stopping threshold for fraction of *chosen* completions that are selected over *rejected*.
|
||||
# More than 0.85 may be overtrained. Set to > 1.0 to disable early stopping.
|
||||
finetune_early_stopping_loss: null # Loss threshold for early stopping. Set to null to disable.
|
||||
|
||||
# --- LoRA Configuration ---
|
||||
finetune_lora_r: 200 # the ftpo trainer works best with a high lora rank
|
||||
finetune_lora_alpha: 64
|
||||
finetune_lora_dropout: 0.05
|
||||
finetune_weight_decay: 0.01
|
||||
finetune_target_modules: ["up_proj", "down_proj", "lm_head"]
|
||||
|
||||
# --- Layer Freezing ---
|
||||
finetune_freeze_early_layers: true
|
||||
finetune_n_layers_unfrozen: 3
|
||||
|
||||
# --- Training Process ---
|
||||
finetune_gradient_checkpointing: "unsloth"
|
||||
finetune_chat_template: "" # e.g. "gemma-3" -- get the chat template from unsloth's helper if required, otherwise leave the string blank to use the tokeniser>
|
||||
finetune_batch_size: 1
|
||||
finetune_gradient_accumulation_steps: 16
|
||||
finetune_warmup_ratio: 0.1
|
||||
finetune_num_epochs: 1
|
||||
|
||||
# --- Learning Rate ---
|
||||
finetune_learning_rate: 0.000001
|
||||
finetune_auto_learning_rate: true # true: automatically determine learning rate based on dataset size, effective batch size & lora rank
|
||||
finetune_auto_learning_rate_adjustment_scaling: 0.04 # scale the auto-lr by this factor
|
||||
|
||||
# --- DPO/FTPO Specific ---
|
||||
finetune_beta: 0.1 # DPO beta
|
||||
|
||||
# --- Output & Saving ---
|
||||
finetune_output_dir_suffix: "_ftpo_exp01" # Appended to experiment run dir
|
||||
finetune_save_merged_16bit: true
|
||||
finetune_save_gguf_q8_0: false
|
||||
|
||||
# --- Dataset Handling for Finetuning ---
|
||||
finetune_max_train_examples: 14000 # adjust as needed
|
||||
finetune_shuffle_seed: 666
|
||||
|
||||
# --- FTPO Sample Regularization ---
|
||||
# 0 = off; 0.9 strongly downsamples overrepresented rule violations
|
||||
# (this is useful because the raw generated dataset is typically very skewed)
|
||||
ftpo_sample_rejected_regularisation_strength: 0.7
|
||||
ftpo_sample_chosen_regularisation_strength: 0.2
|
||||
ftpo_sample_min_chosen_tokens: 3 # filter out ftpo samples that have fewer than this number in the chosen tokens list
|
||||
|
||||
|
||||
# ── FTPO-specific hyper-parameters ─────────────────────────────────────────
|
||||
# Leave any of these out (or set to null) to fall back to FTPOTrainer defaults.
|
||||
|
||||
# Loss terms are computed separately for the target (chosen + rejected) tokens vs the remainder of the vocab.
|
||||
# This is because we want to allow more freedom of movement for the target tokens.
|
||||
|
||||
# MSE loss term 1: light mse loss applied tokenwise on target tokens
|
||||
ftpo_lambda_mse_target: 0.05 # Strength of MSE loss tether on the individual logits in the
|
||||
# chosen+rejected set vs reference.
|
||||
ftpo_tau_mse_target: 0.5 # Grace bandwidth (logits) before the above MSE loss kicks in.
|
||||
|
||||
# MSE loss term 2: stronger mse term applied to remaining (non-target) vocab
|
||||
ftpo_lambda_mse: 0.4
|
||||
|
||||
ftpo_clip_epsilon_logits: 2 # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off"
|
||||
279
configs/default.yaml
Normal file
279
configs/default.yaml
Normal file
@@ -0,0 +1,279 @@
|
||||
################################################################################
|
||||
# MAIN AUTO-ANTISLOP CONFIGURATION
|
||||
################################################################################
|
||||
|
||||
################################################################################
|
||||
# RUN SETUP
|
||||
################################################################################
|
||||
experiment_base_dir: "results/auto_antislop_runs" # Base for timestamped run directories
|
||||
human_profile_path: "data/human_writing_profile.json"
|
||||
log_level: "INFO"
|
||||
# Iteration 0: Generates the baseline dataset & computes slop strings/ngrams to ban
|
||||
# Iteration 1: Generates a dataset using antislop, banning those strings & ngrams. Recomputes the slop strings/ngrams at the end & adds any new slop to the ban lists
|
||||
# Iteration 2+: Extra iterations catch slop that emerges after the initial set is banned
|
||||
num_iterations: 2 # Minimum 2 iterations (this is enough to catch most slop)
|
||||
model_id: "unsloth/gemma-3-4b-it" # Global model id for the pipeline. Can be overridden on individual steps.
|
||||
|
||||
################################################################################
|
||||
# VLLM SERVER MANAGEMENT (Conditional: if --manage-vllm is True)
|
||||
################################################################################
|
||||
manage_vllm: true
|
||||
vllm_model_id: null # Model served by vLLM (if unset, will use model_id)
|
||||
vllm_port: 8000
|
||||
vllm_hf_token: null # Optional: Your Hugging Face token if model is gated
|
||||
vllm_cuda_visible_devices: "0" # set to e.g. "0,1,2,3" for multiple gpus
|
||||
vllm_gpu_memory_utilization: 0.85 # leave some room for the refusal classifier if you are using it (about 3gb)
|
||||
vllm_max_model_len: 2500
|
||||
vllm_dtype: "bfloat16"
|
||||
# Additional raw CLI arguments for vLLM server, e.g., ["--tensor-parallel-size", "4"] for multiple gpus
|
||||
vllm_extra_args: [] # each param as a separate string, e.g. ["--quantization", "bitsandbytes"]
|
||||
vllm_env: # env vars for the vLLM process
|
||||
# VLLM_USE_V1: "1" # may be needed for amd gpus
|
||||
|
||||
|
||||
################################################################################
|
||||
# GENERATION PARAMETERS (using antislop-vllm)
|
||||
################################################################################
|
||||
generation_step_enabled: true
|
||||
|
||||
# --- API & Model Configuration ---
|
||||
# If you set manage_vllm=true, leave the base url unset
|
||||
#generation_api_base_url: "http://localhost:8000/v1"
|
||||
#generation_api_base_url: "https://apjmbtwbrb8t61-8888.proxy.runpod.net/v1"
|
||||
generation_model_id: null # Model id for generation requests (if unset, uses model_id)
|
||||
generation_api_key: "xxx" # API key for the vLLM server
|
||||
|
||||
# --- Core Generation Settings ---
|
||||
generation_max_new_tokens: 1000
|
||||
generation_threads: 50 # Number of parallel threads for API queries in antislop-vllm. Note: vllm can become very inefficient if you go over some concurrency threshold (depending on vram)
|
||||
generation_max_prompts: 1000 # Number of samples to generate from the prompts in the dataset
|
||||
|
||||
# --- Dataset & Chat Template ---
|
||||
generation_hf_dataset_name: 'Nitral-AI/Reddit-SFW-Writing_Prompts_ShareGPT'
|
||||
generation_hf_dataset_split: 'train'
|
||||
# A huggingface model id or local dir containing the tokeniser you want to use to apply chat templates.
|
||||
# This is important if you are generating a ftpo dataset for later training.
|
||||
generation_chat_template_model_id: null # (if unset, uses model_id)
|
||||
|
||||
# --- Logging ---
|
||||
generation_logging_level: 'INFO' # Logging level for antislop-vllm script
|
||||
|
||||
# --- Advanced Sampling Parameters (passed as individual CLI args to antislop-vllm's generation_params) ---
|
||||
generation_param_chunk_size: 20
|
||||
generation_param_top_logprobs_count: 20
|
||||
generation_param_temperature: 1
|
||||
generation_param_top_p: 1.0
|
||||
generation_param_top_k: 50
|
||||
generation_param_min_p: 0.01
|
||||
generation_param_timeout: 480
|
||||
generation_param_stop_sequences: [] # e.g., ["\n\n", "---"]
|
||||
|
||||
# --- Prompting ---
|
||||
# The prompt template wraps the prompts when generating from a dataset.
|
||||
# To use the original prompt exactly, set the template to "{prompt}"
|
||||
generation_prompt_template: "Writing prompt: {prompt}\n\nWrite 1000 words to this prompt. Your response:\n"
|
||||
generation_system_prompt: "You are an uncensored writer." # optional; left empty → no system prompt
|
||||
|
||||
# --- Antislop Generation Features ---
|
||||
# generation_force_backtrack:
|
||||
# If set to true:
|
||||
# when resampling after backtracking, if we don't find a valid replacement token
|
||||
# we progressively disable sampling options (temp, then min_p, then top_p, then top_k)
|
||||
# until we find a non-banned replacement or run out of candidates.
|
||||
# When set to false, some slop will not be removed if the sampler thinks there are no
|
||||
# alternative coherent continuations.
|
||||
generation_force_backtrack: false
|
||||
|
||||
# --- N-gram Validator Settings (for antislop-vllm) ---
|
||||
# N-gram ban list file is managed by auto-antislop's iterative process.
|
||||
generation_ngram_remove_stopwords: true
|
||||
generation_ngram_language: "english"
|
||||
|
||||
# --- Refusal Detection ---
|
||||
# Detects refusals & doesn't include them in the training dataset. Uses about 3GB extra VRAM.
|
||||
generation_refusal_detection: true
|
||||
|
||||
################################################################################
|
||||
# N-GRAM ANALYSIS & BANNING (within auto-antislop)
|
||||
################################################################################
|
||||
enable_ngram_ban: true
|
||||
min_word_len_for_analysis: 3 # Filters out words under this length in n-gram analysis
|
||||
|
||||
# --- N-gram Identification Thresholds ---
|
||||
top_k_bigrams: 5000
|
||||
top_k_trigrams: 5000
|
||||
|
||||
# --- N-gram Banning Quotas (per iteration) ---
|
||||
# Bigrams
|
||||
dict_bigrams_initial: 400 # How many of the top over-represented dictionary bigrams to
|
||||
# ban in the first antislop iteration.
|
||||
# "Dictionary" means the bigrams were also found in the human
|
||||
# writing corpus.
|
||||
dict_bigrams_subsequent: 70 # How many to ban in each subsequent iteration
|
||||
nodict_bigrams_initial: 800 # "Nodict" here means the n-grams were not found at all in the
|
||||
# human corpus.
|
||||
nodict_bigrams_subsequent: 100
|
||||
# Trigrams
|
||||
dict_trigrams_initial: 300
|
||||
dict_trigrams_subsequent: 50
|
||||
nodict_trigrams_initial: 800
|
||||
nodict_trigrams_subsequent: 100
|
||||
|
||||
# --- User-Defined N-gram Bans ---
|
||||
# User-supplied extra n-grams to always ban (processed by auto-antislop)
|
||||
extra_ngrams_to_ban: [
|
||||
# "voice barely whisper",
|
||||
]
|
||||
|
||||
################################################################################
|
||||
# OVER-REPRESENTED WORD ANALYSIS & BANNING
|
||||
################################################################################
|
||||
compute_overrep_words: true
|
||||
top_k_words_for_overrep_analysis: 200000
|
||||
|
||||
# --- Quotas for Adding Over-represented Words to Slop Phrase Ban List ---
|
||||
dict_overrep_initial: 800 # How many of the top over-represented dictionary words to
|
||||
# ban in the first antislop iteration.
|
||||
# "Dictionary" means the words were also found in the human
|
||||
# writing corpus.
|
||||
dict_overrep_subsequent: 200 # How many to ban in each subsequent iteration
|
||||
nodict_overrep_initial: 80 # "Nodict" here means the n-grams were not found at all in the
|
||||
# human corpus.
|
||||
nodict_overrep_subsequent: 20
|
||||
|
||||
################################################################################
|
||||
# SLOP PHRASE BANNING
|
||||
################################################################################
|
||||
|
||||
# Slop phrases are over-represented whole phrases extracted from the generated texts.
|
||||
enable_slop_phrase_ban: true
|
||||
min_phrase_freq_to_keep: 2 # Min frequency for a new phrase from slop-forensics to be considered
|
||||
top_n_initial_slop_ban: 600 # New slop phrases from slop-forensics to ban in iter 0
|
||||
top_n_subsequent_slop_ban: 100 # New slop phrases from slop-forensics to ban in later iters
|
||||
|
||||
# --- User-Defined Slop Phrase Bans ---
|
||||
# User supplied list of strings to always ban
|
||||
# - case insensitive
|
||||
# To trigger a ban, the sequence must not have a word-like character
|
||||
# (not punctuation or whitespace) directly on either side. That is to say, we
|
||||
# are not banning disallowed sequences that occur as substrings in longer
|
||||
# words. The exception is if the banned string is already bookended by
|
||||
# a non-word character.
|
||||
#
|
||||
# Examples:
|
||||
# banned string "cat"
|
||||
# - won't trigger a ban for "cation"
|
||||
# - will trigger a ban on "cat[morecat]"
|
||||
# banned string "cat["
|
||||
# - *will* trigger a ban on "cat[morecat]", because the banned string
|
||||
# ends with a non-word character.
|
||||
extra_slop_phrases_to_ban: [
|
||||
# "testament to",
|
||||
#"…", "*", " –", "–", "#",
|
||||
]
|
||||
|
||||
# --- Whitelisted Strings ---
|
||||
# These will be excluded from the list of slop strings that the pipeline finds.
|
||||
# Note: special tokens in the tokenizer and parts of the chat template are
|
||||
# automatically whitelisted.
|
||||
whitelist_strings: [
|
||||
# "think", "thinking"
|
||||
]
|
||||
|
||||
################################################################################
|
||||
# REGEX BANNING
|
||||
################################################################################
|
||||
# User-supplied regex patterns to ban
|
||||
# Note: unoptimised regex patterns can slow down antislop generation, as they will be called often on large texts.
|
||||
extra_regex_patterns: [
|
||||
# These ones ban "it's not x, it's y" type patterns:
|
||||
|
||||
#"\\b(?:\\w+n(?:['’]t)|not\\s+(?:just|only|merely|because))\\s+(?:(?![.;:?!…]).){1,100}?[.;:?!…]\\s*(?:it|they|you)(?:['’](?:s|re|m))?\\b(?!\\s+(?:was|were|is|are|wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t)\\b)(?:\\s*[*…]?\\s*)?(?!when\\b|then\\b|but\\b|and\\b|yet\\b)(?!right\\b)(?!normal\\b)(?!true\\b)(?!sure\\b)(?!only\\b)(?!still\\b)(?!rarely\\b)(?!already\\b)(?!wrong\\b)(?!want\\b)(?!just\\b)(?!couldn\\b)(?!could\\b)(?!saw\\b)(?!started\\b)(?!remember\\b)(?!struggled\\b)(?!watched\\b)(?!goal\\b)(?!took\\b)(?!kept\\b)(?!reminded\\b)(?!time\\b)(?!have\\b)(?!acted\\b)(?!smiled\\b)(?!think\\b)(?!give\\b)(?!grab\\b)(?!gave\\b)(?!turn\\b)(?!justify\\b)(?!\\w+ly\\b)(?=[a-z]{4,}\\b)[a-z]+\\w*",
|
||||
|
||||
#"\\b(?:\\w+n(?:['’]t)|not)\\s+(?:just|only|merely)?\\s*(?:(?![-–—]|[.?!…]).){1,80}?[-–—]{1,2}\\s*\\w+(?:['’]\\w+)?\\s+",
|
||||
|
||||
#"\\b(?:wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t|not)\\s+(?!\\b(?:minute|minutes|hour|hours|day|days|year|years|second|seconds)\\b)(?!with\\b)(?!even\\b)(?:(?![.;:?!…]).){2,120}?[.;:?!…]\\s*(?:it|they|you|that)(?:\\s+(?:was|were|is|are)\\b(?:\\s+[*_~]?\\w+[*_~]?)?|(?:['’](?:s|re|m))\\b(?:\\s+[*_~]?\\w+[*_~]?)?)",
|
||||
|
||||
#"\\bno\\s+longer\\s+(?:just|only|merely)?\\s+[^.;:?!…]{1,120}[.;:?!…]\\s*(?:it|they|you)\\s+(?:is|are|was|were)\\b(?:\\s+[*_~]?\\w+[*_~]?)?",
|
||||
|
||||
#"\\b(?:wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t|not)\\s+(?:just|only|merely)?\\s*(?:(?!\\bbut\\b|[.?!…]).){1,80}?[,;:\\-–—]\\s*but\\s+(?!I\\b)(?:also\\s+)?"
|
||||
|
||||
]
|
||||
|
||||
################################################################################
|
||||
# FINETUNING
|
||||
################################################################################
|
||||
finetune_enabled: true
|
||||
|
||||
# --- General Finetuning Setup ---
|
||||
finetune_use_unsloth: false
|
||||
finetune_mode: "ftpo" # dpo / ftpo (final token preference optimisation)
|
||||
finetune_ftpo_dataset: "" # you can specify an existing ftpo dataset, or leave unset to let the
|
||||
# pipeline use the one produced in the generation step
|
||||
finetune_base_model_id: null # Base model for DPO (if unset, uses model_id)
|
||||
finetune_max_seq_length: 3500 # this may truncate some outputs
|
||||
finetune_load_in_4bit: false # qlora
|
||||
|
||||
# --- Early Stopping ---
|
||||
finetune_early_stopping_wins: 0.85 # Early stopping threshold for fraction of *chosen* completions that are selected over *rejected*.
|
||||
# More than 0.85 may be overtrained. Set to > 1.0 to disable early stopping.
|
||||
finetune_early_stopping_loss: null # Loss threshold for early stopping. Set to null to disable.
|
||||
|
||||
# --- LoRA Configuration ---
|
||||
finetune_lora_r: 128 # the ftpo trainer works best with a high lora rank
|
||||
finetune_lora_alpha: 128
|
||||
finetune_lora_dropout: 0.05
|
||||
finetune_weight_decay: 0.01
|
||||
finetune_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "lm_head"]
|
||||
|
||||
# --- Layer Freezing ---
|
||||
finetune_freeze_early_layers: true
|
||||
finetune_n_layers_unfrozen: 10
|
||||
|
||||
# --- Training Process ---
|
||||
finetune_gradient_checkpointing: "unsloth"
|
||||
finetune_chat_template: "" # e.g. "gemma-3" -- get the chat template from unsloth's helper if required, otherwise leave the string blank to use the tokeniser's chat template
|
||||
finetune_batch_size: 1
|
||||
finetune_gradient_accumulation_steps: 16
|
||||
finetune_warmup_ratio: 0.1
|
||||
finetune_num_epochs: 1
|
||||
|
||||
# --- Learning Rate ---
|
||||
finetune_learning_rate: 0.000001
|
||||
finetune_auto_learning_rate: true # true: automatically determine learning rate based on dataset size, effective batch size & lora rank
|
||||
finetune_auto_learning_rate_adjustment_scaling: 0.15 # scale the auto-lr by this factor
|
||||
|
||||
# --- DPO/FTPO Specific ---
|
||||
finetune_beta: 0.1 # DPO beta
|
||||
|
||||
# --- Output & Saving ---
|
||||
finetune_output_dir_suffix: "_ftpo_exp01" # Appended to experiment run dir
|
||||
finetune_save_merged_16bit: true
|
||||
finetune_save_gguf_q8_0: false
|
||||
|
||||
# --- Dataset Handling for Finetuning ---
|
||||
finetune_max_train_examples: 1000 # adjust as needed
|
||||
finetune_shuffle_seed: 666
|
||||
|
||||
# --- FTPO Sample Regularization ---
|
||||
# 0 = off; 0.9 strongly downsamples overrepresented rule violations
|
||||
# (this is useful because the raw generated dataset is typically very skewed)
|
||||
ftpo_sample_rejected_regularisation_strength: 0.8
|
||||
ftpo_sample_chosen_regularisation_strength: 0.2
|
||||
ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list
|
||||
|
||||
# ── FTPO-specific hyper-parameters ─────────────────────────────────────────
|
||||
# Leave any of these out (or set to null) to fall back to FTPOTrainer defaults.
|
||||
|
||||
# Loss terms are computed separately for the target (chosen + rejected) tokens vs the remainder of the vocab.
|
||||
# This is because we want to allow more freedom of movement for the target tokens.
|
||||
|
||||
# MSE loss term 1: light mse loss applied tokenwise on target tokens
|
||||
ftpo_lambda_mse_target: 0.05 # Strength of MSE loss tether on the individual logits in the
|
||||
# chosen+rejected set vs reference.
|
||||
ftpo_tau_mse_target: 0.5 # Grace bandwidth (logits) before the above MSE loss kicks in.
|
||||
|
||||
# MSE loss term 2: stronger mse term applied to remaining (non-target) vocab
|
||||
ftpo_lambda_mse: 0.4
|
||||
|
||||
ftpo_clip_epsilon_logits: 2 # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off"
|
||||
279
configs/gemma-3-12b-it.yaml
Normal file
279
configs/gemma-3-12b-it.yaml
Normal file
@@ -0,0 +1,279 @@
|
||||
################################################################################
|
||||
# MAIN AUTO-ANTISLOP CONFIGURATION
|
||||
################################################################################
|
||||
|
||||
################################################################################
|
||||
# RUN SETUP
|
||||
################################################################################
|
||||
experiment_base_dir: "results/auto_antislop_runs" # Base for timestamped run directories
|
||||
human_profile_path: "data/human_writing_profile.json"
|
||||
log_level: "INFO"
|
||||
# Iteration 0: Generates the baseline dataset & computes slop strings/ngrams to ban
|
||||
# Iteration 1: Generates a dataset using antislop, banning those strings & ngrams. Recomputes the slop strings/ngrams at the end & adds any new slop to the ban lists
|
||||
# Iteration 2+: Extra iterations catch slop that emerges after the initial set is banned
|
||||
num_iterations: 2 # Minimum 2 iterations (this is enough to catch most slop)
|
||||
model_id: "unsloth/gemma-3-12b-it" # Global model id for the pipeline. Can be overridden on individual steps.
|
||||
|
||||
################################################################################
|
||||
# VLLM SERVER MANAGEMENT (Conditional: if --manage-vllm is True)
|
||||
################################################################################
|
||||
manage_vllm: true
|
||||
vllm_model_id: null # Model served by vLLM (if unset, will use model_id)
|
||||
vllm_port: 8000
|
||||
vllm_hf_token: null # Optional: Your Hugging Face token if model is gated
|
||||
vllm_cuda_visible_devices: "0" # set to e.g. "0,1,2,3" for multiple gpus
|
||||
vllm_gpu_memory_utilization: 0.85 # leave some room for the refusal classifier if you are using it (about 3gb)
|
||||
vllm_max_model_len: 4500
|
||||
vllm_dtype: "bfloat16"
|
||||
# Additional raw CLI arguments for vLLM server, e.g., ["--tensor-parallel-size", "4"] for multiple gpus
|
||||
vllm_extra_args: [] # each param as a separate string, e.g. ["--quantization", "bitsandbytes"]
|
||||
vllm_env: # env vars for the vLLM process
|
||||
# VLLM_USE_V1: "1" # may be needed for amd gpus
|
||||
|
||||
|
||||
################################################################################
|
||||
# GENERATION PARAMETERS (using antislop-vllm)
|
||||
################################################################################
|
||||
generation_step_enabled: true
|
||||
|
||||
# --- API & Model Configuration ---
|
||||
# If you set manage_vllm=true, leave the base url unset
|
||||
#generation_api_base_url: "http://localhost:8000/v1"
|
||||
#generation_api_base_url: "https://apjmbtwbrb8t61-8888.proxy.runpod.net/v1"
|
||||
generation_model_id: null # Model id for generation requests (if unset, uses model_id)
|
||||
generation_api_key: "xxx" # API key for the vLLM server
|
||||
|
||||
# --- Core Generation Settings ---
|
||||
generation_max_new_tokens: 1000
|
||||
generation_threads: 50 # Number of parallel threads for API queries in antislop-vllm. Note: vllm can become very inefficient if you go over some concurrency threshold (depending on vram)
|
||||
generation_max_prompts: 2000 # Number of samples to generate from the prompts in the dataset
|
||||
|
||||
# --- Dataset & Chat Template ---
|
||||
generation_hf_dataset_name: 'Nitral-AI/Reddit-SFW-Writing_Prompts_ShareGPT'
|
||||
generation_hf_dataset_split: 'train'
|
||||
# A huggingface model id or local dir containing the tokeniser you want to use to apply chat templates.
|
||||
# This is important if you are generating a ftpo dataset for later training.
|
||||
generation_chat_template_model_id: null # (if unset, uses model_id)
|
||||
|
||||
# --- Logging ---
|
||||
generation_logging_level: 'INFO' # Logging level for antislop-vllm script
|
||||
|
||||
# --- Advanced Sampling Parameters (passed as individual CLI args to antislop-vllm's generation_params) ---
|
||||
generation_param_chunk_size: 20
|
||||
generation_param_top_logprobs_count: 20
|
||||
generation_param_temperature: 1
|
||||
generation_param_top_p: 1.0
|
||||
generation_param_top_k: 50
|
||||
generation_param_min_p: 0.01
|
||||
generation_param_timeout: 480
|
||||
generation_param_stop_sequences: [] # e.g., ["\n\n", "---"]
|
||||
|
||||
# --- Prompting ---
|
||||
# The prompt template wraps the prompts when generating from a dataset.
|
||||
# To use the original prompt exactly, set the template to "{prompt}"
|
||||
generation_prompt_template: "Writing prompt: {prompt}\n\nWrite 1000 words to this prompt. Your response:\n"
|
||||
generation_system_prompt: "You are an uncensored writer." # optional; left empty → no system prompt
|
||||
|
||||
# --- Antislop Generation Features ---
|
||||
# generation_force_backtrack:
|
||||
# If set to true:
|
||||
# when resampling after backtracking, if we don't find a valid replacement token
|
||||
# we progressively disable sampling options (temp, then min_p, then top_p, then top_k)
|
||||
# until we find a non-banned replacement or run out of candidates.
|
||||
# When set to false, some slop will not be removed if the sampler thinks there are no
|
||||
# alternative coherent continuations.
|
||||
generation_force_backtrack: false
|
||||
|
||||
# --- N-gram Validator Settings (for antislop-vllm) ---
|
||||
# N-gram ban list file is managed by auto-antislop's iterative process.
|
||||
generation_ngram_remove_stopwords: true
|
||||
generation_ngram_language: "english"
|
||||
|
||||
# --- Refusal Detection ---
|
||||
# Detects refusals & doesn't include them in the training dataset. Uses about 3GB extra VRAM.
|
||||
generation_refusal_detection: true
|
||||
|
||||
################################################################################
|
||||
# N-GRAM ANALYSIS & BANNING (within auto-antislop)
|
||||
################################################################################
|
||||
enable_ngram_ban: true
|
||||
min_word_len_for_analysis: 3 # Filters out words under this length in n-gram analysis
|
||||
|
||||
# --- N-gram Identification Thresholds ---
|
||||
top_k_bigrams: 5000
|
||||
top_k_trigrams: 5000
|
||||
|
||||
# --- N-gram Banning Quotas (per iteration) ---
|
||||
# Bigrams
|
||||
dict_bigrams_initial: 400 # How many of the top over-represented dictionary bigrams to
|
||||
# ban in the first antislop iteration.
|
||||
# "Dictionary" means the bigrams were also found in the human
|
||||
# writing corpus.
|
||||
dict_bigrams_subsequent: 70 # How many to ban in each subsequent iteration
|
||||
nodict_bigrams_initial: 800 # "Nodict" here means the n-grams were not found at all in the
|
||||
# human corpus.
|
||||
nodict_bigrams_subsequent: 100
|
||||
# Trigrams
|
||||
dict_trigrams_initial: 300
|
||||
dict_trigrams_subsequent: 50
|
||||
nodict_trigrams_initial: 800
|
||||
nodict_trigrams_subsequent: 100
|
||||
|
||||
# --- User-Defined N-gram Bans ---
|
||||
# User-supplied extra n-grams to always ban (processed by auto-antislop)
|
||||
extra_ngrams_to_ban: [
|
||||
# "voice barely whisper",
|
||||
]
|
||||
|
||||
################################################################################
|
||||
# OVER-REPRESENTED WORD ANALYSIS & BANNING
|
||||
################################################################################
|
||||
compute_overrep_words: true
|
||||
top_k_words_for_overrep_analysis: 200000
|
||||
|
||||
# --- Quotas for Adding Over-represented Words to Slop Phrase Ban List ---
|
||||
dict_overrep_initial: 800 # How many of the top over-represented dictionary words to
|
||||
# ban in the first antislop iteration.
|
||||
# "Dictionary" means the words were also found in the human
|
||||
# writing corpus.
|
||||
dict_overrep_subsequent: 200 # How many to ban in each subsequent iteration
|
||||
nodict_overrep_initial: 80 # "Nodict" here means the n-grams were not found at all in the
|
||||
# human corpus.
|
||||
nodict_overrep_subsequent: 20
|
||||
|
||||
################################################################################
|
||||
# SLOP PHRASE BANNING
|
||||
################################################################################
|
||||
|
||||
# Slop phrases are over-represented whole phrases extracted from the generated texts.
|
||||
enable_slop_phrase_ban: true
|
||||
min_phrase_freq_to_keep: 2 # Min frequency for a new phrase from slop-forensics to be considered
|
||||
top_n_initial_slop_ban: 600 # New slop phrases from slop-forensics to ban in iter 0
|
||||
top_n_subsequent_slop_ban: 100 # New slop phrases from slop-forensics to ban in later iters
|
||||
|
||||
# --- User-Defined Slop Phrase Bans ---
|
||||
# User supplied list of strings to always ban
|
||||
# - case insensitive
|
||||
# To trigger a ban, the sequence must not have a word-like character
|
||||
# (not punctuation or whitespace) directly on either side. That is to say, we
|
||||
# are not banning disallowed sequences that occur as substrings in longer
|
||||
# words. The exception is if the banned string is already bookended by
|
||||
# a non-word character.
|
||||
#
|
||||
# Examples:
|
||||
# banned string "cat"
|
||||
# - won't trigger a ban for "cation"
|
||||
# - will trigger a ban on "cat[morecat]"
|
||||
# banned string "cat["
|
||||
# - *will* trigger a ban on "cat[morecat]", because the banned string
|
||||
# ends with a non-word character.
|
||||
extra_slop_phrases_to_ban: [
|
||||
"…", "...", "rain", "tapestry", "static", "regret", "rust"
|
||||
]
|
||||
|
||||
# --- Whitelisted Strings ---
|
||||
# These will be excluded from the list of slop strings that the pipeline finds.
|
||||
# Note: special tokens in the tokenizer and parts of the chat template are
|
||||
# automatically whitelisted.
|
||||
whitelist_strings: [
|
||||
# "think", "thinking"
|
||||
]
|
||||
|
||||
################################################################################
|
||||
# REGEX BANNING
|
||||
################################################################################
|
||||
# User-supplied regex patterns to ban
|
||||
# Note: unoptimised regex patterns can slow down antislop generation, as they will be called often on large texts.
|
||||
extra_regex_patterns: [
|
||||
# These ones ban "it's not x, it's y" type patterns:
|
||||
|
||||
#"\\b(?:\\w+n(?:['’]t)|not\\s+(?:just|only|merely|because))\\s+(?:(?![.;:?!…]).){1,100}?[.;:?!…]\\s*(?:it|they|you)(?:['’](?:s|re|m))?\\b(?!\\s+(?:was|were|is|are|wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t)\\b)(?:\\s*[*…]?\\s*)?(?!when\\b|then\\b|but\\b|and\\b|yet\\b)(?!right\\b)(?!normal\\b)(?!true\\b)(?!sure\\b)(?!only\\b)(?!still\\b)(?!rarely\\b)(?!already\\b)(?!wrong\\b)(?!want\\b)(?!just\\b)(?!couldn\\b)(?!could\\b)(?!saw\\b)(?!started\\b)(?!remember\\b)(?!struggled\\b)(?!watched\\b)(?!goal\\b)(?!took\\b)(?!kept\\b)(?!reminded\\b)(?!time\\b)(?!have\\b)(?!acted\\b)(?!smiled\\b)(?!think\\b)(?!give\\b)(?!grab\\b)(?!gave\\b)(?!turn\\b)(?!justify\\b)(?!\\w+ly\\b)(?=[a-z]{4,}\\b)[a-z]+\\w*",
|
||||
|
||||
#"\\b(?:\\w+n(?:['’]t)|not)\\s+(?:just|only|merely)?\\s*(?:(?![-–—]|[.?!…]).){1,80}?[-–—]{1,2}\\s*\\w+(?:['’]\\w+)?\\s+",
|
||||
|
||||
#"\\b(?:wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t|not)\\s+(?!\\b(?:minute|minutes|hour|hours|day|days|year|years|second|seconds)\\b)(?!with\\b)(?!even\\b)(?:(?![.;:?!…]).){2,120}?[.;:?!…]\\s*(?:it|they|you|that)(?:\\s+(?:was|were|is|are)\\b(?:\\s+[*_~]?\\w+[*_~]?)?|(?:['’](?:s|re|m))\\b(?:\\s+[*_~]?\\w+[*_~]?)?)",
|
||||
|
||||
#"\\bno\\s+longer\\s+(?:just|only|merely)?\\s+[^.;:?!…]{1,120}[.;:?!…]\\s*(?:it|they|you)\\s+(?:is|are|was|were)\\b(?:\\s+[*_~]?\\w+[*_~]?)?",
|
||||
|
||||
#"\\b(?:wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t|not)\\s+(?:just|only|merely)?\\s*(?:(?!\\bbut\\b|[.?!…]).){1,80}?[,;:\\-–—]\\s*but\\s+(?!I\\b)(?:also\\s+)?"
|
||||
|
||||
]
|
||||
|
||||
################################################################################
|
||||
# FINETUNING
|
||||
################################################################################
|
||||
finetune_enabled: true
|
||||
|
||||
# --- General Finetuning Setup ---
|
||||
finetune_use_unsloth: false
|
||||
finetune_mode: "ftpo" # dpo / ftpo (final token preference optimisation)
|
||||
finetune_ftpo_dataset: "" # you can specify an existing ftpo dataset, or leave unset to let the
|
||||
# pipeline use the one produced in the generation step
|
||||
finetune_base_model_id: null # Base model for DPO (if unset, uses model_id)
|
||||
finetune_max_seq_length: 2500 # this may truncate some outputs
|
||||
finetune_load_in_4bit: true # qlora
|
||||
|
||||
# --- Early Stopping ---
|
||||
finetune_early_stopping_wins: 0.85 # Early stopping threshold for fraction of *chosen* completions that are selected over *rejected*.
|
||||
# More than 0.85 may be overtrained. Set to > 1.0 to disable early stopping.
|
||||
finetune_early_stopping_loss: null # Loss threshold for early stopping. Set to null to disable.
|
||||
|
||||
# --- LoRA Configuration ---
|
||||
finetune_lora_r: 256 # the ftpo trainer works best with a high lora rank
|
||||
finetune_lora_alpha: 256
|
||||
finetune_lora_dropout: 0.05
|
||||
finetune_weight_decay: 0.01
|
||||
finetune_target_modules: ["up_proj", "down_proj", "lm_head"] #["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "lm_head"]
|
||||
|
||||
# --- Layer Freezing ---
|
||||
finetune_freeze_early_layers: true
|
||||
finetune_n_layers_unfrozen: 5
|
||||
|
||||
# --- Training Process ---
|
||||
finetune_gradient_checkpointing: "unsloth"
|
||||
finetune_chat_template: "" # e.g. "gemma-3" -- get the chat template from unsloth's helper if required, otherwise leave the string blank to use the tokeniser's chat template
|
||||
finetune_batch_size: 3
|
||||
finetune_gradient_accumulation_steps: 5
|
||||
finetune_warmup_ratio: 0.1
|
||||
finetune_num_epochs: 1
|
||||
|
||||
# --- Learning Rate ---
|
||||
finetune_learning_rate: 0.000001
|
||||
finetune_auto_learning_rate: true # true: automatically determine learning rate based on dataset size, effective batch size & lora rank
|
||||
finetune_auto_learning_rate_adjustment_scaling: 0.08 # scale the auto-lr by this factor
|
||||
|
||||
# --- DPO/FTPO Specific ---
|
||||
finetune_beta: 0.1 # DPO beta
|
||||
|
||||
# --- Output & Saving ---
|
||||
finetune_output_dir_suffix: "_ftpo_exp01" # Appended to experiment run dir
|
||||
finetune_save_merged_16bit: true
|
||||
finetune_save_gguf_q8_0: false
|
||||
|
||||
# --- Dataset Handling for Finetuning ---
|
||||
finetune_max_train_examples: 12000 # adjust as needed
|
||||
finetune_shuffle_seed: 666
|
||||
|
||||
# --- FTPO Sample Regularization ---
|
||||
# 0 = off; 0.9 strongly downsamples overrepresented rule violations
|
||||
# (this is useful because the raw generated dataset is typically very skewed)
|
||||
ftpo_sample_rejected_regularisation_strength: 0.8
|
||||
ftpo_sample_chosen_regularisation_strength: 0.2
|
||||
ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list
|
||||
|
||||
|
||||
# ── FTPO-specific hyper-parameters ─────────────────────────────────────────
|
||||
# Leave any of these out (or set to null) to fall back to FTPOTrainer defaults.
|
||||
|
||||
# Loss terms are computed separately for the target (chosen + rejected) tokens vs the remainder of the vocab.
|
||||
# This is because we want to allow more freedom of movement for the target tokens.
|
||||
|
||||
# MSE loss term 1: light mse loss applied tokenwise on target tokens
|
||||
ftpo_lambda_mse_target: 0.05 # Strength of MSE loss tether on the individual logits in the
|
||||
# chosen+rejected set vs reference.
|
||||
ftpo_tau_mse_target: 0.5 # Grace bandwidth (logits) before the above MSE loss kicks in.
|
||||
|
||||
# MSE loss term 2: stronger mse term applied to remaining (non-target) vocab
|
||||
ftpo_lambda_mse: 0.4
|
||||
|
||||
ftpo_clip_epsilon_logits: 2 # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off"
|
||||
279
configs/gemma-3-27b-it.yaml
Normal file
279
configs/gemma-3-27b-it.yaml
Normal file
@@ -0,0 +1,279 @@
|
||||
################################################################################
|
||||
# MAIN AUTO-ANTISLOP CONFIGURATION
|
||||
################################################################################
|
||||
|
||||
################################################################################
|
||||
# RUN SETUP
|
||||
################################################################################
|
||||
experiment_base_dir: "results/auto_antislop_runs" # Base for timestamped run directories
|
||||
human_profile_path: "data/human_writing_profile.json"
|
||||
log_level: "INFO"
|
||||
# Iteration 0: Generates the baseline dataset & computes slop strings/ngrams to ban
|
||||
# Iteration 1: Generates a dataset using antislop, banning those strings & ngrams. Recomputes the slop strings/ngrams at the end & adds any new slop to the ban lists
|
||||
# Iteration 2+: Extra iterations catch slop that emerges after the initial set is banned
|
||||
num_iterations: 2 # Minimum 2 iterations (this is enough to catch most slop)
|
||||
model_id: "unsloth/gemma-3-27b-it" # Global model id for the pipeline. Can be overridden on individual steps.
|
||||
|
||||
################################################################################
|
||||
# VLLM SERVER MANAGEMENT (Conditional: if --manage-vllm is True)
|
||||
################################################################################
|
||||
manage_vllm: true
|
||||
vllm_model_id: null # Model served by vLLM (if unset, will use model_id)
|
||||
vllm_port: 8000
|
||||
vllm_hf_token: null # Optional: Your Hugging Face token if model is gated
|
||||
vllm_cuda_visible_devices: "0" # set to e.g. "0,1,2,3" for multiple gpus
|
||||
vllm_gpu_memory_utilization: 0.92 # leave some room for the refusal classifier if you are using it (about 3gb)
|
||||
vllm_max_model_len: 4500
|
||||
vllm_dtype: "bfloat16"
|
||||
# Additional raw CLI arguments for vLLM server, e.g., ["--tensor-parallel-size", "4"] for multiple gpus
|
||||
vllm_extra_args: [] # each param as a separate string, e.g. ["--quantization", "bitsandbytes"]
|
||||
vllm_env: # env vars for the vLLM process
|
||||
# VLLM_USE_V1: "1" # may be needed for amd gpus
|
||||
|
||||
|
||||
################################################################################
|
||||
# GENERATION PARAMETERS (using antislop-vllm)
|
||||
################################################################################
|
||||
generation_step_enabled: true
|
||||
|
||||
# --- API & Model Configuration ---
|
||||
# If you set manage_vllm=true, leave the base url unset
|
||||
#generation_api_base_url: "http://localhost:8000/v1"
|
||||
#generation_api_base_url: "https://apjmbtwbrb8t61-8888.proxy.runpod.net/v1"
|
||||
generation_model_id: null # Model id for generation requests (if unset, uses model_id)
|
||||
generation_api_key: "xxx" # API key for the vLLM server
|
||||
|
||||
# --- Core Generation Settings ---
|
||||
generation_max_new_tokens: 1000
|
||||
generation_threads: 200 # Number of parallel threads for API queries in antislop-vllm. Note: vllm can become very inefficient if you go over some concurrency threshold (depending on vram)
|
||||
generation_max_prompts: 2000 # Number of samples to generate from the prompts in the dataset
|
||||
|
||||
# --- Dataset & Chat Template ---
|
||||
generation_hf_dataset_name: 'Nitral-AI/Reddit-SFW-Writing_Prompts_ShareGPT'
|
||||
generation_hf_dataset_split: 'train'
|
||||
# A huggingface model id or local dir containing the tokeniser you want to use to apply chat templates.
|
||||
# This is important if you are generating a ftpo dataset for later training.
|
||||
generation_chat_template_model_id: null # (if unset, uses model_id)
|
||||
|
||||
# --- Logging ---
|
||||
generation_logging_level: 'INFO' # Logging level for antislop-vllm script
|
||||
|
||||
# --- Advanced Sampling Parameters (passed as individual CLI args to antislop-vllm's generation_params) ---
|
||||
generation_param_chunk_size: 20
|
||||
generation_param_top_logprobs_count: 20
|
||||
generation_param_temperature: 1
|
||||
generation_param_top_p: 1.0
|
||||
generation_param_top_k: 50
|
||||
generation_param_min_p: 0.01
|
||||
generation_param_timeout: 480
|
||||
generation_param_stop_sequences: [] # e.g., ["\n\n", "---"]
|
||||
|
||||
# --- Prompting ---
|
||||
# The prompt template wraps the prompts when generating from a dataset.
|
||||
# To use the original prompt exactly, set the template to "{prompt}"
|
||||
generation_prompt_template: "Writing prompt: {prompt}\n\nWrite 1000 words to this prompt. Your response:\n"
|
||||
generation_system_prompt: "" # optional; left empty → no system prompt
|
||||
|
||||
# --- Antislop Generation Features ---
|
||||
# generation_force_backtrack:
|
||||
# If set to true:
|
||||
# when resampling after backtracking, if we don't find a valid replacement token
|
||||
# we progressively disable sampling options (temp, then min_p, then top_p, then top_k)
|
||||
# until we find a non-banned replacement or run out of candidates.
|
||||
# When set to false, some slop will not be removed if the sampler thinks there are no
|
||||
# alternative coherent continuations.
|
||||
generation_force_backtrack: false
|
||||
|
||||
# --- N-gram Validator Settings (for antislop-vllm) ---
|
||||
# N-gram ban list file is managed by auto-antislop's iterative process.
|
||||
generation_ngram_remove_stopwords: true
|
||||
generation_ngram_language: "english"
|
||||
|
||||
# --- Refusal Detection ---
|
||||
# Detects refusals & doesn't include them in the training dataset. Uses about 3GB extra VRAM.
|
||||
generation_refusal_detection: true
|
||||
|
||||
################################################################################
|
||||
# N-GRAM ANALYSIS & BANNING (within auto-antislop)
|
||||
################################################################################
|
||||
enable_ngram_ban: true
|
||||
min_word_len_for_analysis: 3 # Filters out words under this length in n-gram analysis
|
||||
|
||||
# --- N-gram Identification Thresholds ---
|
||||
top_k_bigrams: 5000
|
||||
top_k_trigrams: 5000
|
||||
|
||||
# --- N-gram Banning Quotas (per iteration) ---
|
||||
# Bigrams
|
||||
dict_bigrams_initial: 400 # How many of the top over-represented dictionary bigrams to
|
||||
# ban in the first antislop iteration.
|
||||
# "Dictionary" means the bigrams were also found in the human
|
||||
# writing corpus.
|
||||
dict_bigrams_subsequent: 70 # How many to ban in each subsequent iteration
|
||||
nodict_bigrams_initial: 800 # "Nodict" here means the n-grams were not found at all in the
|
||||
# human corpus.
|
||||
nodict_bigrams_subsequent: 100
|
||||
# Trigrams
|
||||
dict_trigrams_initial: 300
|
||||
dict_trigrams_subsequent: 50
|
||||
nodict_trigrams_initial: 800
|
||||
nodict_trigrams_subsequent: 100
|
||||
|
||||
# --- User-Defined N-gram Bans ---
|
||||
# User-supplied extra n-grams to always ban (processed by auto-antislop)
|
||||
extra_ngrams_to_ban: [
|
||||
# "voice barely whisper",
|
||||
]
|
||||
|
||||
################################################################################
|
||||
# OVER-REPRESENTED WORD ANALYSIS & BANNING
|
||||
################################################################################
|
||||
compute_overrep_words: true
|
||||
top_k_words_for_overrep_analysis: 200000
|
||||
|
||||
# --- Quotas for Adding Over-represented Words to Slop Phrase Ban List ---
|
||||
dict_overrep_initial: 800 # How many of the top over-represented dictionary words to
|
||||
# ban in the first antislop iteration.
|
||||
# "Dictionary" means the words were also found in the human
|
||||
# writing corpus.
|
||||
dict_overrep_subsequent: 200 # How many to ban in each subsequent iteration
|
||||
nodict_overrep_initial: 80 # "Nodict" here means the n-grams were not found at all in the
|
||||
# human corpus.
|
||||
nodict_overrep_subsequent: 20
|
||||
|
||||
################################################################################
|
||||
# SLOP PHRASE BANNING
|
||||
################################################################################
|
||||
|
||||
# Slop phrases are over-represented whole phrases extracted from the generated texts.
|
||||
enable_slop_phrase_ban: true
|
||||
min_phrase_freq_to_keep: 2 # Min frequency for a new phrase from slop-forensics to be considered
|
||||
top_n_initial_slop_ban: 600 # New slop phrases from slop-forensics to ban in iter 0
|
||||
top_n_subsequent_slop_ban: 100 # New slop phrases from slop-forensics to ban in later iters
|
||||
|
||||
# --- User-Defined Slop Phrase Bans ---
|
||||
# User supplied list of strings to always ban
|
||||
# - case insensitive
|
||||
# To trigger a ban, the sequence must not have a word-like character
|
||||
# (not punctuation or whitespace) directly on either side. That is to say, we
|
||||
# are not banning disallowed sequences that occur as substrings in longer
|
||||
# words. The exception is if the banned string is already bookended by
|
||||
# a non-word character.
|
||||
#
|
||||
# Examples:
|
||||
# banned string "cat"
|
||||
# - won't trigger a ban for "cation"
|
||||
# - will trigger a ban on "cat[morecat]"
|
||||
# banned string "cat["
|
||||
# - *will* trigger a ban on "cat[morecat]", because the banned string
|
||||
# ends with a non-word character.
|
||||
extra_slop_phrases_to_ban: [
|
||||
"…", "...", "rain", "tapestry", "static", "regret", "rust"
|
||||
]
|
||||
|
||||
# --- Whitelisted Strings ---
|
||||
# These will be excluded from the list of slop strings that the pipeline finds.
|
||||
# Note: special tokens in the tokenizer and parts of the chat template are
|
||||
# automatically whitelisted.
|
||||
whitelist_strings: [
|
||||
# "think", "thinking"
|
||||
]
|
||||
|
||||
################################################################################
|
||||
# REGEX BANNING
|
||||
################################################################################
|
||||
# User-supplied regex patterns to ban
|
||||
# Note: unoptimised regex patterns can slow down antislop generation, as they will be called often on large texts.
|
||||
extra_regex_patterns: [
|
||||
# These ones ban "it's not x, it's y" type patterns:
|
||||
|
||||
#"\\b(?:\\w+n(?:['’]t)|not\\s+(?:just|only|merely|because))\\s+(?:(?![.;:?!…]).){1,100}?[.;:?!…]\\s*(?:it|they|you)(?:['’](?:s|re|m))?\\b(?!\\s+(?:was|were|is|are|wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t)\\b)(?:\\s*[*…]?\\s*)?(?!when\\b|then\\b|but\\b|and\\b|yet\\b)(?!right\\b)(?!normal\\b)(?!true\\b)(?!sure\\b)(?!only\\b)(?!still\\b)(?!rarely\\b)(?!already\\b)(?!wrong\\b)(?!want\\b)(?!just\\b)(?!couldn\\b)(?!could\\b)(?!saw\\b)(?!started\\b)(?!remember\\b)(?!struggled\\b)(?!watched\\b)(?!goal\\b)(?!took\\b)(?!kept\\b)(?!reminded\\b)(?!time\\b)(?!have\\b)(?!acted\\b)(?!smiled\\b)(?!think\\b)(?!give\\b)(?!grab\\b)(?!gave\\b)(?!turn\\b)(?!justify\\b)(?!\\w+ly\\b)(?=[a-z]{4,}\\b)[a-z]+\\w*",
|
||||
|
||||
#"\\b(?:\\w+n(?:['’]t)|not)\\s+(?:just|only|merely)?\\s*(?:(?![-–—]|[.?!…]).){1,80}?[-–—]{1,2}\\s*\\w+(?:['’]\\w+)?\\s+",
|
||||
|
||||
#"\\b(?:wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t|not)\\s+(?!\\b(?:minute|minutes|hour|hours|day|days|year|years|second|seconds)\\b)(?!with\\b)(?!even\\b)(?:(?![.;:?!…]).){2,120}?[.;:?!…]\\s*(?:it|they|you|that)(?:\\s+(?:was|were|is|are)\\b(?:\\s+[*_~]?\\w+[*_~]?)?|(?:['’](?:s|re|m))\\b(?:\\s+[*_~]?\\w+[*_~]?)?)",
|
||||
|
||||
#"\\bno\\s+longer\\s+(?:just|only|merely)?\\s+[^.;:?!…]{1,120}[.;:?!…]\\s*(?:it|they|you)\\s+(?:is|are|was|were)\\b(?:\\s+[*_~]?\\w+[*_~]?)?",
|
||||
|
||||
#"\\b(?:wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t|not)\\s+(?:just|only|merely)?\\s*(?:(?!\\bbut\\b|[.?!…]).){1,80}?[,;:\\-–—]\\s*but\\s+(?!I\\b)(?:also\\s+)?"
|
||||
|
||||
]
|
||||
|
||||
################################################################################
|
||||
# FINETUNING
|
||||
################################################################################
|
||||
finetune_enabled: true
|
||||
|
||||
# --- General Finetuning Setup ---
|
||||
finetune_use_unsloth: false
|
||||
finetune_mode: "ftpo" # dpo / ftpo (final token preference optimisation)
|
||||
finetune_ftpo_dataset: "" # you can specify an existing ftpo dataset, or leave unset to let the
|
||||
# pipeline use the one produced in the generation step
|
||||
finetune_base_model_id: null # Base model for DPO (if unset, uses model_id)
|
||||
finetune_max_seq_length: 2500 # this may truncate some outputs
|
||||
finetune_load_in_4bit: true # qlora
|
||||
|
||||
# --- Early Stopping ---
|
||||
finetune_early_stopping_wins: 0.85 # Early stopping threshold for fraction of *chosen* completions that are selected over *rejected*.
|
||||
# More than 0.85 may be overtrained. Set to > 1.0 to disable early stopping.
|
||||
finetune_early_stopping_loss: null # Loss threshold for early stopping. Set to null to disable.
|
||||
|
||||
# --- LoRA Configuration ---
|
||||
finetune_lora_r: 128 # the ftpo trainer works best with a high lora rank
|
||||
finetune_lora_alpha: 128
|
||||
finetune_lora_dropout: 0.05
|
||||
finetune_weight_decay: 0.01
|
||||
finetune_target_modules: ["up_proj", "down_proj", "lm_head"] #["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "lm_head"]
|
||||
|
||||
# --- Layer Freezing ---
|
||||
finetune_freeze_early_layers: true
|
||||
finetune_n_layers_unfrozen: 5
|
||||
|
||||
# --- Training Process ---
|
||||
finetune_gradient_checkpointing: "unsloth"
|
||||
finetune_chat_template: "" # e.g. "gemma-3" -- get the chat template from unsloth's helper if required, otherwise leave the string blank to use the tokeniser's chat template
|
||||
finetune_batch_size: 1
|
||||
finetune_gradient_accumulation_steps: 16
|
||||
finetune_warmup_ratio: 0.1
|
||||
finetune_num_epochs: 1
|
||||
|
||||
# --- Learning Rate ---
|
||||
finetune_learning_rate: 0.000001
|
||||
finetune_auto_learning_rate: true # true: automatically determine learning rate based on dataset size, effective batch size & lora rank
|
||||
finetune_auto_learning_rate_adjustment_scaling: 0.08 # scale the auto-lr by this factor
|
||||
|
||||
# --- DPO/FTPO Specific ---
|
||||
finetune_beta: 0.1 # DPO beta
|
||||
|
||||
# --- Output & Saving ---
|
||||
finetune_output_dir_suffix: "_ftpo_exp01" # Appended to experiment run dir
|
||||
finetune_save_merged_16bit: true
|
||||
finetune_save_gguf_q8_0: false
|
||||
|
||||
# --- Dataset Handling for Finetuning ---
|
||||
finetune_max_train_examples: 12000 # adjust as needed
|
||||
finetune_shuffle_seed: 666
|
||||
|
||||
# --- FTPO Sample Regularization ---
|
||||
# 0 = off; 0.9 strongly downsamples overrepresented rule violations
|
||||
# (this is useful because the raw generated dataset is typically very skewed)
|
||||
ftpo_sample_rejected_regularisation_strength: 0.8
|
||||
ftpo_sample_chosen_regularisation_strength: 0.2
|
||||
ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list
|
||||
|
||||
|
||||
# ── FTPO-specific hyper-parameters ─────────────────────────────────────────
|
||||
# Leave any of these out (or set to null) to fall back to FTPOTrainer defaults.
|
||||
|
||||
# Loss terms are computed separately for the target (chosen + rejected) tokens vs the remainder of the vocab.
|
||||
# This is because we want to allow more freedom of movement for the target tokens.
|
||||
|
||||
# MSE loss term 1: light mse loss applied tokenwise on target tokens
|
||||
ftpo_lambda_mse_target: 0.05 # Strength of MSE loss tether on the individual logits in the
|
||||
# chosen+rejected set vs reference.
|
||||
ftpo_tau_mse_target: 0.5 # Grace bandwidth (logits) before the above MSE loss kicks in.
|
||||
|
||||
# MSE loss term 2: stronger mse term applied to remaining (non-target) vocab
|
||||
ftpo_lambda_mse: 0.4
|
||||
|
||||
ftpo_clip_epsilon_logits: 2 # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off"
|
||||
279
configs/gemma-3-4b-it.yaml
Normal file
279
configs/gemma-3-4b-it.yaml
Normal file
@@ -0,0 +1,279 @@
|
||||
################################################################################
|
||||
# MAIN AUTO-ANTISLOP CONFIGURATION
|
||||
################################################################################
|
||||
|
||||
################################################################################
|
||||
# RUN SETUP
|
||||
################################################################################
|
||||
experiment_base_dir: "results/auto_antislop_runs" # Base for timestamped run directories
|
||||
human_profile_path: "data/human_writing_profile.json"
|
||||
log_level: "INFO"
|
||||
# Iteration 0: Generates the baseline dataset & computes slop strings/ngrams to ban
|
||||
# Iteration 1: Generates a dataset using antislop, banning those strings & ngrams. Recomputes the slop strings/ngrams at the end & adds any new slop to the ban lists
|
||||
# Iteration 2+: Extra iterations catch slop that emerges after the initial set is banned
|
||||
num_iterations: 2 # Minimum 2 iterations (this is enough to catch most slop)
|
||||
model_id: "unsloth/gemma-3-4b-it" # Global model id for the pipeline. Can be overridden on individual steps.
|
||||
|
||||
################################################################################
|
||||
# VLLM SERVER MANAGEMENT (Conditional: if --manage-vllm is True)
|
||||
################################################################################
|
||||
manage_vllm: true
|
||||
vllm_model_id: null # Model served by vLLM (if unset, will use model_id)
|
||||
vllm_port: 8000
|
||||
vllm_hf_token: null # Optional: Your Hugging Face token if model is gated
|
||||
vllm_cuda_visible_devices: "0" # set to e.g. "0,1,2,3" for multiple gpus
|
||||
vllm_gpu_memory_utilization: 0.85 # leave some room for the refusal classifier if you are using it (about 3gb)
|
||||
vllm_max_model_len: 4500
|
||||
vllm_dtype: "bfloat16"
|
||||
# Additional raw CLI arguments for vLLM server, e.g., ["--tensor-parallel-size", "4"] for multiple gpus
|
||||
vllm_extra_args: [] # each param as a separate string, e.g. ["--quantization", "bitsandbytes"]
|
||||
vllm_env: # env vars for the vLLM process
|
||||
# VLLM_USE_V1: "1" # may be needed for amd gpus
|
||||
|
||||
|
||||
################################################################################
|
||||
# GENERATION PARAMETERS (using antislop-vllm)
|
||||
################################################################################
|
||||
generation_step_enabled: true
|
||||
|
||||
# --- API & Model Configuration ---
|
||||
# If you set manage_vllm=true, leave the base url unset
|
||||
#generation_api_base_url: "http://localhost:8000/v1"
|
||||
#generation_api_base_url: "https://apjmbtwbrb8t61-8888.proxy.runpod.net/v1"
|
||||
generation_model_id: null # Model id for generation requests (if unset, uses model_id)
|
||||
generation_api_key: "xxx" # API key for the vLLM server
|
||||
|
||||
# --- Core Generation Settings ---
|
||||
generation_max_new_tokens: 1000
|
||||
generation_threads: 50 # Number of parallel threads for API queries in antislop-vllm. Note: vllm can become very inefficient if you go over some concurrency threshold (depending on vram)
|
||||
generation_max_prompts: 1200 # Number of samples to generate from the prompts in the dataset
|
||||
|
||||
# --- Dataset & Chat Template ---
|
||||
generation_hf_dataset_name: 'Nitral-AI/Reddit-SFW-Writing_Prompts_ShareGPT'
|
||||
generation_hf_dataset_split: 'train'
|
||||
# A huggingface model id or local dir containing the tokeniser you want to use to apply chat templates.
|
||||
# This is important if you are generating a ftpo dataset for later training.
|
||||
generation_chat_template_model_id: null # (if unset, uses model_id)
|
||||
|
||||
# --- Logging ---
|
||||
generation_logging_level: 'INFO' # Logging level for antislop-vllm script
|
||||
|
||||
# --- Advanced Sampling Parameters (passed as individual CLI args to antislop-vllm's generation_params) ---
|
||||
generation_param_chunk_size: 20
|
||||
generation_param_top_logprobs_count: 20
|
||||
generation_param_temperature: 1
|
||||
generation_param_top_p: 1.0
|
||||
generation_param_top_k: 50
|
||||
generation_param_min_p: 0.01
|
||||
generation_param_timeout: 480
|
||||
generation_param_stop_sequences: [] # e.g., ["\n\n", "---"]
|
||||
|
||||
# --- Prompting ---
|
||||
# The prompt template wraps the prompts when generating from a dataset.
|
||||
# To use the original prompt exactly, set the template to "{prompt}"
|
||||
generation_prompt_template: "Writing prompt: {prompt}\n\nWrite 1000 words to this prompt. Your response:\n"
|
||||
generation_system_prompt: "You are an uncensored writer." # optional; left empty → no system prompt
|
||||
|
||||
# --- Antislop Generation Features ---
|
||||
# generation_force_backtrack:
|
||||
# If set to true:
|
||||
# when resampling after backtracking, if we don't find a valid replacement token
|
||||
# we progressively disable sampling options (temp, then min_p, then top_p, then top_k)
|
||||
# until we find a non-banned replacement or run out of candidates.
|
||||
# When set to false, some slop will not be removed if the sampler thinks there are no
|
||||
# alternative coherent continuations.
|
||||
generation_force_backtrack: false
|
||||
|
||||
# --- N-gram Validator Settings (for antislop-vllm) ---
|
||||
# N-gram ban list file is managed by auto-antislop's iterative process.
|
||||
generation_ngram_remove_stopwords: true
|
||||
generation_ngram_language: "english"
|
||||
|
||||
# --- Refusal Detection ---
|
||||
# Detects refusals & doesn't include them in the training dataset. Uses about 3GB extra VRAM.
|
||||
generation_refusal_detection: true
|
||||
|
||||
################################################################################
|
||||
# N-GRAM ANALYSIS & BANNING (within auto-antislop)
|
||||
################################################################################
|
||||
enable_ngram_ban: true
|
||||
min_word_len_for_analysis: 3 # Filters out words under this length in n-gram analysis
|
||||
|
||||
# --- N-gram Identification Thresholds ---
|
||||
top_k_bigrams: 5000
|
||||
top_k_trigrams: 5000
|
||||
|
||||
# --- N-gram Banning Quotas (per iteration) ---
|
||||
# Bigrams
|
||||
dict_bigrams_initial: 400 # How many of the top over-represented dictionary bigrams to
|
||||
# ban in the first antislop iteration.
|
||||
# "Dictionary" means the bigrams were also found in the human
|
||||
# writing corpus.
|
||||
dict_bigrams_subsequent: 70 # How many to ban in each subsequent iteration
|
||||
nodict_bigrams_initial: 800 # "Nodict" here means the n-grams were not found at all in the
|
||||
# human corpus.
|
||||
nodict_bigrams_subsequent: 100
|
||||
# Trigrams
|
||||
dict_trigrams_initial: 300
|
||||
dict_trigrams_subsequent: 50
|
||||
nodict_trigrams_initial: 800
|
||||
nodict_trigrams_subsequent: 100
|
||||
|
||||
# --- User-Defined N-gram Bans ---
|
||||
# User-supplied extra n-grams to always ban (processed by auto-antislop)
|
||||
extra_ngrams_to_ban: [
|
||||
# "voice barely whisper",
|
||||
]
|
||||
|
||||
################################################################################
|
||||
# OVER-REPRESENTED WORD ANALYSIS & BANNING
|
||||
################################################################################
|
||||
compute_overrep_words: true
|
||||
top_k_words_for_overrep_analysis: 200000
|
||||
|
||||
# --- Quotas for Adding Over-represented Words to Slop Phrase Ban List ---
|
||||
dict_overrep_initial: 800 # How many of the top over-represented dictionary words to
|
||||
# ban in the first antislop iteration.
|
||||
# "Dictionary" means the words were also found in the human
|
||||
# writing corpus.
|
||||
dict_overrep_subsequent: 200 # How many to ban in each subsequent iteration
|
||||
nodict_overrep_initial: 80 # "Nodict" here means the n-grams were not found at all in the
|
||||
# human corpus.
|
||||
nodict_overrep_subsequent: 20
|
||||
|
||||
################################################################################
|
||||
# SLOP PHRASE BANNING
|
||||
################################################################################
|
||||
|
||||
# Slop phrases are over-represented whole phrases extracted from the generated texts.
|
||||
enable_slop_phrase_ban: true
|
||||
min_phrase_freq_to_keep: 2 # Min frequency for a new phrase from slop-forensics to be considered
|
||||
top_n_initial_slop_ban: 600 # New slop phrases from slop-forensics to ban in iter 0
|
||||
top_n_subsequent_slop_ban: 100 # New slop phrases from slop-forensics to ban in later iters
|
||||
|
||||
# --- User-Defined Slop Phrase Bans ---
|
||||
# User supplied list of strings to always ban
|
||||
# - case insensitive
|
||||
# To trigger a ban, the sequence must not have a word-like character
|
||||
# (not punctuation or whitespace) directly on either side. That is to say, we
|
||||
# are not banning disallowed sequences that occur as substrings in longer
|
||||
# words. The exception is if the banned string is already bookended by
|
||||
# a non-word character.
|
||||
#
|
||||
# Examples:
|
||||
# banned string "cat"
|
||||
# - won't trigger a ban for "cation"
|
||||
# - will trigger a ban on "cat[morecat]"
|
||||
# banned string "cat["
|
||||
# - *will* trigger a ban on "cat[morecat]", because the banned string
|
||||
# ends with a non-word character.
|
||||
extra_slop_phrases_to_ban: [
|
||||
"…", "...", "rain", "tapestry", "static", "regret", "rust"
|
||||
]
|
||||
|
||||
# --- Whitelisted Strings ---
|
||||
# These will be excluded from the list of slop strings that the pipeline finds.
|
||||
# Note: special tokens in the tokenizer and parts of the chat template are
|
||||
# automatically whitelisted.
|
||||
whitelist_strings: [
|
||||
# "think", "thinking"
|
||||
]
|
||||
|
||||
################################################################################
|
||||
# REGEX BANNING
|
||||
################################################################################
|
||||
# User-supplied regex patterns to ban
|
||||
# Note: unoptimised regex patterns can slow down antislop generation, as they will be called often on large texts.
|
||||
extra_regex_patterns: [
|
||||
# These ones ban "it's not x, it's y" type patterns:
|
||||
|
||||
#"\\b(?:\\w+n(?:['’]t)|not\\s+(?:just|only|merely|because))\\s+(?:(?![.;:?!…]).){1,100}?[.;:?!…]\\s*(?:it|they|you)(?:['’](?:s|re|m))?\\b(?!\\s+(?:was|were|is|are|wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t)\\b)(?:\\s*[*…]?\\s*)?(?!when\\b|then\\b|but\\b|and\\b|yet\\b)(?!right\\b)(?!normal\\b)(?!true\\b)(?!sure\\b)(?!only\\b)(?!still\\b)(?!rarely\\b)(?!already\\b)(?!wrong\\b)(?!want\\b)(?!just\\b)(?!couldn\\b)(?!could\\b)(?!saw\\b)(?!started\\b)(?!remember\\b)(?!struggled\\b)(?!watched\\b)(?!goal\\b)(?!took\\b)(?!kept\\b)(?!reminded\\b)(?!time\\b)(?!have\\b)(?!acted\\b)(?!smiled\\b)(?!think\\b)(?!give\\b)(?!grab\\b)(?!gave\\b)(?!turn\\b)(?!justify\\b)(?!\\w+ly\\b)(?=[a-z]{4,}\\b)[a-z]+\\w*",
|
||||
|
||||
#"\\b(?:\\w+n(?:['’]t)|not)\\s+(?:just|only|merely)?\\s*(?:(?![-–—]|[.?!…]).){1,80}?[-–—]{1,2}\\s*\\w+(?:['’]\\w+)?\\s+",
|
||||
|
||||
#"\\b(?:wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t|not)\\s+(?!\\b(?:minute|minutes|hour|hours|day|days|year|years|second|seconds)\\b)(?!with\\b)(?!even\\b)(?:(?![.;:?!…]).){2,120}?[.;:?!…]\\s*(?:it|they|you|that)(?:\\s+(?:was|were|is|are)\\b(?:\\s+[*_~]?\\w+[*_~]?)?|(?:['’](?:s|re|m))\\b(?:\\s+[*_~]?\\w+[*_~]?)?)",
|
||||
|
||||
#"\\bno\\s+longer\\s+(?:just|only|merely)?\\s+[^.;:?!…]{1,120}[.;:?!…]\\s*(?:it|they|you)\\s+(?:is|are|was|were)\\b(?:\\s+[*_~]?\\w+[*_~]?)?",
|
||||
|
||||
#"\\b(?:wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t|not)\\s+(?:just|only|merely)?\\s*(?:(?!\\bbut\\b|[.?!…]).){1,80}?[,;:\\-–—]\\s*but\\s+(?!I\\b)(?:also\\s+)?"
|
||||
|
||||
]
|
||||
|
||||
################################################################################
|
||||
# FINETUNING
|
||||
################################################################################
|
||||
finetune_enabled: true
|
||||
|
||||
# --- General Finetuning Setup ---
|
||||
finetune_use_unsloth: true
|
||||
finetune_mode: "ftpo" # dpo / ftpo (final token preference optimisation)
|
||||
finetune_ftpo_dataset: "" # you can specify an existing ftpo dataset, or leave unset to let the
|
||||
# pipeline use the one produced in the generation step
|
||||
finetune_base_model_id: null # Base model for DPO (if unset, uses model_id)
|
||||
finetune_max_seq_length: 4000 # this may truncate some outputs
|
||||
finetune_load_in_4bit: true # qlora
|
||||
|
||||
# --- Early Stopping ---
|
||||
finetune_early_stopping_wins: 0.85 # Early stopping threshold for fraction of *chosen* completions that are selected over *rejected*.
|
||||
# More than 0.85 may be overtrained. Set to > 1.0 to disable early stopping.
|
||||
finetune_early_stopping_loss: null # Loss threshold for early stopping. Set to null to disable.
|
||||
|
||||
# --- LoRA Configuration ---
|
||||
finetune_lora_r: 256 # the ftpo trainer works best with a high lora rank
|
||||
finetune_lora_alpha: 256
|
||||
finetune_lora_dropout: 0.05
|
||||
finetune_weight_decay: 0.01
|
||||
finetune_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "lm_head"]
|
||||
|
||||
# --- Layer Freezing ---
|
||||
finetune_freeze_early_layers: true
|
||||
finetune_n_layers_unfrozen: 10
|
||||
|
||||
# --- Training Process ---
|
||||
finetune_gradient_checkpointing: "unsloth"
|
||||
finetune_chat_template: "" # e.g. "gemma-3" -- get the chat template from unsloth's helper if required, otherwise leave the string blank to use the tokeniser's chat template
|
||||
finetune_batch_size: 1
|
||||
finetune_gradient_accumulation_steps: 16
|
||||
finetune_warmup_ratio: 0.1
|
||||
finetune_num_epochs: 1
|
||||
|
||||
# --- Learning Rate ---
|
||||
finetune_learning_rate: 0.000001
|
||||
finetune_auto_learning_rate: true # true: automatically determine learning rate based on dataset size, effective batch size & lora rank
|
||||
finetune_auto_learning_rate_adjustment_scaling: 0.1 # scale the auto-lr by this factor
|
||||
|
||||
# --- DPO/FTPO Specific ---
|
||||
finetune_beta: 0.1 # DPO beta
|
||||
|
||||
# --- Output & Saving ---
|
||||
finetune_output_dir_suffix: "_ftpo_exp01" # Appended to experiment run dir
|
||||
finetune_save_merged_16bit: true
|
||||
finetune_save_gguf_q8_0: false
|
||||
|
||||
# --- Dataset Handling for Finetuning ---
|
||||
finetune_max_train_examples: 12000 # adjust as needed
|
||||
finetune_shuffle_seed: 666
|
||||
|
||||
# --- FTPO Sample Regularization ---
|
||||
# 0 = off; 0.9 strongly downsamples overrepresented rule violations
|
||||
# (this is useful because the raw generated dataset is typically very skewed)
|
||||
ftpo_sample_rejected_regularisation_strength: 0.8
|
||||
ftpo_sample_chosen_regularisation_strength: 0.2
|
||||
ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list
|
||||
|
||||
|
||||
# ── FTPO-specific hyper-parameters ─────────────────────────────────────────
|
||||
# Leave any of these out (or set to null) to fall back to FTPOTrainer defaults.
|
||||
|
||||
# Loss terms are computed separately for the target (chosen + rejected) tokens vs the remainder of the vocab.
|
||||
# This is because we want to allow more freedom of movement for the target tokens.
|
||||
|
||||
# MSE loss term 1: light mse loss applied tokenwise on target tokens
|
||||
ftpo_lambda_mse_target: 0.05 # Strength of MSE loss tether on the individual logits in the
|
||||
# chosen+rejected set vs reference.
|
||||
ftpo_tau_mse_target: 0.5 # Grace bandwidth (logits) before the above MSE loss kicks in.
|
||||
|
||||
# MSE loss term 2: stronger mse term applied to remaining (non-target) vocab
|
||||
ftpo_lambda_mse: 0.4
|
||||
|
||||
ftpo_clip_epsilon_logits: 2 # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off"
|
||||
0
core/__init__.py
Normal file
0
core/__init__.py
Normal file
470
core/analysis.py
Normal file
470
core/analysis.py
Normal file
@@ -0,0 +1,470 @@
|
||||
import json
|
||||
import math
|
||||
import re
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from collections import Counter
|
||||
import pandas as pd
|
||||
import nltk
|
||||
import numpy as np
|
||||
from typing import List, Tuple, Dict, Optional
|
||||
import unicodedata
|
||||
|
||||
# Assuming slop-forensics is in sys.path via main.py
|
||||
from slop_forensics import config as _sf_cfg # For SLOP_PHRASES_NGRAM_SIZE etc.
|
||||
from slop_forensics.analysis import (
|
||||
get_word_counts, filter_mostly_numeric, merge_plural_possessive_s,
|
||||
filter_stopwords, filter_common_words, analyze_word_rarity,
|
||||
find_over_represented_words
|
||||
)
|
||||
from slop_forensics.utils import normalize_text as normalise_keep_marks
|
||||
from slop_forensics.utils import extract_words
|
||||
# from slop_forensics.utils import load_jsonl_file, normalize_text, extract_words # Using local versions for now
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Local implementations of utils if slop_forensics.utils is problematic
|
||||
# These should ideally come from the submodule if its structure allows easy import
|
||||
def local_load_jsonl_file(file_path_str: str):
|
||||
data = []
|
||||
with open(file_path_str, 'r', encoding='utf-8') as f:
|
||||
for line in f:
|
||||
try:
|
||||
data.append(json.loads(line))
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Skipping malformed JSON line in {file_path_str}: {line.strip()}")
|
||||
return data
|
||||
|
||||
def local_normalize_text(text: str) -> str:
|
||||
return normalise_keep_marks(text)
|
||||
|
||||
def local_extract_words(normalized_text: str, min_len: int):
|
||||
return [
|
||||
w for w in extract_words(normalized_text)
|
||||
if len(w) >= min_len or "'" in w
|
||||
]
|
||||
|
||||
def _token_is_letter_or_mark(token: str) -> bool:
|
||||
"""
|
||||
True if every code-point in *token* is either a Unicode Letter (L*)
|
||||
or Mark (M*). Apostrophes / hyphens are not allowed here because
|
||||
`normalise_keep_marks` has already removed them.
|
||||
"""
|
||||
for ch in token:
|
||||
if unicodedata.category(ch)[0] not in ("L", "M"):
|
||||
return False
|
||||
return True
|
||||
|
||||
# --- Over-Represented Word Analysis ---
|
||||
BOOST_EXPONENT = 0.75
|
||||
ATTEN_EXPONENT = 0.75
|
||||
|
||||
def build_overrep_word_csv(
|
||||
texts: List[str],
|
||||
out_csv: Path,
|
||||
top_n_words_analysis: int,
|
||||
stop_words_set: Optional[set] = None, # keeps the caller happy
|
||||
) -> Tuple[pd.DataFrame, List[str], List[str]]:
|
||||
"""
|
||||
Notebook-faithful implementation that ALSO accepts *stop_words_set* so the
|
||||
CLI call `build_overrep_word_csv(..., stop_words_set=…)` keeps working.
|
||||
Returns (df, dict_words, nodict_words).
|
||||
"""
|
||||
# ------------------------------------------------- plain-file logging ---
|
||||
log_path = out_csv.with_suffix(".log")
|
||||
|
||||
def _log(msg: str) -> None:
|
||||
with log_path.open("a", encoding="utf-8") as fh:
|
||||
fh.write(f"{datetime.datetime.now():%Y-%m-%d %H:%M:%S} {msg}\n")
|
||||
|
||||
import datetime, traceback
|
||||
_log(f"build_overrep_word_csv ▶ {len(texts)} input texts")
|
||||
|
||||
try:
|
||||
# ---------- flatten + count (identical to the notebook) -------------
|
||||
counts = get_word_counts(texts) # ← **no extra kwargs**
|
||||
_log(f"after get_word_counts: {len(counts)} types")
|
||||
|
||||
counts = filter_mostly_numeric(counts)
|
||||
counts = merge_plural_possessive_s(counts)
|
||||
counts = filter_stopwords(counts)
|
||||
|
||||
_log(f"after filters: {len(counts)} types")
|
||||
|
||||
# ---------- rarity + over-rep score ---------------------------------
|
||||
corpus_freqs, wf_freqs, *_ = analyze_word_rarity(counts)
|
||||
overrep = find_over_represented_words(
|
||||
corpus_freqs, wf_freqs, top_n=top_n_words_analysis
|
||||
)
|
||||
_log(f"find_over_represented_words → {len(overrep)} rows")
|
||||
|
||||
# ---------- DataFrame ----------------------------------------------
|
||||
df = pd.DataFrame(
|
||||
overrep,
|
||||
columns=[
|
||||
"word", "ratio_corpus/wordfreq", "corpus_freq", "wordfreq_freq"
|
||||
],
|
||||
)
|
||||
num_cols = ["ratio_corpus/wordfreq", "corpus_freq", "wordfreq_freq"]
|
||||
df[num_cols] = df[num_cols].apply(pd.to_numeric, errors="coerce")
|
||||
|
||||
# ---------- modulated_score for dictionary words --------------------
|
||||
dict_mask = df["wordfreq_freq"] > 0
|
||||
if dict_mask.any():
|
||||
df_dict = df[dict_mask].copy()
|
||||
boost = np.power(df_dict["corpus_freq"], BOOST_EXPONENT)
|
||||
atten = np.power(df_dict["wordfreq_freq"], ATTEN_EXPONENT)
|
||||
atten_safe = np.where(atten == 0, 1, atten)
|
||||
df.loc[dict_mask, "modulated_score"] = (
|
||||
df_dict["ratio_corpus/wordfreq"] * boost / atten_safe
|
||||
)
|
||||
|
||||
# ---------- write CSV ----------------------------------------------
|
||||
df.to_csv(out_csv, index=False)
|
||||
_log(f"CSV written → {out_csv} ({len(df)} rows)")
|
||||
|
||||
# ---------- split & sort -------------------------------------------
|
||||
dict_words_df = df[dict_mask]
|
||||
dict_words = (
|
||||
dict_words_df.sort_values(
|
||||
"modulated_score", ascending=False)["word"].tolist()
|
||||
if "modulated_score" in dict_words_df.columns
|
||||
else dict_words_df["word"].tolist()
|
||||
)
|
||||
nodict_words = df[~dict_mask]["word"].tolist()
|
||||
_log(f"returning {len(dict_words)} dict words, {len(nodict_words)} non-dict")
|
||||
|
||||
return df, dict_words, nodict_words
|
||||
|
||||
except Exception as exc:
|
||||
_log("ERROR:\n" + "".join(traceback.format_exception(type(exc), exc, exc.__traceback__)))
|
||||
raise
|
||||
|
||||
|
||||
def select_overrep_words_for_ban(dict_words: list[str],
|
||||
nodict_words: list[str],
|
||||
is_first_iteration: bool,
|
||||
config: dict,
|
||||
*,
|
||||
whitelist: set[str]) -> list[str]:
|
||||
if is_first_iteration:
|
||||
dict_q, nodict_q = config['dict_overrep_initial'], config['nodict_overrep_initial']
|
||||
else:
|
||||
dict_q, nodict_q = config['dict_overrep_subsequent'], config['nodict_overrep_subsequent']
|
||||
|
||||
selected = []
|
||||
for w in dict_words:
|
||||
if len(selected) >= dict_q: break
|
||||
if w.lower() not in whitelist: selected.append(w)
|
||||
for w in nodict_words:
|
||||
if len(selected) >= dict_q + nodict_q: break
|
||||
if w.lower() not in whitelist: selected.append(w)
|
||||
logger.info(f"Selected {len(selected)} over-rep words for ban ({dict_q}/{nodict_q} quotas).")
|
||||
return selected
|
||||
|
||||
|
||||
# --- Slop Phrase Banning ---
|
||||
def update_banned_slop_phrases(
|
||||
json_path: Path,
|
||||
texts: list[str],
|
||||
how_many_new: int,
|
||||
tmp_dir: Path,
|
||||
config: dict,
|
||||
*,
|
||||
whitelist: set[str],
|
||||
over_represented_words: Optional[list[str]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Unchanged logic EXCEPT: any candidate phrase that contains a
|
||||
*whitelisted word (case-insensitive)* is skipped, and the final
|
||||
merged list drops any legacy items that are now whitelisted.
|
||||
"""
|
||||
logger.info(f"Updating slop-phrase ban list ({json_path.name}) …")
|
||||
|
||||
def is_whitelisted(phrase: str) -> bool:
|
||||
return any(w == phrase.lower() for w in whitelist)
|
||||
|
||||
# --------------------------------------------------------------- #
|
||||
# 1. run extractor (identical to previous body) #
|
||||
# --------------------------------------------------------------- #
|
||||
from slop_forensics.slop_lists import extract_and_save_slop_phrases as _extract
|
||||
from slop_forensics import config as _sf_cfg
|
||||
tmp_dir.mkdir(parents=True, exist_ok=True)
|
||||
_extract(
|
||||
texts=texts,
|
||||
output_dir=tmp_dir,
|
||||
n=_sf_cfg.SLOP_PHRASES_NGRAM_SIZE,
|
||||
top_k_ngrams=max(1000, how_many_new * 5),
|
||||
top_phrases_to_save=max(how_many_new * 3, 100),
|
||||
chunksize=_sf_cfg.SLOP_PHRASES_CHUNKSIZE,
|
||||
)
|
||||
|
||||
cand_phrases: List[str] = []
|
||||
try:
|
||||
with (tmp_dir / "slop_list_phrases.jsonl").open(encoding="utf-8") as fh:
|
||||
for line in fh:
|
||||
item = json.loads(line)
|
||||
phrase = item[0] if isinstance(item, list) else str(item)
|
||||
if phrase and not is_whitelisted(phrase):
|
||||
cand_phrases.append(phrase)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
# --------------------------------------------------------------- #
|
||||
# 2. merge with existing #
|
||||
# --------------------------------------------------------------- #
|
||||
existing: set[str] = set()
|
||||
if json_path.exists():
|
||||
try:
|
||||
for entry in json.loads(json_path.read_text("utf-8")):
|
||||
p = entry[0] if isinstance(entry, list) else str(entry)
|
||||
if p and not is_whitelisted(p):
|
||||
existing.add(p)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# keep requested quota only
|
||||
cand_phrases = cand_phrases[:how_many_new]
|
||||
if over_represented_words:
|
||||
for w in over_represented_words:
|
||||
if w not in whitelist:
|
||||
existing.add(w)
|
||||
|
||||
merged = sorted((existing | set(cand_phrases)) - whitelist)
|
||||
json_path.write_text(json.dumps([[p, 1] for p in merged], indent=2, ensure_ascii=False), "utf-8")
|
||||
logger.info(f"🚫 Slop-phrase ban list now {len(merged)} entries "
|
||||
f"(+{len(merged)-len(existing)} this iter)")
|
||||
|
||||
|
||||
# --- N-Gram Analysis ---
|
||||
def _convert_and_normalize_human_ngram_list(ngram_list_of_dicts, n_value: int):
|
||||
if not isinstance(ngram_list_of_dicts, list): return {}
|
||||
converted_dict = {}
|
||||
for item in ngram_list_of_dicts:
|
||||
if not isinstance(item, dict): continue
|
||||
ngram_str, frequency = item.get("ngram"), item.get("frequency")
|
||||
if ngram_str is None or frequency is None: continue
|
||||
try: freq_int = int(frequency)
|
||||
except ValueError: continue
|
||||
|
||||
tokens = [
|
||||
t.lower()
|
||||
for t in nltk.word_tokenize(local_normalize_text(str(ngram_str)))
|
||||
if _token_is_letter_or_mark(t)
|
||||
]
|
||||
if len(tokens) == n_value:
|
||||
processed_ngram_key = " ".join(tokens)
|
||||
if processed_ngram_key:
|
||||
converted_dict[processed_ngram_key] = converted_dict.get(processed_ngram_key, 0) + freq_int
|
||||
return converted_dict
|
||||
|
||||
def norm_per_freq_denom(raw_count: int, char_total: float, freq_norm_denom: int) -> float:
|
||||
if char_total == 0: return 0.0 if raw_count == 0 else math.inf
|
||||
return (raw_count / char_total) * freq_norm_denom
|
||||
|
||||
def build_norm_dict(counter: Counter, char_total: float, top_k: int, freq_norm_denom: int):
|
||||
char_total_float = float(char_total)
|
||||
return {
|
||||
term: {"gen_count": counter[term], "gen_freq_norm": norm_per_freq_denom(counter[term], char_total_float, freq_norm_denom)}
|
||||
for term, _ in counter.most_common(top_k) if term
|
||||
}
|
||||
|
||||
def compare_to_human(gen_norm: dict, human_counts: dict, human_total_chars: float, freq_norm_denom: int):
|
||||
both, gen_only = {}, {}
|
||||
human_total_chars_float = float(human_total_chars)
|
||||
for term, data in gen_norm.items():
|
||||
if not term: continue
|
||||
if term in human_counts:
|
||||
h_raw_count = human_counts[term]
|
||||
h_freq_norm = norm_per_freq_denom(h_raw_count, human_total_chars_float, freq_norm_denom)
|
||||
gen_freq = data["gen_freq_norm"]
|
||||
ratio = math.inf if h_freq_norm == 0 and gen_freq > 0 else (gen_freq / h_freq_norm if h_freq_norm > 0 else (1.0 if gen_freq == 0 else 0.0) )
|
||||
both[term] = {**data, "human_count": h_raw_count, "human_freq_norm": h_freq_norm, "freq_ratio_gen/hu": ratio}
|
||||
else:
|
||||
gen_only[term] = {**data, "human_count": 0, "human_freq_norm": 0.0, "freq_ratio_gen/hu": math.inf if data["gen_freq_norm"] > 0 else 0.0}
|
||||
return both, gen_only
|
||||
|
||||
def _is_refusal(rec: dict) -> bool:
|
||||
"""
|
||||
Returns True if this JSONL record represents a refused / skipped prompt.
|
||||
Recognises all variants produced by main.py / auto_unslop.py.
|
||||
"""
|
||||
if rec.get("refusal_detected") is True:
|
||||
return True
|
||||
status = rec.get("status", "").lower()
|
||||
if status in {"skipped"}:
|
||||
return True
|
||||
if status == "failed" and isinstance(rec.get("error"), str):
|
||||
err = rec["error"].lower()
|
||||
if err.startswith("refusal detected") or err.startswith("skipped -- prior refusal"):
|
||||
return True
|
||||
return False
|
||||
|
||||
def analyze_iteration_outputs_core(
|
||||
generated_jsonl_path: Path, human_profile_full: dict,
|
||||
iter_analysis_output_dir: Path, config: dict, stop_words_set: set
|
||||
):
|
||||
logger.info(f"--- Analyzing Outputs for {generated_jsonl_path.name} ---")
|
||||
iter_analysis_output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
gen_rows = local_load_jsonl_file(str(generated_jsonl_path))
|
||||
|
||||
gen_texts: List[str] = [
|
||||
rec["generation"]
|
||||
for rec in gen_rows
|
||||
if isinstance(rec, dict)
|
||||
and isinstance(rec.get("generation"), str)
|
||||
and rec["generation"].strip() # non-empty
|
||||
and not _is_refusal(rec) # ⬅️ new guard
|
||||
]
|
||||
|
||||
if not gen_texts:
|
||||
logger.warning(f"No usable text in {generated_jsonl_path}. Skipping analysis.")
|
||||
return None, None, None, None, [], 0
|
||||
|
||||
human_profile = human_profile_full.get('human-authored', human_profile_full.get(next(iter(human_profile_full), None))) # Robust key access
|
||||
if not human_profile: raise ValueError("Human profile data not found in JSON.")
|
||||
|
||||
human_bigrams = _convert_and_normalize_human_ngram_list(human_profile.get("top_bigrams", []), 2)
|
||||
human_trigrams = _convert_and_normalize_human_ngram_list(human_profile.get("top_trigrams", []), 3)
|
||||
|
||||
h_chars_total = float(
|
||||
human_profile.get("total_chars")
|
||||
or human_profile.get("chars_total")
|
||||
or (human_profile.get("num_texts_analyzed", 0) * human_profile.get("avg_length", 0))
|
||||
)
|
||||
if h_chars_total == 0:
|
||||
logger.warning("Human total characters is 0. Ratios might be infinite.")
|
||||
|
||||
|
||||
total_chars = sum(len(txt) for txt in gen_texts)
|
||||
bigram_counter, trigram_counter = Counter(), Counter()
|
||||
|
||||
min_word_len = config['min_word_len_for_analysis']
|
||||
|
||||
for txt in gen_texts:
|
||||
tokens_raw = nltk.word_tokenize(local_normalize_text(txt))
|
||||
tokens = [
|
||||
t.lower()
|
||||
for t in tokens_raw
|
||||
if _token_is_letter_or_mark(t)
|
||||
]
|
||||
tokens_filtered = [
|
||||
tok for tok in tokens
|
||||
if tok not in stop_words_set and (len(tok) >= min_word_len or tok in {"it's", "i'm"})
|
||||
]
|
||||
bigram_counter.update(" ".join(bg) for bg in nltk.ngrams(tokens_filtered, 2) if all(bg))
|
||||
trigram_counter.update(" ".join(tg) for tg in nltk.ngrams(tokens_filtered, 3) if all(tg))
|
||||
|
||||
freq_norm_denom = config.get('freq_norm_denom_for_analysis', 100000)
|
||||
gen_bigrams_norm = build_norm_dict(bigram_counter, float(total_chars), config['top_k_bigrams'], freq_norm_denom)
|
||||
gen_trigrams_norm = build_norm_dict(trigram_counter, float(total_chars), config['top_k_trigrams'], freq_norm_denom)
|
||||
|
||||
bigrams_dict, bigrams_nondict = compare_to_human(gen_bigrams_norm, human_bigrams, h_chars_total, freq_norm_denom)
|
||||
trigrams_dict, trigrams_nondict = compare_to_human(gen_trigrams_norm, human_trigrams, h_chars_total, freq_norm_denom)
|
||||
|
||||
df_bi_dict = pd.DataFrame.from_dict(bigrams_dict, orient="index").rename_axis('ngram').reset_index()
|
||||
df_bi_nondct = pd.DataFrame.from_dict(bigrams_nondict, orient="index").rename_axis('ngram').reset_index()
|
||||
df_tri_dict = pd.DataFrame.from_dict(trigrams_dict, orient="index").rename_axis('ngram').reset_index()
|
||||
df_tri_nondct = pd.DataFrame.from_dict(trigrams_nondict, orient="index").rename_axis('ngram').reset_index()
|
||||
|
||||
for df, sort_col in [(df_bi_dict, "freq_ratio_gen/hu"), (df_tri_dict, "freq_ratio_gen/hu")]:
|
||||
if not df.empty and sort_col in df.columns: df.sort_values(by=sort_col, ascending=False, inplace=True)
|
||||
for df, sort_col in [(df_bi_nondct, "gen_freq_norm"), (df_tri_nondct, "gen_freq_norm")]:
|
||||
if not df.empty and sort_col in df.columns: df.sort_values(by=sort_col, ascending=False, inplace=True)
|
||||
|
||||
df_bi_dict.to_csv(iter_analysis_output_dir / "bigrams__dictionary_sorted.csv", index=False)
|
||||
df_bi_nondct.to_csv(iter_analysis_output_dir / "bigrams__non_dictionary_sorted.csv", index=False)
|
||||
df_tri_dict.to_csv(iter_analysis_output_dir / "trigrams__dictionary_sorted.csv", index=False)
|
||||
df_tri_nondct.to_csv(iter_analysis_output_dir / "trigrams__non_dictionary_sorted.csv", index=False)
|
||||
logger.info(f"N-gram analysis CSVs written to {iter_analysis_output_dir.resolve()}")
|
||||
|
||||
return df_bi_dict, df_bi_nondct, df_tri_dict, df_tri_nondct, gen_texts, total_chars
|
||||
|
||||
|
||||
def update_banned_ngrams_list(
|
||||
banned_ngrams_json_path: Path,
|
||||
dfs: list, # bi/tri, dict / non-dict
|
||||
is_first_iteration: bool,
|
||||
config: dict,
|
||||
*,
|
||||
whitelist: set[str],
|
||||
):
|
||||
newly: set[str] = set()
|
||||
def _take(df, n): # helper
|
||||
return {
|
||||
row["ngram"] for _, row in (df.head(n)).iterrows()
|
||||
if "ngram" in row and row["ngram"] and row["ngram"] not in whitelist
|
||||
} if df is not None and not df.empty and n > 0 else set()
|
||||
|
||||
if is_first_iteration:
|
||||
quotas = (
|
||||
config['dict_bigrams_initial'], config['nodict_bigrams_initial'],
|
||||
config['dict_trigrams_initial'], config['nodict_trigrams_initial'],
|
||||
)
|
||||
else:
|
||||
quotas = (
|
||||
config['dict_bigrams_subsequent'], config['nodict_bigrams_subsequent'],
|
||||
config['dict_trigrams_subsequent'], config['nodict_trigrams_subsequent'],
|
||||
)
|
||||
for df, q in zip(dfs, quotas):
|
||||
newly |= _take(df, q)
|
||||
|
||||
newly |= {s for s in config.get('extra_ngrams_to_ban', []) if s not in whitelist}
|
||||
|
||||
current = set()
|
||||
if banned_ngrams_json_path.exists():
|
||||
try:
|
||||
current = set(json.loads(banned_ngrams_json_path.read_text("utf-8")))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
final = sorted((current | newly) - whitelist)
|
||||
banned_ngrams_json_path.write_text(json.dumps(final, indent=2, ensure_ascii=False), "utf-8")
|
||||
logger.info(f"📄 N-gram ban list updated → {banned_ngrams_json_path} "
|
||||
f"(+{len(final)-len(current)} new, total {len(final)})")
|
||||
|
||||
|
||||
# --- Metrics Calculation ---
|
||||
def calculate_lexical_diversity_stats(gen_texts: list, min_word_len: int):
|
||||
if not gen_texts: return 0.0, 0.0
|
||||
all_words = []
|
||||
for text in gen_texts:
|
||||
tokens = [t.lower() for t in nltk.word_tokenize(local_normalize_text(text)) if t.isalpha() and (len(t) >= min_word_len or t in {"a", "i"})]
|
||||
all_words.extend(tokens)
|
||||
if not all_words: return 0.0, 0.0
|
||||
num_tokens, num_types = len(all_words), len(set(all_words))
|
||||
ttr = num_types / num_tokens if num_tokens > 0 else 0.0
|
||||
rttr = num_types / math.sqrt(num_tokens) if num_tokens > 0 else 0.0
|
||||
return ttr, rttr
|
||||
|
||||
def calculate_repetition_score(gen_texts: list, total_chars: int, iteration_dfs: list, config: dict, stop_words_set: set):
|
||||
if not gen_texts or total_chars == 0: return 0.0
|
||||
|
||||
target_ngrams = set()
|
||||
top_n_rep = config.get('top_n_repetition_stat', 50)
|
||||
min_word_len = config['min_word_len_for_analysis']
|
||||
freq_norm_denom = config.get('freq_norm_denom_for_analysis', 100000)
|
||||
|
||||
for df in iteration_dfs: # df_bi_dict, df_bi_nondct, df_tri_dict, df_tri_nondct
|
||||
if df is not None and not df.empty and 'ngram' in df.columns:
|
||||
target_ngrams.update(df.head(top_n_rep)['ngram'].tolist())
|
||||
if not target_ngrams: return 0.0
|
||||
|
||||
total_repetition_instances = 0
|
||||
for text in gen_texts:
|
||||
tokens_all = [
|
||||
t.lower()
|
||||
for t in nltk.word_tokenize(local_normalize_text(text))
|
||||
if _token_is_letter_or_mark(t)
|
||||
]
|
||||
tokens = [tok for tok in tokens_all if tok not in stop_words_set and (len(tok) >= min_word_len or tok in {"it's", "i'm"})]
|
||||
|
||||
current_bigrams = [" ".join(bg) for bg in nltk.ngrams(tokens, 2) if all(bg)]
|
||||
current_trigrams = [" ".join(tg) for tg in nltk.ngrams(tokens, 3) if all(tg)]
|
||||
for bg in current_bigrams:
|
||||
if bg in target_ngrams: total_repetition_instances += 1
|
||||
for tg in current_trigrams:
|
||||
if tg in target_ngrams: total_repetition_instances += 1
|
||||
|
||||
return norm_per_freq_denom(total_repetition_instances, float(total_chars), freq_norm_denom)
|
||||
105
core/dpo.py
Normal file
105
core/dpo.py
Normal file
@@ -0,0 +1,105 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
def create_dpo_dataset(
|
||||
iter0_jsonl: Path,
|
||||
final_iter_jsonl: Path,
|
||||
output_jsonl: Path,
|
||||
) -> None:
|
||||
logger.info(f"Creating DPO dataset from {iter0_jsonl.name} and {final_iter_jsonl.name} -> {output_jsonl.name}")
|
||||
|
||||
KEY_PROMPT = "prompt"
|
||||
KEY_GENERATION = "generation"
|
||||
KEY_PROMPT_ID = "prompt_id" # Assuming antislop-vllm output includes this
|
||||
|
||||
def _strip_wrapping(text: str) -> str:
|
||||
# This needs to match how prompts are formatted by antislop-vllm/main.py
|
||||
# If main.py's HF dataset loading adds "Writing prompt: ... Your response:\n", strip it.
|
||||
# For now, assume prompts in the JSONL are the "actual" prompts.
|
||||
# If antislop-vllm's output `prompt` field is already clean, this might not be needed.
|
||||
# The example in the notebook was:
|
||||
# prefix = "Writing prompt: "
|
||||
# if text.startswith(prefix): text = text[len(prefix):]
|
||||
# return text.strip()
|
||||
return text # Assuming prompt field in JSONL is already the core prompt
|
||||
|
||||
def _load_file(path: Path) -> dict[str, dict[str, str]]:
|
||||
out_data: dict[str, dict[str, str]] = {}
|
||||
if not path.exists():
|
||||
logger.warning(f"DPO source file not found: {path}")
|
||||
return out_data
|
||||
|
||||
with path.open(encoding="utf-8") as fh:
|
||||
for i, line_raw in enumerate(fh):
|
||||
try:
|
||||
row = json.loads(line_raw)
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Skipping malformed JSON line {i+1} in {path}")
|
||||
continue
|
||||
|
||||
if not isinstance(row, dict):
|
||||
logger.warning(f"Skipping non-dict row {i+1} in {path}")
|
||||
continue
|
||||
|
||||
prompt_raw = row.get(KEY_PROMPT)
|
||||
gen = row.get(KEY_GENERATION)
|
||||
prompt_id = row.get(KEY_PROMPT_ID) # Use prompt_id as the primary key
|
||||
|
||||
if not isinstance(prompt_raw, str) or not isinstance(gen, str) or prompt_id is None:
|
||||
# logger.debug(f"Skipping row {i+1} in {path} due to missing/invalid prompt, generation, or prompt_id.")
|
||||
continue
|
||||
|
||||
prompt_clean = _strip_wrapping(prompt_raw)
|
||||
key = str(prompt_id) # Use prompt_id as the key
|
||||
|
||||
if not key:
|
||||
logger.warning(f"Skipping row {i+1} in {path} due to empty key (prompt_id).")
|
||||
continue
|
||||
out_data[key] = {"prompt": prompt_clean, "generation": gen}
|
||||
return out_data
|
||||
|
||||
data_iter0 = _load_file(iter0_jsonl)
|
||||
data_final = _load_file(final_iter_jsonl)
|
||||
|
||||
if not data_iter0 or not data_final:
|
||||
logger.error("DPO dataset not created: one or both input files were empty or could not be loaded.")
|
||||
return
|
||||
|
||||
common_keys = data_iter0.keys() & data_final.keys()
|
||||
|
||||
if not common_keys:
|
||||
logger.warning("No overlapping prompt_ids between iteration-0 and final iteration; DPO dataset not written.")
|
||||
logger.warning(f"Iter0 keys: {len(data_iter0)}, Final keys: {len(data_final)}")
|
||||
return
|
||||
|
||||
count_written = 0
|
||||
with output_jsonl.open("w", encoding="utf-8") as out_fh:
|
||||
for key in common_keys:
|
||||
# Use the prompt from iter0 as canonical, assuming prompt_id ensures they are fundamentally the same.
|
||||
prompt_for_dpo = data_iter0[key]["prompt"]
|
||||
|
||||
# Sanity check: if prompts differ significantly despite same ID, log it.
|
||||
if prompt_for_dpo != data_final[key]["prompt"]:
|
||||
logger.debug(f"Prompt text mismatch for prompt_id '{key}'. Using iter0 prompt for DPO pair.")
|
||||
|
||||
rec = {
|
||||
"prompt": prompt_for_dpo,
|
||||
"chosen": data_final[key]["generation"],
|
||||
"rejected": data_iter0[key]["generation"],
|
||||
}
|
||||
# Ensure chosen and rejected are not identical
|
||||
if rec["chosen"] == rec["rejected"]:
|
||||
logger.debug(f"Skipping DPO pair for prompt_id '{key}' as chosen and rejected generations are identical.")
|
||||
continue
|
||||
|
||||
json.dump(rec, out_fh, ensure_ascii=False)
|
||||
out_fh.write("\n")
|
||||
count_written +=1
|
||||
|
||||
if count_written > 0:
|
||||
logger.info(f"📁 DPO dataset written -> {output_jsonl} ({count_written} prompt pairs from {len(common_keys)} common prompt_ids)")
|
||||
else:
|
||||
logger.warning(f"No DPO pairs written. Common prompt_ids found: {len(common_keys)}, but all might have had identical chosen/rejected texts.")
|
||||
141
core/dpo_trainer.py
Normal file
141
core/dpo_trainer.py
Normal file
@@ -0,0 +1,141 @@
|
||||
# core/dpo_trainer.py
|
||||
|
||||
from __future__ import annotations
|
||||
from typing import List, Optional
|
||||
|
||||
from core.finetuning import DPOTrainer, torch, F
|
||||
|
||||
|
||||
class DPOTrainerWithChoiceWin(DPOTrainer):
|
||||
"""
|
||||
Drop-in DPO trainer that adds a token-level `chosen_win` metric for
|
||||
dpo_final_token batches without altering DPO loss/behavior.
|
||||
|
||||
Metric (per-row):
|
||||
Given the prompt context, compare the model's next-token probabilities:
|
||||
win = 1{ log p(chosen_first | prompt) > log p(rejected_first | prompt) }
|
||||
chosen_win = mean over rows of win.
|
||||
|
||||
Expected batch keys (dpo_final_token):
|
||||
- chosen_input_ids, chosen_attention_mask
|
||||
- rejected_input_ids, rejected_attention_mask
|
||||
- prompt_input_ids (+ prompt_attention_mask) OR
|
||||
prompt_ids (+ attention_mask)
|
||||
"""
|
||||
|
||||
def compute_loss(self, model, inputs, return_outputs: bool = False, **kwargs):
|
||||
# 1) standard DPO loss (unchanged)
|
||||
loss = super().compute_loss(model, inputs, return_outputs=False, **kwargs)
|
||||
|
||||
# 2) metric (only if dpo_final_token fields are present)
|
||||
needed = {
|
||||
"chosen_input_ids", "chosen_attention_mask",
|
||||
"rejected_input_ids", "rejected_attention_mask",
|
||||
}
|
||||
if needed.issubset(inputs.keys()):
|
||||
chosen_win = self._metric_dpo_final_token(model, inputs)
|
||||
if chosen_win is not None:
|
||||
self.store_metrics({"chosen_win": chosen_win}, train_eval="train")
|
||||
|
||||
if return_outputs:
|
||||
return loss, {}
|
||||
return loss
|
||||
|
||||
# ----------------------------- helpers ---------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _first_real_index(mask_row: torch.Tensor) -> Optional[int]:
|
||||
"""Index of the first non-pad token according to an attention mask row."""
|
||||
nz = mask_row.nonzero(as_tuple=False)
|
||||
return int(nz[0].item()) if nz.numel() else None
|
||||
|
||||
# ---------------------- core metric (dpo_final_token) -------------------
|
||||
|
||||
def _metric_dpo_final_token(self, model, inputs) -> Optional[torch.Tensor]:
|
||||
device = next(model.parameters()).device
|
||||
tok = getattr(self, "tokenizer", getattr(self, "processing_class", None))
|
||||
pad_id = getattr(tok, "pad_token_id", 0)
|
||||
|
||||
# Continuations
|
||||
ch_ids = inputs["chosen_input_ids"].to(device) # [B, Lc] (continuation only)
|
||||
ch_am = inputs["chosen_attention_mask"].to(device) # [B, Lc]
|
||||
rj_ids = inputs["rejected_input_ids"].to(device) # [B, Lr]
|
||||
rj_am = inputs["rejected_attention_mask"].to(device) # [B, Lr]
|
||||
|
||||
# Prompt (prefer explicit prompt_input_ids → fall back to prompt_ids)
|
||||
pr_ids_key = "prompt_input_ids" if "prompt_input_ids" in inputs else (
|
||||
"prompt_ids" if "prompt_ids" in inputs else None
|
||||
)
|
||||
pr_am_key = "prompt_attention_mask" if "prompt_attention_mask" in inputs else (
|
||||
"attention_mask" if "attention_mask" in inputs else None
|
||||
)
|
||||
|
||||
if pr_ids_key is None:
|
||||
# No prompt in the batch → cannot compute the metric reliably
|
||||
return None
|
||||
|
||||
pr_ids_full = inputs[pr_ids_key].to(device) # [B, Lp]
|
||||
if pr_am_key in inputs:
|
||||
pr_am_full = inputs[pr_am_key].to(device) # [B, Lp]
|
||||
else:
|
||||
# derive attention mask from pad id
|
||||
pr_am_full = pr_ids_full.ne(pad_id).to(pr_ids_full.dtype)
|
||||
|
||||
B = ch_ids.size(0)
|
||||
|
||||
# Assemble per-row prompt and first continuation tokens
|
||||
prompts: List[torch.Tensor] = []
|
||||
chosen_first: List[int] = []
|
||||
rejected_first: List[int] = []
|
||||
last_idx: List[int] = []
|
||||
|
||||
for b in range(B):
|
||||
# first token of each continuation (continuations are typically 1 token + EOS)
|
||||
k_ch = self._first_real_index(ch_am[b])
|
||||
k_rj = self._first_real_index(rj_am[b])
|
||||
if k_ch is None or k_rj is None:
|
||||
continue
|
||||
|
||||
ch_first = int(ch_ids[b, k_ch].item())
|
||||
rj_first = int(rj_ids[b, k_rj].item())
|
||||
|
||||
# prompt (all real tokens)
|
||||
pr_mask = pr_am_full[b].bool()
|
||||
pr_seq = pr_ids_full[b][pr_mask] # 1D tensor with real prompt ids
|
||||
if pr_seq.numel() == 0:
|
||||
continue
|
||||
|
||||
prompts.append(pr_seq)
|
||||
chosen_first.append(ch_first)
|
||||
rejected_first.append(rj_first)
|
||||
last_idx.append(pr_seq.numel() - 1)
|
||||
|
||||
n = len(prompts)
|
||||
if n == 0:
|
||||
return None
|
||||
|
||||
# Right-pad prompts into a dense batch for one forward pass
|
||||
max_pr = max(p.numel() for p in prompts)
|
||||
prompt_ids = pr_ids_full.new_full((n, max_pr), pad_id)
|
||||
prompt_am = pr_am_full.new_zeros((n, max_pr))
|
||||
|
||||
for i, p in enumerate(prompts):
|
||||
Lp = p.numel()
|
||||
prompt_ids[i, :Lp] = p
|
||||
prompt_am [i, :Lp] = 1
|
||||
|
||||
last_idx_t = torch.tensor(last_idx, device=device, dtype=torch.long)
|
||||
chosen_tok = torch.tensor(chosen_first, device=device, dtype=torch.long)
|
||||
reject_tok = torch.tensor(rejected_first, device=device, dtype=torch.long)
|
||||
|
||||
# One no-grad forward on prompts; evaluate next-token distribution at last prompt token
|
||||
with torch.no_grad():
|
||||
out = model(prompt_ids, attention_mask=prompt_am, use_cache=False, return_dict=True)
|
||||
logits_last = out.logits[torch.arange(n, device=device), last_idx_t, :] # [n, V]
|
||||
logp_last = F.log_softmax(logits_last, dim=-1)
|
||||
|
||||
lp_good = logp_last.gather(1, chosen_tok.unsqueeze(1)).squeeze(1) # [n]
|
||||
lp_bad = logp_last.gather(1, reject_tok.unsqueeze(1)).squeeze(1) # [n]
|
||||
wins = (lp_good > lp_bad).float()
|
||||
|
||||
return wins.mean().detach()
|
||||
1054
core/finetuning.py
Normal file
1054
core/finetuning.py
Normal file
File diff suppressed because it is too large
Load Diff
332
core/ftpo_trainer.py
Normal file
332
core/ftpo_trainer.py
Normal file
@@ -0,0 +1,332 @@
|
||||
from core.finetuning import DPOTrainer, torch, pad_sequence, F
|
||||
|
||||
# ---------------------------------------------------------------
|
||||
# Add Adaptive Gradient Clipping to *every* parameter in-place
|
||||
# ---------------------------------------------------------------
|
||||
def attach_agc(model, clip: float = 0.01, eps: float = 1e-3):
|
||||
"""
|
||||
Registers a per-parameter hook that applies Brock et al.’s
|
||||
Adaptive Gradient Clipping:
|
||||
|
||||
||g||₂ > clip * (||θ||₂ + eps) → g ← g * (threshold / ||g||₂)
|
||||
|
||||
Works with params in fp32, bf16, or bitsandbytes int4.
|
||||
"""
|
||||
|
||||
def _agc_hook(grad, param):
|
||||
#print('agc hook')
|
||||
if grad is None:
|
||||
return grad
|
||||
param_norm = param.detach().norm() # ||θ||
|
||||
grad_norm = grad.norm() # ||g||
|
||||
max_norm = clip * (param_norm + eps)
|
||||
if grad_norm > max_norm:
|
||||
grad = grad * (max_norm / (grad_norm + 1e-6))
|
||||
return grad
|
||||
|
||||
for p in model.parameters():
|
||||
if p.requires_grad:
|
||||
p.register_hook(lambda g, p=p: _agc_hook(g, p))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------
|
||||
# Early-stopping on any logged scalar (loss, chosen_win, etc.)
|
||||
# ---------------------------------------------------------------------
|
||||
from transformers.trainer_callback import TrainerCallback
|
||||
|
||||
|
||||
class ThresholdStop(TrainerCallback):
|
||||
"""
|
||||
Stop training immediately when `monitor` crosses `threshold`.
|
||||
|
||||
If `higher_is_better` is True → stop when metric >= threshold.
|
||||
If False → stop when metric <= threshold.
|
||||
"""
|
||||
def __init__(self, monitor: str, threshold: float, higher_is_better: bool):
|
||||
self.monitor = monitor
|
||||
self.threshold = threshold
|
||||
self.higher_is_better = higher_is_better
|
||||
|
||||
def on_log(self, args, state, control, logs=None, **kwargs):
|
||||
if logs is None or self.monitor not in logs:
|
||||
return
|
||||
value = logs[self.monitor]
|
||||
stop = (value >= self.threshold) if self.higher_is_better else (value <= self.threshold)
|
||||
if stop:
|
||||
control.should_training_stop = True
|
||||
print(f"[ThresholdStop] {self.monitor}={value:.4f} "
|
||||
f"crossed {'≥' if self.higher_is_better else '≤'} "
|
||||
f"{self.threshold} – stopping.")
|
||||
|
||||
class EarlyStoppingByMetric(TrainerCallback):
|
||||
"""
|
||||
Stop training when a monitored metric has stopped improving.
|
||||
|
||||
Args
|
||||
----
|
||||
monitor: str
|
||||
Key that appears in the `logs` dict (e.g. "loss", "chosen_win").
|
||||
higher_is_better: bool
|
||||
True → metric should increase (e.g. chosen_win)
|
||||
False → metric should decrease (e.g. loss / pref_loss)
|
||||
patience: int
|
||||
How many *log events* with no improvement to wait before stopping.
|
||||
min_delta: float
|
||||
Minimum change that counts as an improvement.
|
||||
"""
|
||||
def __init__(self,
|
||||
monitor: str,
|
||||
higher_is_better: bool,
|
||||
patience: int = 10,
|
||||
min_delta: float = 0.0):
|
||||
self.monitor = monitor
|
||||
self.higher_is_better = higher_is_better
|
||||
self.patience = patience
|
||||
self.min_delta = min_delta
|
||||
self.best = None
|
||||
self.counter = 0 # events since last improv.
|
||||
|
||||
# ── invoked every time trainer logs metrics ─────────────────────
|
||||
def on_log(self, args, state, control, logs=None, **kwargs):
|
||||
if logs is None or self.monitor not in logs:
|
||||
return
|
||||
|
||||
current = logs[self.monitor]
|
||||
|
||||
# first observation
|
||||
if self.best is None:
|
||||
self.best = current
|
||||
return
|
||||
|
||||
# compute signed improvement
|
||||
if self.higher_is_better:
|
||||
improvement = current - self.best
|
||||
else:
|
||||
improvement = self.best - current
|
||||
|
||||
# has the metric improved “enough”?
|
||||
if improvement > self.min_delta:
|
||||
self.best = current
|
||||
self.counter = 0
|
||||
else:
|
||||
self.counter += 1
|
||||
if self.counter >= self.patience:
|
||||
# signal the Trainer to halt
|
||||
control.should_training_stop = True
|
||||
print(f"[EarlyStopping] '{self.monitor}' plateaued "
|
||||
f"(best={self.best:.5f}) – stopping training.")
|
||||
|
||||
|
||||
|
||||
class FTPOTrainer(DPOTrainer):
|
||||
"""
|
||||
Trainer for final token preference optimisation (ftpo).
|
||||
Replaces TRL’s standard loss with a log-ratio on the **last**
|
||||
autoregressive position.
|
||||
"""
|
||||
|
||||
# ──────────────────────────────────────────────────────────────
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.remove_unused_columns = False
|
||||
self.data_collator = self.ftpo_collator # override
|
||||
|
||||
# ──────────────────────────────────────────────────────────────
|
||||
@staticmethod
|
||||
def _get_proj(model):
|
||||
"""
|
||||
Return the output-projection module in a model-agnostic way.
|
||||
Falls back to `lm_head` if `get_output_embeddings()` is None.
|
||||
"""
|
||||
proj = model.get_output_embeddings()
|
||||
if proj is None:
|
||||
proj = getattr(model, "lm_head", None)
|
||||
if proj is None:
|
||||
raise AttributeError(
|
||||
"Model lacks both get_output_embeddings() and lm_head."
|
||||
)
|
||||
return proj
|
||||
|
||||
# ──────────────────────────────────────────────────────────────
|
||||
def ftpo_collator(self, features):
|
||||
"""
|
||||
Left-pads every prompt to `self.args.max_length`, so the last real
|
||||
token is always at position -1. That lets the loss read logits
|
||||
with a single slice ([:, -1, :]).
|
||||
"""
|
||||
pad_id = self.padding_value
|
||||
max_len = self.args.max_length
|
||||
batch_sz = len(features)
|
||||
|
||||
# ── build [B, L] prompt tensor ───────────────────────────────
|
||||
prompt_ids = torch.full((batch_sz, max_len), pad_id, dtype=torch.long)
|
||||
attention_ms = torch.zeros_like(prompt_ids, dtype=torch.bool)
|
||||
|
||||
for i, feat in enumerate(features):
|
||||
seq = torch.tensor(feat["prompt_ids"], dtype=torch.long)
|
||||
if seq.size(0) > max_len:
|
||||
seq = seq[-max_len:] # truncate left if over-long
|
||||
prompt_ids[i, -seq.size(0):] = seq # left-pad
|
||||
attention_ms[i, -seq.size(0):] = True
|
||||
|
||||
# ── universal fields ─────────────────────────────────────────
|
||||
batch = dict(
|
||||
prompt_ids = prompt_ids,
|
||||
attention_mask = attention_ms,
|
||||
rejected_token_id = torch.tensor([f["rejected_token_id"] for f in features]),
|
||||
)
|
||||
|
||||
# ── ftpo vs single-token branch ────────────────────────
|
||||
max_c = max(len(f["chosen_ids"]) for f in features)
|
||||
chosen_pad = torch.full((batch_sz, max_c), pad_id, dtype=torch.long)
|
||||
chosen_mask = torch.zeros_like(chosen_pad, dtype=torch.bool)
|
||||
for i, f in enumerate(features):
|
||||
ids = torch.tensor(f["chosen_ids"], dtype=torch.long)
|
||||
chosen_pad [i, :ids.size(0)] = ids
|
||||
chosen_mask[i, :ids.size(0)] = True
|
||||
batch.update(chosen_ids = chosen_pad,
|
||||
chosen_mask = chosen_mask)
|
||||
|
||||
return batch
|
||||
|
||||
def compute_loss(self, model, inputs, return_outputs=False, **_):
|
||||
# We use 2 separate MSE loss terms (aggregate removed):
|
||||
|
||||
# 1. A lightly applied tokenwise MSE loss applied to only the target tokens
|
||||
lambda_mse_target = getattr(self, "lambda_mse_target", 0.05) # strength
|
||||
tau_mse_target = getattr(self, "tau_mse_target", 1.0) # grace region (zero cost movement)
|
||||
|
||||
# 2. A strongly applied tokenwise MSE loss applied to the remaining (non-target) vocab
|
||||
lambda_mse = getattr(self, "lambda_mse", 0.4) # how strongly the remaining vocab (other than chosen/rejected) is tethered to reference via mse loss
|
||||
|
||||
# loss contribution is clipped if (chosen - rejected) logits delta is above this
|
||||
clip_epsilon_logits = getattr(self, "clip_epsilon_logits", 2)
|
||||
|
||||
USE_MSE_LOSS=True # tether all the logits other than the ones we are interested in moving to the reference
|
||||
|
||||
# ── unpack ---------------------------------------------------------
|
||||
device = next(model.parameters()).device # works for DP / DDP
|
||||
ids = inputs["prompt_ids"].to(device) # [B,L]
|
||||
attn = inputs["attention_mask"].to(device) # [B,L]
|
||||
B, L = ids.shape
|
||||
|
||||
seq_len = attn.sum(1)
|
||||
pad_off = (L - seq_len).unsqueeze(1)
|
||||
arange_L = torch.arange(L, device=ids.device).unsqueeze(0)
|
||||
pos_full = (arange_L - pad_off).clamp(min=0)
|
||||
pos_full = pos_full.masked_fill(attn == 0, 0)
|
||||
|
||||
outputs = model(
|
||||
ids,
|
||||
attention_mask=attn,
|
||||
position_ids=pos_full,
|
||||
use_cache=False,
|
||||
return_dict=True,
|
||||
)
|
||||
|
||||
logits_last = outputs.logits[:, -1, :] # [B, V]
|
||||
logp_all = F.log_softmax(logits_last, dim=-1) # [B, V]
|
||||
|
||||
ch_ids = inputs["chosen_ids"].to(device)
|
||||
ch_mask = inputs["chosen_mask"].to(device)
|
||||
rejected = inputs["rejected_token_id"].to(device)
|
||||
logp_bad = logp_all.gather(-1, rejected.unsqueeze(-1)).squeeze(-1)
|
||||
|
||||
batch_rows = torch.arange(B, device=logp_all.device).unsqueeze(1)
|
||||
gathered = logits_last[batch_rows, ch_ids]
|
||||
logit_bad = logits_last.gather(-1, rejected.unsqueeze(-1))
|
||||
margin = gathered - logit_bad
|
||||
weights = torch.clamp((clip_epsilon_logits - margin) / clip_epsilon_logits, 0.0, 1.0) * ch_mask
|
||||
|
||||
zero_row = weights.sum(dim=-1, keepdim=True) < 1e-12
|
||||
weights = torch.where(zero_row, ch_mask.float(), weights)
|
||||
|
||||
weights_sum = weights.sum(dim=-1, keepdim=True)
|
||||
batch_rows = torch.arange(B, device=ids.device).unsqueeze(1)
|
||||
|
||||
l_chosen = logits_last[batch_rows, ch_ids]
|
||||
l_bad = logits_last.gather(-1, rejected.unsqueeze(-1))
|
||||
delta_tok = l_chosen - l_bad
|
||||
|
||||
margin = clip_epsilon_logits
|
||||
tau = 1.0
|
||||
gap = margin - delta_tok
|
||||
per_tok_loss = F.softplus(gap / tau)
|
||||
|
||||
pref_loss = (per_tok_loss * weights).sum() / weights_sum.sum()
|
||||
|
||||
extra_metrics = {}
|
||||
|
||||
if USE_MSE_LOSS:
|
||||
with torch.no_grad():
|
||||
if self.ref_model is None:
|
||||
with self.null_ref_context():
|
||||
ref_logits_last = model(
|
||||
ids, attention_mask=attn, position_ids=pos_full,
|
||||
use_cache=False, return_dict=True,
|
||||
).logits[:, -1, :]
|
||||
else:
|
||||
ref_logits_last = self.ref_model(
|
||||
ids, attention_mask=attn, position_ids=pos_full,
|
||||
use_cache=False, return_dict=True,
|
||||
).logits[:, -1, :]
|
||||
|
||||
freeze_mask = torch.ones_like(logits_last, dtype=torch.bool)
|
||||
rows = torch.arange(B, device=ch_ids.device).unsqueeze(1).expand_as(ch_ids)
|
||||
freeze_mask[rows[ch_mask], ch_ids[ch_mask]] = False
|
||||
freeze_mask.scatter_(1, rejected.unsqueeze(-1), False)
|
||||
|
||||
diff = logits_last - ref_logits_last
|
||||
mse_elem_raw = (freeze_mask * diff.pow(2)).sum() / freeze_mask.sum()
|
||||
|
||||
tgt_mask = torch.zeros_like(logits_last, dtype=torch.bool)
|
||||
rows = torch.arange(B, device=ch_ids.device).unsqueeze(1).expand_as(ch_ids)
|
||||
tgt_mask[rows[ch_mask], ch_ids[ch_mask]] = True
|
||||
tgt_mask.scatter_(1, rejected.unsqueeze(-1), True)
|
||||
|
||||
if lambda_mse_target:
|
||||
diff_tok = logits_last - ref_logits_last
|
||||
diff_tok = diff_tok * tgt_mask
|
||||
excess_tok = torch.clamp(diff_tok.abs() - tau_mse_target, min=0.0)
|
||||
mse_target_raw = (excess_tok.pow(2)).sum() / tgt_mask.sum()
|
||||
else:
|
||||
mse_target_raw = logits_last.new_tensor(0.0)
|
||||
|
||||
mse_loss = (
|
||||
lambda_mse * mse_elem_raw
|
||||
+ lambda_mse_target * mse_target_raw
|
||||
)
|
||||
loss = pref_loss + mse_loss
|
||||
|
||||
extra_metrics.update({
|
||||
"mse_elem" : mse_elem_raw.detach(),
|
||||
"mse_tgt_tokenwise" : mse_target_raw.detach(),
|
||||
})
|
||||
|
||||
else:
|
||||
loss = pref_loss
|
||||
|
||||
lp_chosen = logp_all.gather(-1, ch_ids)
|
||||
lp_bad = logp_bad.unsqueeze(-1)
|
||||
|
||||
wins_tok = (lp_chosen > lp_bad) & ch_mask
|
||||
frac_win = wins_tok.float().sum(-1) / ch_mask.sum(-1).clamp(min=1e-8)
|
||||
chosen_win = frac_win.mean().detach()
|
||||
|
||||
metrics = {
|
||||
"pref_loss": pref_loss.detach(),
|
||||
"chosen_win": chosen_win,
|
||||
**extra_metrics,
|
||||
}
|
||||
self.store_metrics(metrics, train_eval="train")
|
||||
|
||||
if return_outputs:
|
||||
return loss, metrics
|
||||
return loss
|
||||
|
||||
|
||||
|
||||
|
||||
# ----------------------------------------------------------
|
||||
def _prepare_dataset(self, dataset, *args, **_):
|
||||
return dataset
|
||||
703
core/orchestration.py
Normal file
703
core/orchestration.py
Normal file
@@ -0,0 +1,703 @@
|
||||
import json
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import os
|
||||
import datetime
|
||||
import logging
|
||||
from pathlib import Path
|
||||
import pandas as pd
|
||||
from typing import Optional, Dict, Any, List
|
||||
import traceback
|
||||
|
||||
from utils.fs_helpers import merge_custom_bans_into_file, set_from_json
|
||||
|
||||
from core.analysis import (
|
||||
build_overrep_word_csv, select_overrep_words_for_ban,
|
||||
update_banned_slop_phrases, analyze_iteration_outputs_core,
|
||||
update_banned_ngrams_list, calculate_lexical_diversity_stats,
|
||||
calculate_repetition_score
|
||||
)
|
||||
from core.dpo import create_dpo_dataset
|
||||
from utils.whitelist import WhitelistBuilder
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# --- RESUME HELPERS -------------------------------------------------
|
||||
def _load_prompt_ids(path: Path) -> set[int]:
|
||||
"""Return the set of prompt_id ints found in an existing generation file."""
|
||||
ids = set()
|
||||
if not path.is_file(): # nothing yet
|
||||
return ids
|
||||
with path.open(encoding="utf-8") as fh:
|
||||
for ln in fh:
|
||||
try:
|
||||
row = json.loads(ln)
|
||||
pid = row.get("prompt_id")
|
||||
if isinstance(pid, int):
|
||||
ids.add(pid)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
return ids
|
||||
|
||||
def _copy_if_exists(src: Path, dst: Path) -> None:
|
||||
if src and src.is_file():
|
||||
dst.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copy2(src, dst)
|
||||
|
||||
def _patch_length_errors(jsonl_path: Path) -> None:
|
||||
"""
|
||||
Re-write rows whose error msg contains 'maximum context length'
|
||||
so that later iterations treat them like refusals.
|
||||
"""
|
||||
if not jsonl_path.is_file():
|
||||
return
|
||||
changed = False
|
||||
out_lines = []
|
||||
with jsonl_path.open(encoding="utf-8") as fh:
|
||||
for ln in fh:
|
||||
try:
|
||||
row = json.loads(ln)
|
||||
except json.JSONDecodeError:
|
||||
out_lines.append(ln); continue
|
||||
if (
|
||||
row.get("status") == "failed"
|
||||
and isinstance(row.get("error"), str)
|
||||
and "maximum context length" in row["error"]
|
||||
):
|
||||
row["status"] = "skipped -- too long"
|
||||
row["refusal_detected"] = True
|
||||
changed = True
|
||||
out_lines.append(json.dumps(row, ensure_ascii=False) + "\n")
|
||||
if changed:
|
||||
jsonl_path.write_text("".join(out_lines), encoding="utf-8")
|
||||
|
||||
|
||||
def _build_generation_command(
|
||||
main_script_path: Path,
|
||||
config: Dict[str, Any],
|
||||
output_jsonl_path: Path,
|
||||
iter_idx: int,
|
||||
banned_ngrams_file_for_iter: Optional[Path],
|
||||
slop_phrases_file_for_iter: Optional[Path],
|
||||
regex_blocklist_file_for_iter: Optional[Path]
|
||||
) -> List[str]:
|
||||
"""
|
||||
Constructs the command list for invoking the antislop-vllm generation script.
|
||||
|
||||
For iter_idx == 0 (baseline generation), all file-based banning mechanisms
|
||||
in antislop-vllm are explicitly disabled by passing empty strings for file paths
|
||||
and zero for counts, overriding any defaults in antislop-vllm's local config.yaml.
|
||||
|
||||
For subsequent iterations (iter_idx > 0), it uses the provided ban list file paths.
|
||||
Paths for file arguments are resolved to absolute paths.
|
||||
|
||||
Args:
|
||||
main_script_path: Absolute path to antislop-vllm/main.py.
|
||||
config: The main configuration dictionary for auto-antislop.
|
||||
output_jsonl_path: Absolute path for the generation output of this iteration.
|
||||
iter_idx: The current iteration index (0-based).
|
||||
banned_ngrams_file_for_iter: Path to the n-gram ban list to use for this iteration (if iter_idx > 0).
|
||||
slop_phrases_file_for_iter: Path to the slop phrase ban list to use for this iteration (if iter_idx > 0).
|
||||
regex_blocklist_file_for_iter: Path to the regex blocklist to use for this iteration (if iter_idx > 0).
|
||||
|
||||
Returns:
|
||||
A list of strings representing the command and its arguments.
|
||||
"""
|
||||
|
||||
def get_abs_path_str(p: Optional[Path]) -> Optional[str]:
|
||||
"""Resolves a Path object to an absolute path string, or returns None."""
|
||||
return str(p.resolve()) if p else None
|
||||
|
||||
ftpo_pairs_jsonl_path_str = get_abs_path_str(output_jsonl_path.parent / f"iter_{str(iter_idx)}_ftpo_pairs.jsonl")
|
||||
experiment_dir = output_jsonl_path.parent.resolve()
|
||||
|
||||
# Determine the API base URL for generation requests
|
||||
gen_api_base_url = config.get('generation_api_base_url')
|
||||
if not gen_api_base_url:
|
||||
vllm_port = config.get('vllm_port', 8000)
|
||||
gen_api_base_url = f"http://127.0.0.1:{vllm_port}/v1"
|
||||
logger.debug(
|
||||
f"generation_api_base_url not explicitly configured, defaulting to {gen_api_base_url} "
|
||||
f"based on vllm_port ({vllm_port})."
|
||||
)
|
||||
|
||||
# Core command arguments that are always present
|
||||
command_base = [
|
||||
sys.executable, str(main_script_path),
|
||||
"--api-base-url", gen_api_base_url,
|
||||
"--api-key", config['generation_api_key'],
|
||||
"--model-name", config['generation_model_id'],
|
||||
"--config", str((main_script_path.parent / "config-example.yaml").resolve()), # provides pipeline defaults that we aren't passing here
|
||||
"--output-jsonl", get_abs_path_str(output_jsonl_path),
|
||||
"--input-hf-dataset", config['generation_hf_dataset_name'],
|
||||
"--hf-dataset-split", config['generation_hf_dataset_split'],
|
||||
"--threads", str(config['generation_threads']),
|
||||
"--max-prompts", str(config['generation_max_prompts']),
|
||||
"--logging-level", config['generation_logging_level'],
|
||||
"--max-new-tokens", str(config['generation_max_new_tokens']),
|
||||
"--top-logprobs-count", str(config['generation_param_top_logprobs_count']),
|
||||
"--temperature", str(config['generation_param_temperature']),
|
||||
"--top-p", str(config['generation_param_top_p']),
|
||||
"--top-k", str(config['generation_param_top_k']),
|
||||
"--min-p", str(config['generation_param_min_p']),
|
||||
"--timeout", str(config['generation_param_timeout']),
|
||||
"--force-backtrack", str(config['generation_force_backtrack']),
|
||||
"--ngram-remove-stopwords", str(config['generation_ngram_remove_stopwords']).lower(),
|
||||
"--ngram-language", config['generation_ngram_language'],
|
||||
"--enable-refusal-detection", str(config.get("generation_refusal_detection", False)),
|
||||
"--prompt-template", config['generation_prompt_template'],
|
||||
"--system-prompt", config['generation_system_prompt'],
|
||||
]
|
||||
command = list(command_base) # Create a mutable copy
|
||||
|
||||
if iter_idx > 0:
|
||||
prev_iter_jsonl_path = experiment_dir / f"iter_{iter_idx-1}_creative_writing_generations.jsonl"
|
||||
command.extend(["--refusals-file", str(prev_iter_jsonl_path)])
|
||||
|
||||
# Use full-length chunks for the baseline run (iter_idx == 0);
|
||||
# fall back to the configured chunk size for every later iteration.
|
||||
chunk_size = (
|
||||
config['generation_max_new_tokens']
|
||||
if iter_idx == 0
|
||||
else config['generation_param_chunk_size']
|
||||
)
|
||||
command.extend(["--chunk-size", str(chunk_size)])
|
||||
|
||||
# Optional command arguments based on configuration
|
||||
if config.get('generation_param_stop_sequences'):
|
||||
stop_sequences_str = ",".join(config['generation_param_stop_sequences'])
|
||||
if stop_sequences_str: # Only add if there are actual sequences
|
||||
command.extend(["--stop-sequences", stop_sequences_str])
|
||||
|
||||
if config.get('generation_chat_template_model_id'):
|
||||
command.extend(["--chat-template-model-id", config['generation_chat_template_model_id']])
|
||||
|
||||
# --- Ban list arguments: behavior depends on iteration index ---
|
||||
if iter_idx == 0:
|
||||
# For iteration 0 (baseline), explicitly disable all file-based banning in antislop-vllm
|
||||
# by passing empty strings for file paths and zero for counts.
|
||||
# This overrides any defaults in antislop-vllm's local config.yaml.
|
||||
logger.debug("Iteration 0: Configuring antislop-vllm for baseline generation (no ban lists).")
|
||||
command.extend(["--ngram-banned-file", ""])
|
||||
command.extend(["--slop-phrases-file", ""])
|
||||
command.extend(["--top-n-slop-phrases", "0"])
|
||||
command.extend(["--regex-blocklist-file", ""])
|
||||
else:
|
||||
# this seems to overlap with another param, should fix
|
||||
command.extend(["--ftpo-pairs-jsonl", ftpo_pairs_jsonl_path_str])
|
||||
|
||||
# For iterations > 0, use the ban lists determined by the orchestrate_pipeline function.
|
||||
if banned_ngrams_file_for_iter:
|
||||
command.extend(["--ngram-banned-file", get_abs_path_str(banned_ngrams_file_for_iter)])
|
||||
|
||||
if slop_phrases_file_for_iter:
|
||||
command.extend(["--slop-phrases-file", get_abs_path_str(slop_phrases_file_for_iter)])
|
||||
# When providing a slop phrases file, instruct antislop-vllm to use all phrases from it.
|
||||
command.extend(["--top-n-slop-phrases", str(999_999)])
|
||||
|
||||
if regex_blocklist_file_for_iter:
|
||||
command.extend(["--regex-blocklist-file", get_abs_path_str(regex_blocklist_file_for_iter)])
|
||||
|
||||
return command
|
||||
|
||||
|
||||
def run_generation_script_wrapper(
|
||||
iter_idx: int,
|
||||
output_jsonl_path: Path,
|
||||
config: Dict[str, Any],
|
||||
banned_ngrams_file_path: Optional[Path] = None,
|
||||
slop_phrases_file_path: Optional[Path] = None,
|
||||
regex_blocklist_file_path: Optional[Path] = None,
|
||||
extra_generation_args: Optional[list[str]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Execute antislop-vllm/main.py for a single iteration, handling all paths,
|
||||
logging, errors, and now arbitrary extra CLI flags.
|
||||
"""
|
||||
project_root = Path(__file__).resolve().parent.parent
|
||||
main_py_script = project_root / "antislop-vllm" / "main.py"
|
||||
if not main_py_script.is_file():
|
||||
raise FileNotFoundError(
|
||||
f"antislop-vllm/main.py not found at {main_py_script}. "
|
||||
"Ensure the submodule is present and initialised."
|
||||
)
|
||||
|
||||
cmd_list = _build_generation_command(
|
||||
main_script_path=main_py_script,
|
||||
config=config,
|
||||
output_jsonl_path=output_jsonl_path,
|
||||
iter_idx=iter_idx,
|
||||
banned_ngrams_file_for_iter=banned_ngrams_file_path,
|
||||
slop_phrases_file_for_iter=slop_phrases_file_path,
|
||||
regex_blocklist_file_for_iter=regex_blocklist_file_path,
|
||||
)
|
||||
|
||||
# ── append any ad-hoc flags (e.g. --prompt-id-file <path>) ─────────────
|
||||
if extra_generation_args:
|
||||
cmd_list.extend(extra_generation_args)
|
||||
|
||||
# pretty-log (truncate very long paths)
|
||||
def _short(s: str) -> str:
|
||||
return f"...{s[-67:]}" if ("/" in s or "\\" in s) and len(s) > 70 else s
|
||||
log_cmd = " ".join(_short(c) for c in cmd_list)
|
||||
|
||||
logger.info(f"\n┏━━ Iteration {iter_idx}: launching antislop-vllm ━━━━━━━━━━━━━┓")
|
||||
logger.info(f"cwd: {main_py_script.parent}")
|
||||
logger.info(log_cmd)
|
||||
logger.info("┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛")
|
||||
|
||||
proc = subprocess.run(
|
||||
cmd_list,
|
||||
cwd=main_py_script.parent,
|
||||
check=False,
|
||||
)
|
||||
|
||||
if proc.returncode != 0:
|
||||
raise RuntimeError(
|
||||
f"antislop-vllm exited with code {proc.returncode} "
|
||||
f"(iteration {iter_idx})"
|
||||
)
|
||||
|
||||
logger.info(f"✅ antislop-vllm completed for iteration {iter_idx}. "
|
||||
f"Output: {output_jsonl_path.name}")
|
||||
|
||||
|
||||
def orchestrate_pipeline(config: Dict[str, Any], experiment_dir: Path, resume_mode: bool):
|
||||
logger.info(f"Starting anti-slop pipeline in directory: {experiment_dir}")
|
||||
generation_enabled = config.get("generation_step_enabled", True)
|
||||
if not generation_enabled:
|
||||
logger.info("⚠️ Generation step disabled by config/CLI flag.")
|
||||
|
||||
|
||||
if generation_enabled:
|
||||
# ------------------------------------------------------------------ #
|
||||
# Build project-wide whitelist #
|
||||
# ------------------------------------------------------------------ #
|
||||
whitelist_set = WhitelistBuilder.build(
|
||||
model_id = config.get("generation_chat_template_model_id") or config["vllm_model_id"],
|
||||
extra_user_items = config.get("whitelist_strings", []),
|
||||
)
|
||||
wl_path = experiment_dir / config.get("whitelist_output_filename", "whitelist_strings.json")
|
||||
WhitelistBuilder.write(wl_path, whitelist_set)
|
||||
logger.info(f"✓ Whitelist compiled → {wl_path} ({len(whitelist_set)} entries)")
|
||||
|
||||
|
||||
|
||||
# --- NLTK Stopwords ---
|
||||
try:
|
||||
from nltk.corpus import stopwords # Import here to keep it local to this function
|
||||
stop_words_set = set(stopwords.words('english'))
|
||||
logger.info(f"Loaded {len(stop_words_set)} NLTK stopwords for 'english'.")
|
||||
except LookupError:
|
||||
logger.error("NLTK 'stopwords' for 'english' not found. Please run fs_helpers.download_nltk_resource or download manually.")
|
||||
logger.error("Pipeline cannot continue without stopwords for analysis.")
|
||||
raise # Critical for analysis
|
||||
|
||||
# --- Human Profile ---
|
||||
human_profile_path = Path(config['human_profile_path'])
|
||||
if not human_profile_path.is_file():
|
||||
logger.error(f"Human profile JSON not found: {human_profile_path.resolve()}")
|
||||
raise FileNotFoundError(f"Human profile not found at {human_profile_path}")
|
||||
try:
|
||||
with human_profile_path.open("r", encoding="utf-8") as f_hp:
|
||||
human_profile_full: dict = json.load(f_hp)
|
||||
except Exception as e:
|
||||
logger.error(f"Could not load or parse human profile JSON '{human_profile_path}': {e}")
|
||||
raise
|
||||
|
||||
# --- Ban Lists Paths (initialized here, files created/updated during iterations) ---
|
||||
banned_ngrams_json_path = experiment_dir / "banned_ngrams.json"
|
||||
if 'banned_slop_phrases_filename' not in config:
|
||||
config['banned_slop_phrases_filename'] = 'banned_slop_phrases.json'
|
||||
banned_slop_phrases_json_path = experiment_dir / config['banned_slop_phrases_filename']
|
||||
|
||||
# Ensure both ban-list files exist so we can always hand them to antislop-vllm,
|
||||
# even if the associated banning feature is turned off.
|
||||
for _p in (banned_ngrams_json_path, banned_slop_phrases_json_path):
|
||||
if not _p.exists():
|
||||
_p.write_text("[]", encoding="utf-8") # write an empty JSON array
|
||||
|
||||
|
||||
# --- Regex Blocklist (user-supplied, written once if provided, used from iter 1+) ---
|
||||
# This file is created before the loop, but only passed to generation from iter 1.
|
||||
user_regex_blocklist_file: Optional[Path] = None # Renamed for clarity
|
||||
extra_regex_patterns = config.get('extra_regex_patterns', [])
|
||||
user_regex_blocklist_file = experiment_dir / "user_defined_regex_blocklist.json"
|
||||
try:
|
||||
# Write it once if resuming and it doesn't exist, or if not resuming.
|
||||
# This ensures it's available for later iterations if resuming.
|
||||
if not resume_mode or (resume_mode and not user_regex_blocklist_file.exists()):
|
||||
user_regex_blocklist_file.write_text(
|
||||
json.dumps(extra_regex_patterns, indent=2, ensure_ascii=False),
|
||||
encoding="utf-8"
|
||||
)
|
||||
logger.info(f"📝 User-defined regex blocklist written to {user_regex_blocklist_file}")
|
||||
elif user_regex_blocklist_file.exists():
|
||||
logger.info(f"📝 User-defined regex blocklist already exists at {user_regex_blocklist_file}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to write user-defined regex blocklist: {e}. It will not be used.")
|
||||
user_regex_blocklist_file = None # Disable if write fails
|
||||
|
||||
iteration_stats: list[dict] = []
|
||||
iter0_output_file_for_dpo: Optional[Path] = None
|
||||
final_iter_output_file_for_dpo: Optional[Path] = None # Tracks the latest successful output
|
||||
|
||||
start_iter_idx = 0
|
||||
if resume_mode:
|
||||
logger.info(f"Attempting to resume from {experiment_dir}...")
|
||||
max_found_iter = -1
|
||||
# Check for successfully completed iterations by looking for their output files
|
||||
max_found_iter = -1
|
||||
need_total = config['generation_max_prompts']
|
||||
|
||||
for i in range(config['num_iterations']):
|
||||
gen_file = experiment_dir / f"iter_{i}_creative_writing_generations.jsonl"
|
||||
|
||||
# Does the file exist at all?
|
||||
if not (gen_file.is_file() and gen_file.stat().st_size):
|
||||
break # nothing (or zero-length) → not complete
|
||||
|
||||
# Does it contain the full prompt set?
|
||||
ids_seen = _load_prompt_ids(gen_file)
|
||||
if len(ids_seen) < need_total:
|
||||
logger.info(
|
||||
f"Iteration {i} resume-check: {len(ids_seen)}/{need_total} prompts present "
|
||||
f"({need_total-len(ids_seen)} still missing).")
|
||||
break # incomplete → resume here
|
||||
|
||||
max_found_iter = i # this one is done, keep going
|
||||
|
||||
|
||||
if max_found_iter >= 0:
|
||||
start_iter_idx = max_found_iter + 1
|
||||
logger.info(f"Resuming from iteration {start_iter_idx}.")
|
||||
# Log presence of existing ban lists if resuming past iter 0
|
||||
if start_iter_idx > 0:
|
||||
if banned_ngrams_json_path.exists(): logger.info(f"Resuming with existing n-gram ban list: {banned_ngrams_json_path}")
|
||||
else: logger.info("No existing n-gram ban list found to resume with for subsequent iterations.")
|
||||
if banned_slop_phrases_json_path.exists(): logger.info(f"Resuming with existing slop phrase ban list: {banned_slop_phrases_json_path}")
|
||||
else: logger.info("No existing slop phrase ban list found to resume with for subsequent iterations.")
|
||||
else:
|
||||
logger.info("No fully completed iterations found to resume. Starting from iteration 0.")
|
||||
# resume_mode = False # No need to change resume_mode, start_iter_idx handles it
|
||||
|
||||
if start_iter_idx >= config['num_iterations']:
|
||||
logger.info(f"All {config['num_iterations']} iterations appear to be completed in {experiment_dir}.")
|
||||
# Attempt to load existing stats for DPO if needed
|
||||
summary_csv_path = experiment_dir / "final_iteration_statistics.csv"
|
||||
if summary_csv_path.exists():
|
||||
try:
|
||||
iteration_stats_df = pd.read_csv(summary_csv_path)
|
||||
iteration_stats = iteration_stats_df.to_dict('records')
|
||||
# Ensure iter0_output_file_for_dpo and final_iter_output_file_for_dpo are set if possible
|
||||
if not iter0_output_file_for_dpo and not iteration_stats_df.empty:
|
||||
iter0_row = iteration_stats_df[iteration_stats_df['iteration'] == 0]
|
||||
if not iter0_row.empty and 'output_file' in iter0_row.columns:
|
||||
path_str = iter0_row.iloc[0]['output_file']
|
||||
if path_str and isinstance(path_str, str): iter0_output_file_for_dpo = experiment_dir / path_str
|
||||
if not final_iter_output_file_for_dpo and not iteration_stats_df.empty:
|
||||
# Find the last completed iteration in the stats
|
||||
last_stat_iter = iteration_stats_df['iteration'].max()
|
||||
final_iter_row = iteration_stats_df[iteration_stats_df['iteration'] == last_stat_iter]
|
||||
if not final_iter_row.empty and 'output_file' in final_iter_row.columns:
|
||||
path_str = final_iter_row.iloc[0]['output_file']
|
||||
if path_str and isinstance(path_str, str): final_iter_output_file_for_dpo = experiment_dir / path_str
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not load or parse existing iteration statistics from {summary_csv_path}: {e}")
|
||||
# Proceed to DPO creation if applicable (handled after the loop)
|
||||
else: # Need to run some or all iterations
|
||||
if generation_enabled:
|
||||
for iter_idx in range(start_iter_idx, config['num_iterations']):
|
||||
current_iter_start_time = datetime.datetime.now()
|
||||
logger.info(f"\n{'='*30} ITERATION {iter_idx} (started at {current_iter_start_time.strftime('%H:%M:%S')}) {'='*30}")
|
||||
|
||||
iter_output_jsonl = experiment_dir / f"iter_{iter_idx}_creative_writing_generations.jsonl"
|
||||
iter_analysis_dir = experiment_dir / f"iter_{iter_idx}_analysis_results"
|
||||
iter_analysis_dir.mkdir(parents=True, exist_ok=True) # Ensure analysis dir exists
|
||||
|
||||
# --- Determine ban lists for the current iteration ---
|
||||
# Iteration 0 (baseline) runs with NO BANNING.
|
||||
# Subsequent iterations use the ban lists accumulated so far.
|
||||
ngram_file_for_generation: Optional[Path] = None
|
||||
slop_file_for_generation: Optional[Path] = None
|
||||
regex_file_for_generation: Optional[Path] = None
|
||||
|
||||
if iter_idx > 0: # Banning starts from iteration 1
|
||||
if banned_ngrams_json_path.exists():
|
||||
ngram_file_for_generation = banned_ngrams_json_path
|
||||
if banned_slop_phrases_json_path.exists():
|
||||
slop_file_for_generation = banned_slop_phrases_json_path
|
||||
if user_regex_blocklist_file and user_regex_blocklist_file.exists(): # User-defined regex
|
||||
regex_file_for_generation = user_regex_blocklist_file
|
||||
|
||||
# If we are resuming and this is the first iteration after the resume,
|
||||
# force-merge any new YAML bans into the existing files *before* generation.
|
||||
if resume_mode and iter_idx == start_iter_idx:
|
||||
if config['enable_ngram_ban'] and config.get('extra_ngrams_to_ban'):
|
||||
merge_custom_bans_into_file(banned_ngrams_json_path,
|
||||
config['extra_ngrams_to_ban'])
|
||||
if config['enable_slop_phrase_ban'] and config.get('extra_slop_phrases_to_ban'):
|
||||
merge_custom_bans_into_file(banned_slop_phrases_json_path,
|
||||
config['extra_slop_phrases_to_ban'])
|
||||
|
||||
_copy_if_exists(ngram_file_for_generation,
|
||||
iter_analysis_dir / "banned_ngrams_used.json")
|
||||
_copy_if_exists(slop_file_for_generation,
|
||||
iter_analysis_dir / "banned_slop_phrases_used.json")
|
||||
_copy_if_exists(regex_file_for_generation,
|
||||
iter_analysis_dir / "regex_blocklist_used.json")
|
||||
|
||||
# remember their current contents so we can diff later
|
||||
before_ngrams = set_from_json(ngram_file_for_generation)
|
||||
before_slop = set_from_json(slop_file_for_generation)
|
||||
|
||||
else:
|
||||
before_ngrams, before_slop = set(), set()
|
||||
|
||||
if iter_idx == 0:
|
||||
logger.info("Iteration 0: Running baseline generation with NO ban lists.")
|
||||
else:
|
||||
logger.info(f"Iteration {iter_idx}: Using ban lists - N-grams: {ngram_file_for_generation}, Slop: {slop_file_for_generation}, Regex: {regex_file_for_generation}")
|
||||
|
||||
# ────────────────────────────────────────────────────────────────────
|
||||
# A. fast-path – is generation already complete?
|
||||
# ────────────────────────────────────────────────────────────────────
|
||||
existing_ids = _load_prompt_ids(iter_output_jsonl)
|
||||
need_total = config['generation_max_prompts']
|
||||
missing_ids = sorted(set(range(need_total)) - existing_ids)
|
||||
|
||||
if not missing_ids:
|
||||
logger.info(f"Iteration {iter_idx}: found {need_total} / {need_total} prompts "
|
||||
f"in {iter_output_jsonl.name} – skipping generation step.")
|
||||
else:
|
||||
logger.info(f"Iteration {iter_idx}: {len(missing_ids)} / {need_total} prompts "
|
||||
f"still missing – resuming generation.")
|
||||
# antislop-vllm already supports '--prompt-id-file' (one id per line)
|
||||
# we can skip this as antislop-vllm automatically resumes now
|
||||
#miss_file = _write_missing_prompt_file(missing_ids, experiment_dir, iter_idx)
|
||||
|
||||
try:
|
||||
run_generation_script_wrapper(
|
||||
iter_idx = iter_idx,
|
||||
output_jsonl_path = iter_output_jsonl,
|
||||
config = config,
|
||||
banned_ngrams_file_path= ngram_file_for_generation,
|
||||
slop_phrases_file_path = slop_file_for_generation,
|
||||
regex_blocklist_file_path = regex_file_for_generation,
|
||||
#extra_generation_args = ["--prompt-id-file", str(miss_file)]
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"❌ Generation script failed for iteration {iter_idx}: {e}")
|
||||
# identical failure-handling block as before …
|
||||
iteration_stats.append({
|
||||
"iteration": iter_idx, "status": "generation_failed",
|
||||
"error": str(e), "output_file": str(iter_output_jsonl.name)
|
||||
})
|
||||
if iter_idx == 0:
|
||||
iter0_output_file_for_dpo = None
|
||||
continue
|
||||
|
||||
# turn “max-context” failures into skips so later iterations don’t retry them
|
||||
_patch_length_errors(iter_output_jsonl)
|
||||
|
||||
|
||||
if not iter_output_jsonl.exists() or iter_output_jsonl.stat().st_size == 0:
|
||||
logger.error(f"❌ Generation output file {iter_output_jsonl} is missing or empty for iteration {iter_idx}.")
|
||||
iteration_stats.append({
|
||||
"iteration": iter_idx, "status": "output_file_missing_or_empty",
|
||||
"output_file": str(iter_output_jsonl.name)
|
||||
})
|
||||
if iter_idx == 0: iter0_output_file_for_dpo = None
|
||||
continue
|
||||
|
||||
# Update DPO file pointers
|
||||
if iter_idx == 0:
|
||||
iter0_output_file_for_dpo = iter_output_jsonl
|
||||
# final_iter_output_file_for_dpo always points to the latest successfully generated file
|
||||
final_iter_output_file_for_dpo = iter_output_jsonl
|
||||
|
||||
# --- Analysis (runs for all iterations, including iter 0 to find initial slop) ---
|
||||
analysis_results = None
|
||||
try:
|
||||
analysis_results = analyze_iteration_outputs_core(
|
||||
generated_jsonl_path=iter_output_jsonl,
|
||||
human_profile_full=human_profile_full,
|
||||
iter_analysis_output_dir=iter_analysis_dir,
|
||||
config=config,
|
||||
stop_words_set=stop_words_set
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"❌ Text analysis failed for iteration {iter_idx}: {e}", exc_info=True)
|
||||
iteration_stats.append({
|
||||
"iteration": iter_idx, "status": "analysis_failed",
|
||||
"error": str(e), "output_file": str(iter_output_jsonl.name)
|
||||
})
|
||||
continue
|
||||
|
||||
if analysis_results is None or analysis_results[0] is None: # DFs are first part of tuple
|
||||
logger.warning(f"Analysis for iteration {iter_idx} did not produce data. Skipping ban list updates for this iteration.")
|
||||
iteration_stats.append({
|
||||
"iteration": iter_idx, "status": "analysis_no_data",
|
||||
"output_file": str(iter_output_jsonl.name)
|
||||
})
|
||||
continue
|
||||
|
||||
df_bi_dict, df_bi_nondct, df_tri_dict, df_tri_nondct, generated_texts, total_gen_chars = analysis_results
|
||||
if not generated_texts:
|
||||
logger.warning(f"No generated texts found after analysis for iter {iter_idx}. Skipping ban list updates.")
|
||||
iteration_stats.append({
|
||||
"iteration": iter_idx, "status": "no_texts_post_analysis",
|
||||
"output_file": str(iter_output_jsonl.name)
|
||||
})
|
||||
continue
|
||||
|
||||
# --- Update Ban Lists (based on current iteration's analysis) ---
|
||||
# These lists will be used by the *next* iteration's generation step.
|
||||
# --- Update ban lists (based on current iteration's analysis) --------------
|
||||
overrep_tokens_for_ban: list[str] = []
|
||||
iter_log = iter_analysis_dir / "orchestration.log"
|
||||
def _iter_log(msg: str) -> None:
|
||||
with iter_log.open("a", encoding="utf-8") as fh:
|
||||
fh.write(f"{datetime.datetime.now():%Y-%m-%d %H:%M:%S} {msg}\n")
|
||||
|
||||
# (a) over-represented words -------------------------------------------------
|
||||
if config['compute_overrep_words']:
|
||||
try:
|
||||
overrep_csv = iter_analysis_dir / "overrepresented_words.csv"
|
||||
_, dict_words, nodict_words = build_overrep_word_csv(
|
||||
texts=generated_texts,
|
||||
out_csv=overrep_csv,
|
||||
top_n_words_analysis=config['top_k_words_for_overrep_analysis'],
|
||||
stop_words_set=stop_words_set,
|
||||
)
|
||||
overrep_tokens_for_ban = select_overrep_words_for_ban(
|
||||
dict_words, nodict_words, (iter_idx == 0), config, whitelist=whitelist_set
|
||||
)
|
||||
_iter_log(f"overrep_tokens_for_ban = {len(overrep_tokens_for_ban)}")
|
||||
except Exception as exc:
|
||||
_iter_log("❌ build_overrep_word_csv failed:\n" +
|
||||
"".join(traceback.format_exception_only(type(exc), exc)))
|
||||
|
||||
# (b) n-gram ban list --------------------------------------------------------
|
||||
if config['enable_ngram_ban']:
|
||||
try:
|
||||
update_banned_ngrams_list(
|
||||
banned_ngrams_json_path,
|
||||
dfs=[df_bi_dict, df_bi_nondct, df_tri_dict, df_tri_nondct],
|
||||
is_first_iteration=(iter_idx == 0),
|
||||
config=config,
|
||||
whitelist=whitelist_set,
|
||||
)
|
||||
_iter_log("n-gram ban list updated")
|
||||
except Exception as exc:
|
||||
_iter_log("❌ update_banned_ngrams_list failed:\n" +
|
||||
"".join(traceback.format_exception_only(type(exc), exc)))
|
||||
|
||||
# (c) slop-phrase ban list ---------------------------------------------------
|
||||
if config['enable_slop_phrase_ban']:
|
||||
try:
|
||||
phrases_to_add_count = (
|
||||
config['top_n_initial_slop_ban'] if iter_idx == 0
|
||||
else config['top_n_subsequent_slop_ban']
|
||||
)
|
||||
update_banned_slop_phrases(
|
||||
json_path=banned_slop_phrases_json_path,
|
||||
texts=generated_texts,
|
||||
how_many_new=phrases_to_add_count,
|
||||
tmp_dir=iter_analysis_dir / "phrase_tmp",
|
||||
config=config,
|
||||
whitelist=whitelist_set,
|
||||
over_represented_words=(
|
||||
overrep_tokens_for_ban
|
||||
),
|
||||
)
|
||||
_iter_log("slop-phrase ban list updated "
|
||||
f"(+{len(overrep_tokens_for_ban)} over-rep words)")
|
||||
except Exception as exc:
|
||||
_iter_log("❌ update_banned_slop_phrases failed:\n" +
|
||||
"".join(traceback.format_exception_only(type(exc), exc)))
|
||||
|
||||
# ---------- diff → what was *added* this iteration -------------------------
|
||||
if iter_idx > 0:
|
||||
# n-grams
|
||||
if config['enable_ngram_ban'] and banned_ngrams_json_path.exists():
|
||||
after_ngrams = set_from_json(banned_ngrams_json_path)
|
||||
new_ngrams = sorted(after_ngrams - before_ngrams)
|
||||
(iter_analysis_dir / "banned_ngrams_new_this_iter.json"
|
||||
).write_text(json.dumps(new_ngrams, indent=2, ensure_ascii=False), "utf-8")
|
||||
|
||||
# slop phrases
|
||||
if config['enable_slop_phrase_ban'] and banned_slop_phrases_json_path.exists():
|
||||
after_slop = set_from_json(banned_slop_phrases_json_path)
|
||||
new_slop = sorted(after_slop - before_slop)
|
||||
(iter_analysis_dir / "banned_slop_phrases_new_this_iter.json"
|
||||
).write_text(json.dumps(new_slop, indent=2, ensure_ascii=False), "utf-8")
|
||||
|
||||
|
||||
|
||||
# --- Calculate Metrics for this iteration ---
|
||||
ttr, rttr, repetition_norm = 0.0, 0.0, 0.0
|
||||
try:
|
||||
ttr, rttr = calculate_lexical_diversity_stats(generated_texts, config['min_word_len_for_analysis'])
|
||||
repetition_norm = calculate_repetition_score(
|
||||
generated_texts, total_gen_chars,
|
||||
[df_bi_dict, df_bi_nondct, df_tri_dict, df_tri_nondct], config, stop_words_set
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"❌ Error calculating metrics for iteration {iter_idx}: {e}", exc_info=True)
|
||||
|
||||
iteration_stats.append({
|
||||
"iteration": iter_idx, "status": "completed",
|
||||
"generated_text_count": len(generated_texts), "generated_char_count": total_gen_chars,
|
||||
"ttr": ttr, "rttr": rttr, "repetition_per_100k_chars": repetition_norm,
|
||||
"output_file": str(iter_output_jsonl.name), "error": None
|
||||
})
|
||||
iter_duration = datetime.datetime.now() - current_iter_start_time
|
||||
logger.info(f"--- Iteration {iter_idx} completed in {iter_duration} ---")
|
||||
else:
|
||||
logger.info("Skipping generation loop.")
|
||||
|
||||
if generation_enabled:
|
||||
# --- Final Summary & DPO Dataset Creation ---
|
||||
summary_df = pd.DataFrame(iteration_stats)
|
||||
summary_csv = experiment_dir / "final_iteration_statistics.csv"
|
||||
try:
|
||||
summary_df.to_csv(summary_csv, index=False)
|
||||
logger.info(f"\n📊 Final statistics written to {summary_csv.resolve()}")
|
||||
if not summary_df.empty:
|
||||
# Ensure all columns are displayed if possible
|
||||
with pd.option_context('display.max_rows', None, 'display.max_columns', None, 'display.width', 1000):
|
||||
logger.info("\n" + summary_df.to_string(index=False, na_rep="N/A"))
|
||||
else:
|
||||
logger.info("No iteration statistics were generated to summarize.")
|
||||
except Exception as e:
|
||||
logger.error(f"Could not write final statistics CSV to {summary_csv}: {e}")
|
||||
|
||||
# DPO dataset creation logic
|
||||
if config['num_iterations'] >= 1 and iter0_output_file_for_dpo and final_iter_output_file_for_dpo:
|
||||
if iter0_output_file_for_dpo.exists() and final_iter_output_file_for_dpo.exists():
|
||||
if config['num_iterations'] == 1 and iter0_output_file_for_dpo == final_iter_output_file_for_dpo:
|
||||
logger.warning("Only one iteration completed. DPO dataset 'chosen' and 'rejected' will be from the same iter_0 data. This might not be useful for training.")
|
||||
|
||||
dpo_output_jsonl = experiment_dir / "dpo_pairs_dataset.jsonl"
|
||||
try:
|
||||
create_dpo_dataset(iter0_output_file_for_dpo, final_iter_output_file_for_dpo, dpo_output_jsonl)
|
||||
except Exception as e:
|
||||
logger.error(f"❌ ERROR creating DPO dataset: {e}", exc_info=True)
|
||||
else:
|
||||
logger.warning(
|
||||
f"DPO dataset creation skipped: Iteration 0 output file ({iter0_output_file_for_dpo}) "
|
||||
f"or final iteration output file ({final_iter_output_file_for_dpo}) not found or generation failed."
|
||||
)
|
||||
elif config['num_iterations'] < 1:
|
||||
logger.info("No iterations were configured to run. DPO dataset creation skipped.")
|
||||
else: # Cases where DPO files might be None due to errors
|
||||
logger.warning(
|
||||
f"DPO dataset creation skipped due to missing DPO source files. "
|
||||
f"Iter0 source: {iter0_output_file_for_dpo}, Final iter source: {final_iter_output_file_for_dpo}"
|
||||
)
|
||||
|
||||
logger.info("Anti-slop pipeline orchestration finished.")
|
||||
return experiment_dir
|
||||
4000010
data/human_writing_profile.json
Normal file
4000010
data/human_writing_profile.json
Normal file
File diff suppressed because it is too large
Load Diff
59
data/human_writing_profile_info.txt
Normal file
59
data/human_writing_profile_info.txt
Normal file
@@ -0,0 +1,59 @@
|
||||
Generated with https://github.com/sam-paech/slop-forensics
|
||||
|
||||
Using (human written) outputs sourced from: Nitral-AI/Reddit-SFW-Writing_Prompts_ShareGPT
|
||||
|
||||
|
||||
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Convert the Nitral-AI ShareGPT dataset into a JSONL file that the
|
||||
slop-forensics analysis pipeline can read directly.
|
||||
"""
|
||||
|
||||
import os, json, logging
|
||||
from tqdm import tqdm
|
||||
from datasets import load_dataset, disable_caching
|
||||
|
||||
# ---------- config ----------
|
||||
OUTPUT_PATH = "results/datasets/generated_Nitral-AI__human.jsonl"
|
||||
SOURCE_NAME = "Nitral-AI" # kept for consistency with other scripts
|
||||
MODEL_FIELD = "human-authored" # any string is fine – it’s just a label
|
||||
# -----------------------------
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s — %(levelname)s — %(message)s")
|
||||
disable_caching()
|
||||
|
||||
def main():
|
||||
os.makedirs(os.path.dirname(OUTPUT_PATH), exist_ok=True)
|
||||
ds = load_dataset("Nitral-AI/Reddit-SFW-Writing_Prompts_ShareGPT", split="train")
|
||||
|
||||
with open(OUTPUT_PATH, "w", encoding="utf-8") as f_out:
|
||||
kept = 0
|
||||
for idx, row in enumerate(tqdm(ds, desc="Extracting GPT turns")):
|
||||
# locate first GPT message in the ShareGPT conversation
|
||||
gpt_msg = None
|
||||
for msg in row.get("conversations", []):
|
||||
if isinstance(msg, dict) and msg.get("from") == "gpt":
|
||||
gpt_msg = msg.get("value", "").strip()
|
||||
break
|
||||
if not gpt_msg:
|
||||
continue # skip rows without a GPT response
|
||||
|
||||
record = {
|
||||
"source": SOURCE_NAME,
|
||||
"id": idx,
|
||||
"prompt": row.get("title", ""), # not used by analysis but nice to keep
|
||||
"model": MODEL_FIELD,
|
||||
"output": gpt_msg
|
||||
}
|
||||
f_out.write(json.dumps(record, ensure_ascii=False) + "\n")
|
||||
kept += 1
|
||||
|
||||
logging.info(f"Finished. Wrote {kept} records to {OUTPUT_PATH}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
|
||||
# Then:
|
||||
python3 scripts/slop_profile.py --input-dir results/datasets --analysis-output-dir results/analysis_human --combined-output-file results/human_slop_profile.json --top-n 1000000 --max-items 9999999
|
||||
326
main.py
Normal file
326
main.py
Normal file
@@ -0,0 +1,326 @@
|
||||
import argparse
|
||||
import logging
|
||||
import sys
|
||||
import os
|
||||
import json
|
||||
import datetime
|
||||
from pathlib import Path
|
||||
import datetime # For pipeline duration
|
||||
import yaml
|
||||
from pathlib import PurePath # base for PosixPath / WindowsPath
|
||||
|
||||
# register once – covers Path, PosixPath, WindowsPath …
|
||||
yaml.SafeDumper.add_multi_representer(
|
||||
PurePath,
|
||||
lambda dumper, value: dumper.represent_scalar(
|
||||
"tag:yaml.org,2002:str", str(value))
|
||||
)
|
||||
|
||||
# ── make utils importable ────────────────────────────────────────────
|
||||
ROOT_DIR = Path(__file__).resolve().parent
|
||||
sys.path.insert(0, str(ROOT_DIR)) # so "utils" is on sys.path
|
||||
|
||||
# ── guarantee NLTK data is present *before* any other project import ─
|
||||
from utils.fs_helpers import ensure_core_nltk_resources
|
||||
ensure_core_nltk_resources() # downloads punkt, punkt_tab, stopwords
|
||||
|
||||
|
||||
# --- Add project directories to sys.path ---
|
||||
# This allows importing from core, utils, and submodules
|
||||
sys.path.insert(0, str(ROOT_DIR / "slop-forensics"))
|
||||
# antislop-vllm is called as a script, its path for direct import is not strictly needed
|
||||
# unless some of its utils were to be imported by auto-antislop (not the current plan).
|
||||
|
||||
from utils.config_loader import load_pipeline_config, merge_config_with_cli_args
|
||||
from utils.fs_helpers import (
|
||||
create_experiment_dir,
|
||||
ensure_antislop_vllm_config_exists
|
||||
)
|
||||
from utils.vllm_manager import start_vllm_server, stop_vllm_server, is_vllm_server_alive
|
||||
from core.orchestration import orchestrate_pipeline
|
||||
from core.finetuning import run_dpo_finetune
|
||||
|
||||
# --- Basic Logging Setup -------------------------------------------------
|
||||
logging.basicConfig( # root stays at WARNING
|
||||
level=logging.WARNING,
|
||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
||||
)
|
||||
logger = logging.getLogger("auto_antislop_main")
|
||||
|
||||
|
||||
def str2bool(v):
|
||||
if v is None:
|
||||
return None
|
||||
if isinstance(v, bool):
|
||||
return v
|
||||
v = str(v).lower()
|
||||
if v in ("yes", "true", "t", "1", "y"):
|
||||
return True
|
||||
if v in ("no", "false", "f", "0", "n"):
|
||||
return False
|
||||
raise argparse.ArgumentTypeError("Boolean value expected.")
|
||||
|
||||
# ── QUICK CHECK: are *all* generation files already complete? ───────────────
|
||||
def _all_generations_done(cfg: dict, resume_dir: Path | None) -> bool:
|
||||
if not resume_dir or not resume_dir.is_dir():
|
||||
return False
|
||||
|
||||
need = cfg.get("generation_max_prompts", 0)
|
||||
if need <= 0:
|
||||
return False
|
||||
|
||||
def _ids(path: Path) -> int:
|
||||
if not path.is_file():
|
||||
return 0
|
||||
seen = set()
|
||||
for ln in path.read_text(encoding="utf-8").splitlines():
|
||||
try:
|
||||
seen.add(int(json.loads(ln).get("prompt_id", -1)))
|
||||
except Exception:
|
||||
pass
|
||||
return len(seen)
|
||||
|
||||
for i in range(cfg["num_iterations"]):
|
||||
p = resume_dir / f"iter_{i}_creative_writing_generations.jsonl"
|
||||
if _ids(p) < need:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Auto-Antislop: Iterative dataset generation and DPO finetuning.")
|
||||
|
||||
# --- General Arguments ---
|
||||
parser.add_argument(
|
||||
"-c", "--config-file", type=Path, default=Path("auto_antislop_config.yaml"),
|
||||
help="Path to the main YAML configuration file."
|
||||
)
|
||||
parser.add_argument(
|
||||
"-r", "--resume-from-dir", type=Path, default=None,
|
||||
help="Path to an existing experiment run directory to resume."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-level", choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
|
||||
default=None, help="Set the logging level for the auto-antislop script."
|
||||
)
|
||||
|
||||
# --- vLLM Management ---
|
||||
vllm_group = parser.add_argument_group('vLLM Server Management')
|
||||
vllm_group.add_argument(
|
||||
"--manage-vllm",
|
||||
type=str2bool,
|
||||
nargs="?",
|
||||
const=True, # `--manage-vllm` ⇒ True
|
||||
default=None, # fall back to config
|
||||
help="true/false to let this script start/stop a local vLLM server "
|
||||
"(default comes from config)."
|
||||
)
|
||||
vllm_group.add_argument("--vllm-port", type=int, default=None, help="Port for vLLM server. Overrides config.")
|
||||
vllm_group.add_argument("--vllm-model-id", type=str, default=None, help="Model ID for vLLM server. Overrides config.")
|
||||
vllm_group.add_argument(
|
||||
"--generation-api-base-url", type=str,
|
||||
default=None,
|
||||
help="API base URL for generation requests (passed to antislop-vllm). E.g., http://host:port/v1. Overrides config."
|
||||
)
|
||||
|
||||
# --- Pipeline Control ---
|
||||
pipeline_group = parser.add_argument_group('Pipeline Control')
|
||||
pipeline_group.add_argument("--num-iterations", type=int, default=None, help="Number of anti-slop iterations. Overrides config.")
|
||||
pipeline_group.add_argument("--generation-max-prompts", type=int, default=None, help="Max prompts for antislop-vllm. Overrides config.")
|
||||
pipeline_group.add_argument(
|
||||
"--generation-step-enabled",
|
||||
type=str2bool,
|
||||
nargs="?",
|
||||
const=True,
|
||||
default=None,
|
||||
help="true/false to execute the generation step. "
|
||||
"(default from config)."
|
||||
)
|
||||
|
||||
# --- Finetuning Control ---
|
||||
finetune_group = parser.add_argument_group('DPO Finetuning Control')
|
||||
finetune_group.add_argument(
|
||||
"--run-finetune",
|
||||
type=str2bool,
|
||||
nargs="?",
|
||||
const=True,
|
||||
default=None,
|
||||
help="true/false to run DPO finetuning after the pipeline "
|
||||
"(default from config)."
|
||||
)
|
||||
|
||||
finetune_group.add_argument("--finetune-base-model-id", type=str, default=None, help="Base model for DPO. Overrides config.")
|
||||
finetune_group.add_argument("--finetune-num-epochs", type=int, default=None, help="Number of epochs for DPO. Overrides config.")
|
||||
|
||||
finetune_group.add_argument(
|
||||
"--finetune-mode",
|
||||
choices=["dpo", "ftpo"],
|
||||
default=None,
|
||||
help="dpo = vanilla DPO on full continuations (default); "
|
||||
"ftpo = masked Tokenwise-DPO on partial generation pairs, only computing loss for the completion token."
|
||||
)
|
||||
finetune_group.add_argument(
|
||||
"--finetune-ftpo-dataset",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="(Optional) explicit path to a ftpo/last-token JSONL file. "
|
||||
"If omitted and --finetune-mode is ftpo, the script will "
|
||||
"pick the highest iter_*_ftpo_pairs.jsonl in the experiment dir."
|
||||
)
|
||||
finetune_group.add_argument(
|
||||
"--finetune-cuda-visible-devices",
|
||||
type=str,
|
||||
default=None,
|
||||
help='Comma-separated GPU ids for the finetune stage only (e.g. "1,3").'
|
||||
)
|
||||
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# --- Load and Merge Configuration ---
|
||||
config = load_pipeline_config(args.config_file)
|
||||
config = merge_config_with_cli_args(config, args)
|
||||
|
||||
# refine levels once CLI/YAML are merged
|
||||
numeric_log_level = getattr(logging, config['log_level'].upper(), logging.INFO)
|
||||
|
||||
# raise only *our* loggers, keep external libs at WARNING
|
||||
for name in logging.root.manager.loggerDict:
|
||||
if name.startswith(("auto_antislop", "core", "utils")):
|
||||
l = logging.getLogger(name)
|
||||
l.setLevel(numeric_log_level)
|
||||
for h in l.handlers:
|
||||
h.setLevel(min(numeric_log_level, h.level))
|
||||
|
||||
# keep root at WARNING so torch / dynamo INFO spam is hidden
|
||||
logging.getLogger().setLevel(logging.WARNING)
|
||||
logger.info(f"Logging level for project set to: {config['log_level'].upper()}")
|
||||
|
||||
|
||||
|
||||
|
||||
# --- Ensure NLTK resources ---
|
||||
# These are used by core.analysis
|
||||
# --- Ensure *all* NLTK resources are present *before* anything else ---
|
||||
logger.info("Verifying / downloading required NLTK data …")
|
||||
ensure_core_nltk_resources()
|
||||
|
||||
# --- Ensure antislop-vllm config-example is copied (user convenience) ---
|
||||
antislop_vllm_dir = ROOT_DIR / "antislop-vllm"
|
||||
if antislop_vllm_dir.is_dir():
|
||||
ensure_antislop_vllm_config_exists(antislop_vllm_dir)
|
||||
else:
|
||||
logger.warning(f"antislop-vllm submodule directory not found at {antislop_vllm_dir}. Generation will likely fail.")
|
||||
|
||||
|
||||
# --- vLLM Server Management --------------------------------------------------
|
||||
vllm_server_proc = None
|
||||
should_manage_vllm = config.get('manage_vllm', True)
|
||||
|
||||
# Fast-path: if every generation file is already finished, don’t even start vLLM
|
||||
if should_manage_vllm and _all_generations_done(config, args.resume_from_dir):
|
||||
logger.info("✨ All generation files complete – skipping vLLM startup altogether.")
|
||||
should_manage_vllm = False
|
||||
config['manage_vllm'] = False # keep downstream logic consistent
|
||||
|
||||
|
||||
if should_manage_vllm:
|
||||
if not is_vllm_server_alive(config['vllm_port']):
|
||||
logger.info("Attempting to start and manage vLLM server.")
|
||||
vllm_server_proc = start_vllm_server(
|
||||
model_id=config['vllm_model_id'],
|
||||
port=config['vllm_port'],
|
||||
hf_token=config.get('vllm_hf_token'),
|
||||
cuda_visible_devices=config['vllm_cuda_visible_devices'],
|
||||
gpu_memory_utilization=config['vllm_gpu_memory_utilization'],
|
||||
max_model_len=config['vllm_max_model_len'],
|
||||
dtype=config['vllm_dtype'],
|
||||
vllm_extra_args=config.get('vllm_extra_args'),
|
||||
extra_env=config.get('vllm_env'),
|
||||
uvicorn_log_level="error", # <-- cut vllm chatter
|
||||
quiet_stdout=True, # <-- discard server stream
|
||||
)
|
||||
if vllm_server_proc is None: # Failed to start
|
||||
logger.error("Failed to start managed vLLM server. Exiting.")
|
||||
sys.exit(1)
|
||||
else:
|
||||
logger.info(f"vLLM server already running on port {config['vllm_port']}. Script will not manage it.")
|
||||
should_manage_vllm = False # Don't try to stop it later
|
||||
else:
|
||||
logger.info("vLLM server management is disabled by config/CLI.")
|
||||
#if not is_vllm_server_alive(config['vllm_port']):
|
||||
# logger.warning(f"vLLM server management disabled, but no server found on port {config['vllm_port']}. "
|
||||
# "The generation pipeline will likely fail. Please start a vLLM server manually.")
|
||||
|
||||
|
||||
# --- Main Pipeline ---
|
||||
pipeline_start_time = datetime.datetime.now()
|
||||
experiment_run_dir = None
|
||||
try:
|
||||
base_dir = Path(config['experiment_base_dir'])
|
||||
resume_dir_path = Path(config['resume_from_dir']) if config.get('resume_from_dir', None) else None
|
||||
experiment_run_dir = create_experiment_dir(base_dir, resume_dir_path)
|
||||
|
||||
# Pass the actual experiment_run_dir to orchestrate_pipeline
|
||||
config['current_experiment_run_dir'] = str(experiment_run_dir)
|
||||
|
||||
# ---------- persist the exact config used for this run ----------
|
||||
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
cfg_path = experiment_run_dir / f"run_config_{timestamp}.yaml"
|
||||
cfg_path.write_text(
|
||||
yaml.safe_dump(config, sort_keys=False, allow_unicode=True),
|
||||
encoding="utf-8"
|
||||
)
|
||||
logger.info(f"Run configuration written → {cfg_path}")
|
||||
|
||||
orchestrate_pipeline(config, experiment_run_dir, resume_mode=(resume_dir_path is not None))
|
||||
|
||||
except FileNotFoundError as e:
|
||||
logger.error(f"A required file was not found: {e}. Halting pipeline.")
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
logger.error(f"An unexpected error occurred during the anti-slop pipeline: {e}", exc_info=True)
|
||||
sys.exit(1)
|
||||
finally:
|
||||
pipeline_duration = datetime.datetime.now() - pipeline_start_time
|
||||
logger.info(f"Total anti-slop pipeline duration: {pipeline_duration}")
|
||||
|
||||
# --- Finetuning (Optional) ---
|
||||
should_run_finetune = config.get('finetune_enabled', False)
|
||||
|
||||
if should_run_finetune:
|
||||
if experiment_run_dir:
|
||||
# NEW: shut down vLLM so the GPU is free for training
|
||||
if should_manage_vllm and vllm_server_proc:
|
||||
logger.info("Stopping managed vLLM server before finetuning.")
|
||||
stop_vllm_server(vllm_server_proc)
|
||||
vllm_server_proc = None # prevent a second stop later
|
||||
|
||||
logger.info("Proceeding to finetuning.")
|
||||
finetune_start_time = datetime.datetime.now()
|
||||
try:
|
||||
finetune_output_dir = experiment_run_dir / f"finetuned_model{config['finetune_output_dir_suffix']}"
|
||||
if finetune_output_dir.exists():
|
||||
reply = input(f"⚠️ Finetune dir '{finetune_output_dir}' already exists. "
|
||||
"Delete & re-run finetune? [y/N]: ").strip().lower()
|
||||
if reply != "y":
|
||||
logger.info("Finetune stage skipped by user request.")
|
||||
return
|
||||
import shutil
|
||||
shutil.rmtree(finetune_output_dir, ignore_errors=True)
|
||||
logger.info("Old finetune directory removed.")
|
||||
|
||||
run_dpo_finetune(config, experiment_run_dir)
|
||||
except Exception as e:
|
||||
logger.error("An error occurred during finetuning: %s", e, exc_info=True)
|
||||
finally:
|
||||
finetune_duration = datetime.datetime.now() - finetune_start_time
|
||||
logger.info("Total finetuning duration: %s", finetune_duration)
|
||||
else:
|
||||
logger.warning("Skipping finetuning as the main pipeline did not complete successfully or experiment directory is not set.")
|
||||
else:
|
||||
logger.info("inetuning is disabled by config/CLI or due to pipeline issues.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
103
regex-bench.py
Normal file
103
regex-bench.py
Normal file
@@ -0,0 +1,103 @@
|
||||
#!/usr/bin/env python
|
||||
"""
|
||||
regex_mt_bench.py – does Python's re.search() scale across threads?
|
||||
|
||||
• Creates a single master regex (similar to the RegexValidator patch).
|
||||
• Launches N threads; each thread calls .search() a fixed number of times.
|
||||
• Reports wall-time vs process-CPU-time so you can see how many cores
|
||||
the regex engine actually used.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import regex as re, threading, time, os, sys, random, math
|
||||
|
||||
# ---------------------------------------------------------------------
|
||||
# 1. Prepare synthetic workload
|
||||
# ---------------------------------------------------------------------
|
||||
NUM_PATTERNS = 120 # similar to your real list
|
||||
ITERATIONS = 80_000 # 16× more work than before
|
||||
TEXT_LEN_CHARS = 800_000 # force the engine to read a lot
|
||||
|
||||
random.seed(42)
|
||||
|
||||
# Make deterministic-ish patterns: a literal word or a short .* wildcard
|
||||
_PATTERNS: list[str] = []
|
||||
for i in range(NUM_PATTERNS):
|
||||
if i % 3 == 0:
|
||||
_PATTERNS.append(fr"\bword{i}\b")
|
||||
elif i % 3 == 1:
|
||||
_PATTERNS.append(fr"phrase{i}[^ ]+end")
|
||||
else:
|
||||
_PATTERNS.append(fr"token{i}.*?token{i+1}")
|
||||
|
||||
# Build a master alternation with named groups (as in the patch)
|
||||
parts, _group2raw = [], {}
|
||||
for i, p in enumerate(_PATTERNS):
|
||||
gname = f"P{i}"
|
||||
parts.append(f"(?P<{gname}>{p})")
|
||||
_group2raw[gname] = p
|
||||
_BIG_RE = re.compile("|".join(parts), re.IGNORECASE | re.MULTILINE | re.DOTALL)
|
||||
|
||||
# Generate text that *sometimes* matches: sprinkle keywords every ~1000 chars
|
||||
_chunks = []
|
||||
for i in range(TEXT_LEN_CHARS // 50):
|
||||
if i % 20 == 0: # every 20th chunk drop a keyword
|
||||
k = random.randrange(NUM_PATTERNS)
|
||||
tok = f"word{k}" if k % 3 == 0 else f"phrase{k}xxend"
|
||||
_chunks.append(tok)
|
||||
else:
|
||||
_chunks.append("loremipsum")
|
||||
_TEXT = " ".join(_chunks)
|
||||
|
||||
# ---------------------------------------------------------------------
|
||||
# 2. Benchmark helper
|
||||
# ---------------------------------------------------------------------
|
||||
def run_threads(n_threads: int) -> tuple[float, float]:
|
||||
"""
|
||||
Launch n_threads that each call _BIG_RE.search(_TEXT) ITERATIONS times.
|
||||
|
||||
Returns (wall_seconds, cpu_seconds) for the whole job.
|
||||
"""
|
||||
def worker():
|
||||
s = _BIG_RE # local var for speed
|
||||
t = _TEXT
|
||||
for _ in range(ITERATIONS):
|
||||
s.search(t)
|
||||
|
||||
threads = [threading.Thread(target=worker, daemon=True)
|
||||
for _ in range(n_threads)]
|
||||
|
||||
cpu_start = os.times() # returns a 5-tuple
|
||||
t0 = time.perf_counter()
|
||||
|
||||
for th in threads:
|
||||
th.start()
|
||||
for th in threads:
|
||||
th.join()
|
||||
|
||||
wall = time.perf_counter() - t0
|
||||
cpu = (os.times().user + os.times().system) - (cpu_start.user + cpu_start.system)
|
||||
return wall, cpu
|
||||
|
||||
# ---------------------------------------------------------------------
|
||||
# 3. Run for several thread counts
|
||||
# ---------------------------------------------------------------------
|
||||
def main():
|
||||
print(f"patterns : {NUM_PATTERNS}")
|
||||
print(f"text length (chars) : {len(_TEXT):,}")
|
||||
print(f"regex searches/thread: {ITERATIONS}")
|
||||
print()
|
||||
|
||||
for n in (1, 2, 4, 8):
|
||||
wall, cpu = run_threads(n)
|
||||
util = cpu / wall if wall else math.nan
|
||||
print(f"{n:>2} threads → wall {wall:6.2f} s "
|
||||
f"CPU {cpu:6.2f} s ratio {util:4.2f}")
|
||||
|
||||
print("\nInterpretation:")
|
||||
print(" • ratio ≈ 1.0 → work is effectively single-core.")
|
||||
print(" • ratio → N → regex scanning scales across N cores.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
38
requirements.txt
Normal file
38
requirements.txt
Normal file
@@ -0,0 +1,38 @@
|
||||
# Core dependencies
|
||||
pyyaml>=6.0
|
||||
pandas>=1.5
|
||||
numpy>=1.20
|
||||
nltk>=3.6
|
||||
requests>=2.25
|
||||
tqdm>=4.60
|
||||
tiktoken
|
||||
wordfreq>=3.0
|
||||
regex
|
||||
flash-attn
|
||||
dotenv
|
||||
scipy
|
||||
|
||||
# For vLLM (if managed by this script)
|
||||
# vllm # User should install this separately if managing vLLM manually or if specific version needed
|
||||
# torch # Or newer, compatible with vLLM and Unsloth
|
||||
|
||||
|
||||
|
||||
# For DPO Finetuning (Unsloth and its dependencies)
|
||||
unsloth
|
||||
bitsandbytes
|
||||
accelerate
|
||||
peft
|
||||
trl
|
||||
transformers
|
||||
datasets
|
||||
sentencepiece
|
||||
protobuf
|
||||
hf_transfer
|
||||
tensorboard
|
||||
|
||||
|
||||
|
||||
|
||||
# Note: Users will need to ensure compatible versions of torch, CUDA, and vLLM/Unsloth
|
||||
# are installed for their specific hardware if using vLLM management or finetuning.
|
||||
71
test_inference.py
Normal file
71
test_inference.py
Normal file
@@ -0,0 +1,71 @@
|
||||
|
||||
import os
|
||||
import os
|
||||
import torch
|
||||
from pathlib import Path
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
def find_latest_finetuned_model() -> str | None:
|
||||
"""
|
||||
Look under:
|
||||
1. <script dir>/results/auto_antislop_runs
|
||||
2. /results/auto_antislop_runs
|
||||
Return the most recent `…/finetuned_model*/merged_16bit` directory or None.
|
||||
"""
|
||||
candidate_bases = [
|
||||
Path(__file__).resolve().parent / "results" / "auto_antislop_runs",
|
||||
Path("/results/auto_antislop_runs"),
|
||||
]
|
||||
|
||||
latest: tuple[float, Path] | None = None
|
||||
for base in candidate_bases:
|
||||
if not base.is_dir():
|
||||
continue
|
||||
|
||||
# run_*/finetuned_model*/merged_16bit
|
||||
for merged_dir in base.glob("run_*/finetuned_model*/merged_16bit"):
|
||||
if not merged_dir.is_dir():
|
||||
continue
|
||||
mtime = merged_dir.parent.stat().st_mtime # use finetuned_model* dir mtime
|
||||
if latest is None or mtime > latest[0]:
|
||||
latest = (mtime, merged_dir.resolve())
|
||||
|
||||
return str(latest[1]) if latest else None
|
||||
|
||||
|
||||
model_path = find_latest_finetuned_model() or "."
|
||||
print(f"Loading model from: {os.path.abspath(model_path)}")
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
print(f"Using device: {device}")
|
||||
|
||||
try:
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_path,
|
||||
torch_dtype=torch.bfloat16,
|
||||
device_map="auto" if device == "cuda" else None,
|
||||
)
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a creative storyteller."},
|
||||
{"role": "user", "content": "Write a short, engaging story about a princess."}
|
||||
]
|
||||
prompt = tokenizer.apply_chat_template(messages, tokenize=False)
|
||||
print("\nApplied chat template:\n", prompt)
|
||||
|
||||
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
|
||||
generated_ids = model.generate(
|
||||
input_ids,
|
||||
max_new_tokens=500,
|
||||
do_sample=True,
|
||||
temperature=0.7,
|
||||
top_p=0.9,
|
||||
)
|
||||
generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
|
||||
response = generated_text[len(tokenizer.decode(input_ids[0], skip_special_tokens=True)):]
|
||||
print("\n--- Generated Story ---\n", response)
|
||||
print("\nToken count (approximate):", len(generated_ids[0]) - len(input_ids[0]))
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
0
utils/__init__.py
Normal file
0
utils/__init__.py
Normal file
256
utils/config_loader.py
Normal file
256
utils/config_loader.py
Normal file
@@ -0,0 +1,256 @@
|
||||
# utils/config_loader.py
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Sequence, Any
|
||||
|
||||
import yaml
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
_ALWAYS: Sequence[str] = (
|
||||
# minimal required keys for the pipeline to run at all
|
||||
"experiment_base_dir",
|
||||
"human_profile_path",
|
||||
"num_iterations",
|
||||
"min_word_len_for_analysis",
|
||||
"log_level",
|
||||
"model_id",
|
||||
)
|
||||
|
||||
_VLLM: Sequence[str] = (
|
||||
"vllm_model_id",
|
||||
"vllm_port",
|
||||
"vllm_hf_token",
|
||||
"vllm_cuda_visible_devices",
|
||||
"vllm_gpu_memory_utilization",
|
||||
"vllm_max_model_len",
|
||||
"vllm_dtype",
|
||||
"vllm_extra_args",
|
||||
"vllm_env",
|
||||
)
|
||||
|
||||
_GENERATION: Sequence[str] = (
|
||||
"generation_api_key",
|
||||
"generation_api_base_url", # needed if you do local or remote calls
|
||||
"generation_model_id",
|
||||
"generation_max_new_tokens",
|
||||
"generation_threads",
|
||||
"generation_max_prompts",
|
||||
"generation_hf_dataset_name",
|
||||
"generation_hf_dataset_split",
|
||||
"generation_logging_level",
|
||||
"generation_chat_template_model_id",
|
||||
"generation_param_chunk_size",
|
||||
"generation_param_top_logprobs_count",
|
||||
"generation_param_temperature",
|
||||
"generation_param_top_p",
|
||||
"generation_param_top_k",
|
||||
"generation_param_min_p",
|
||||
"generation_param_timeout",
|
||||
"generation_param_stop_sequences",
|
||||
"generation_ngram_remove_stopwords",
|
||||
"generation_ngram_language",
|
||||
"generation_force_backtrack",
|
||||
"generation_prompt_template",
|
||||
"generation_system_prompt"
|
||||
|
||||
)
|
||||
|
||||
_NGRAM: Sequence[str] = (
|
||||
"top_k_bigrams",
|
||||
"top_k_trigrams",
|
||||
"dict_bigrams_initial",
|
||||
"dict_bigrams_subsequent",
|
||||
"nodict_bigrams_initial",
|
||||
"nodict_bigrams_subsequent",
|
||||
"dict_trigrams_initial",
|
||||
"dict_trigrams_subsequent",
|
||||
"nodict_trigrams_initial",
|
||||
"nodict_trigrams_subsequent",
|
||||
"extra_ngrams_to_ban",
|
||||
)
|
||||
|
||||
_SLOP: Sequence[str] = (
|
||||
"min_phrase_freq_to_keep",
|
||||
"top_n_initial_slop_ban",
|
||||
"top_n_subsequent_slop_ban",
|
||||
"extra_slop_phrases_to_ban",
|
||||
"banned_slop_phrases_filename",
|
||||
)
|
||||
|
||||
_OVERREP: Sequence[str] = (
|
||||
"top_k_words_for_overrep_analysis",
|
||||
"dict_overrep_initial",
|
||||
"dict_overrep_subsequent",
|
||||
"nodict_overrep_initial",
|
||||
"nodict_overrep_subsequent",
|
||||
)
|
||||
|
||||
_FINETUNE: Sequence[str] = (
|
||||
"finetune_mode",
|
||||
"finetune_ftpo_dataset",
|
||||
"finetune_base_model_id",
|
||||
"finetune_max_seq_length",
|
||||
"finetune_load_in_4bit",
|
||||
"finetune_lora_r",
|
||||
"finetune_lora_alpha",
|
||||
"finetune_lora_dropout",
|
||||
"finetune_weight_decay",
|
||||
"finetune_target_modules",
|
||||
"finetune_gradient_checkpointing",
|
||||
"finetune_chat_template",
|
||||
"finetune_batch_size",
|
||||
"finetune_gradient_accumulation_steps",
|
||||
"finetune_warmup_ratio",
|
||||
"finetune_num_epochs",
|
||||
"finetune_learning_rate",
|
||||
"finetune_auto_learning_rate",
|
||||
"finetune_beta",
|
||||
"finetune_output_dir_suffix",
|
||||
"finetune_save_merged_16bit",
|
||||
"finetune_save_gguf_q8_0",
|
||||
"finetune_max_train_examples",
|
||||
"finetune_cuda_visible_devices",
|
||||
"ftpo_sample_rejected_regularisation_strength",
|
||||
"ftpo_sample_chosen_regularisation_strength",
|
||||
"ftpo_sample_min_chosen_tokens",
|
||||
)
|
||||
|
||||
|
||||
def _deep_update(dst: Dict, src: Dict) -> Dict:
|
||||
"""Recursively merge src into dst (src wins)."""
|
||||
for k, v in src.items():
|
||||
if k in dst and isinstance(v, dict) and isinstance(dst[k], dict):
|
||||
_deep_update(dst[k], v)
|
||||
else:
|
||||
dst[k] = copy.deepcopy(v)
|
||||
return dst
|
||||
|
||||
def load_pipeline_config(config_path: Path) -> Dict[str, Any]:
|
||||
"""Load config from a YAML file, or return empty dict if missing/invalid."""
|
||||
if config_path and config_path.exists():
|
||||
try:
|
||||
with config_path.open('r', encoding='utf-8') as f:
|
||||
data = yaml.safe_load(f) or {}
|
||||
logger.info("Loaded configuration from %s", config_path)
|
||||
return data
|
||||
except Exception as e:
|
||||
logger.error("Could not load %s: %s – using empty config", config_path, e)
|
||||
else:
|
||||
logger.info("Config file %s not found – using empty config", config_path)
|
||||
return {}
|
||||
|
||||
def merge_config_with_cli_args(config: Dict[str, Any], cli_args: argparse.Namespace) -> Dict[str, Any]:
|
||||
"""
|
||||
Merges every possible CLI parameter from your old DEFAULT_CONFIG
|
||||
into 'config' if the user actually provided it (i.e. it's not None).
|
||||
Also merges housekeeping flags (config_file, resume_from_dir, log_level).
|
||||
"""
|
||||
merged = copy.deepcopy(config)
|
||||
|
||||
# 1. Housekeeping arguments (not originally in DEFAULT_CONFIG, but we keep them if set)
|
||||
if getattr(cli_args, 'config_file', None) is not None:
|
||||
merged['config_file'] = cli_args.config_file
|
||||
if getattr(cli_args, 'resume_from_dir', None) is not None:
|
||||
merged['resume_from_dir'] = cli_args.resume_from_dir
|
||||
if getattr(cli_args, 'log_level', None) is not None:
|
||||
merged['log_level'] = cli_args.log_level
|
||||
|
||||
# 2. Booleans that map from CLI flags to known keys in config
|
||||
if getattr(cli_args, 'run_finetune', None) is not None:
|
||||
merged['finetune_enabled'] = cli_args.run_finetune
|
||||
if getattr(cli_args, 'manage_vllm', None) is not None:
|
||||
merged['manage_vllm'] = cli_args.manage_vllm
|
||||
if getattr(cli_args, 'generation_step_enabled', None) is not None:
|
||||
merged['generation_step_enabled'] = cli_args.generation_step_enabled
|
||||
if getattr(cli_args, "finetune_cuda_visible_devices", None) is not None:
|
||||
merged["finetune_cuda_visible_devices"] = cli_args.finetune_cuda_visible_devices
|
||||
|
||||
|
||||
# 3. All remaining keys from the old DEFAULT_CONFIG
|
||||
_all_groups: Sequence[Sequence[str]] = (
|
||||
_ALWAYS,
|
||||
_VLLM,
|
||||
_GENERATION,
|
||||
_NGRAM,
|
||||
_SLOP,
|
||||
_OVERREP,
|
||||
_FINETUNE,
|
||||
)
|
||||
all_config_keys: List[str] = [k for group in _all_groups for k in group]
|
||||
|
||||
# 4. fallback for per-stage model IDs <-- add this block
|
||||
for key in (
|
||||
"vllm_model_id",
|
||||
"generation_model_id",
|
||||
"generation_chat_template_model_id",
|
||||
"finetune_base_model_id",
|
||||
):
|
||||
if not merged.get(key): # None, "", or missing
|
||||
merged[key] = merged.get("model_id")
|
||||
|
||||
# Overwrite config if user specified a value
|
||||
for key in all_config_keys:
|
||||
cli_val = getattr(cli_args, key, None)
|
||||
if cli_val is not None:
|
||||
merged[key] = cli_val
|
||||
|
||||
return merged
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Validate with partial requirements depending on which features are enabled
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _missing(cfg: Dict[str, Any], keys: Sequence[str]) -> List[str]:
|
||||
return [k for k in keys if k not in cfg or cfg[k] is None]
|
||||
|
||||
def validate_config(cfg: Dict[str, Any]) -> None:
|
||||
"""Raise ValueError if any required config is missing based on pipeline flags."""
|
||||
missing = []
|
||||
# always
|
||||
missing.extend(_missing(cfg, _ALWAYS))
|
||||
|
||||
# vllm
|
||||
if cfg.get("manage_vllm", False):
|
||||
missing.extend(_missing(cfg, _VLLM))
|
||||
|
||||
# generation
|
||||
if cfg.get("generation_step_enabled", True):
|
||||
missing.extend(_missing(cfg, _GENERATION))
|
||||
|
||||
# n-gram ban
|
||||
if cfg.get("enable_ngram_ban", False):
|
||||
missing.extend(_missing(cfg, _NGRAM))
|
||||
|
||||
# slop phrase ban
|
||||
if cfg.get("enable_slop_phrase_ban", False):
|
||||
missing.extend(_missing(cfg, _SLOP))
|
||||
|
||||
# over-rep analysis
|
||||
if cfg.get("compute_overrep_words", False):
|
||||
missing.extend(_missing(cfg, _OVERREP))
|
||||
|
||||
# finetuning
|
||||
if cfg.get("finetune_enabled", False):
|
||||
missing.extend(_missing(cfg, _FINETUNE))
|
||||
|
||||
if missing:
|
||||
raise ValueError(
|
||||
f"Configuration is incomplete; missing these keys: {', '.join(sorted(set(missing)))}"
|
||||
)
|
||||
logger.info("Configuration validated – all required keys present (for enabled features).")
|
||||
|
||||
def load_merge_validate(config_path: Path, cli_args: argparse.Namespace) -> Dict[str, Any]:
|
||||
"""
|
||||
1) Load YAML from config_path,
|
||||
2) Merge in any CLI flags user typed,
|
||||
3) Validate that all needed keys for enabled features are present.
|
||||
"""
|
||||
cfg = load_pipeline_config(config_path)
|
||||
cfg = merge_config_with_cli_args(cfg, cli_args)
|
||||
validate_config(cfg)
|
||||
return cfg
|
||||
290
utils/dataset_helpers.py
Normal file
290
utils/dataset_helpers.py
Normal file
@@ -0,0 +1,290 @@
|
||||
# utils/dataset_helpers.py
|
||||
from __future__ import annotations
|
||||
import logging, os
|
||||
from pathlib import Path
|
||||
from collections import Counter, defaultdict
|
||||
from typing import Collection, Optional
|
||||
from datetime import datetime, timezone
|
||||
import json
|
||||
|
||||
import numpy as np
|
||||
from datasets import load_dataset
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# tokens we want to watch closely
|
||||
_WATCH = [" nodded", " leaned"]
|
||||
|
||||
|
||||
def load_ftpo_multi_dataset(
|
||||
path: Path,
|
||||
tokenizer,
|
||||
*,
|
||||
experiment_run_dir: Path | None = None,
|
||||
max_seq_len: int = 4096,
|
||||
rejected_reg_strength: float = 0.0,
|
||||
chosen_reg_strength: float = 0.0,
|
||||
min_chosen_tokens: int = 1,
|
||||
max_train_examples: int | None = None,
|
||||
stop_words: Optional[Collection[str]] = None,
|
||||
num_proc: int | None = None,
|
||||
batch_size: int = 512,
|
||||
):
|
||||
"""
|
||||
Parallel loader for “multi-chosen” FTPO JSONL with dual regularisation.
|
||||
Logs the counts of `_WATCH` tokens at every major stage.
|
||||
"""
|
||||
|
||||
if min_chosen_tokens < 1:
|
||||
min_chosen_tokens = 1
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# helpers
|
||||
# ------------------------------------------------------------------
|
||||
rng = np.random.default_rng(3407)
|
||||
|
||||
def _median_threshold(cts: Counter[str], strength: float) -> dict[str, float]:
|
||||
if not cts or strength <= 0:
|
||||
return {}
|
||||
med = float(np.median(list(cts.values())))
|
||||
return {t: 1.0 if c <= med else (med / c) ** strength for t, c in cts.items()}
|
||||
|
||||
def _log_top(cts: Counter[str], what: str) -> None:
|
||||
head = ", ".join(f"{tok!r}:{cnt}" for tok, cnt in cts.most_common(20))
|
||||
logger.info(f"[ftpo-loader] {what} top-20 → {head}")
|
||||
logger.info(
|
||||
" ↳ watch «%s»: %s «%s»: %s",
|
||||
_WATCH[0], cts[_WATCH[0]],
|
||||
_WATCH[1], cts[_WATCH[1]],
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# stop-word list (unchanged)
|
||||
# ------------------------------------------------------------------
|
||||
if stop_words is None:
|
||||
stop_words = {
|
||||
"the","a","an","in","on","at","by","for","to","of","and","or","but",
|
||||
"if","then","else","when","where","how","why","what","who","whom",
|
||||
"this","that","these","those","is","are","was","were","be","being",
|
||||
"been","have","has","had","do","does","did","will","would","shall",
|
||||
"should","can","could","may","might","must"
|
||||
}
|
||||
stop_words = {w.lower() for w in stop_words}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 0️⃣ raw load + shuffle
|
||||
# ------------------------------------------------------------------
|
||||
raw = load_dataset("json", data_files=str(path), split="train").shuffle(seed=3407)
|
||||
rows = list(raw)
|
||||
if not rows:
|
||||
raise ValueError(f"{path} contained no rows")
|
||||
|
||||
rej_counts = Counter(r["rejected_decoded"] for r in rows)
|
||||
_log_top(rej_counts, "BEFORE")
|
||||
|
||||
# ────────────────────────────────────────────────────────────────
|
||||
# 1️⃣ Capture ORIGINAL rejected-token distribution & ratios
|
||||
# (no rows removed, no chosen trimming yet)
|
||||
# ────────────────────────────────────────────────────────────────
|
||||
rej_cts_orig = Counter(r["rejected_decoded"] for r in rows)
|
||||
_log_top(rej_cts_orig, "PRE-NORMALISATION")
|
||||
|
||||
# convert to fractional “weights” via median-threshold regularisation
|
||||
med = float(np.median(list(rej_cts_orig.values())))
|
||||
w_rej = {tok: 1.0 if c <= med else (med / c) ** rejected_reg_strength
|
||||
for tok, c in rej_cts_orig.items()}
|
||||
|
||||
# normalised ratios we *want* to keep in the final dataset
|
||||
total_weighted = sum(w_rej[t] * c for t, c in rej_cts_orig.items())
|
||||
ratio_rej = {tok: (w_rej[tok] * cnt) / total_weighted
|
||||
for tok, cnt in rej_cts_orig.items()}
|
||||
|
||||
# ────────────────────────────────────────────────────────────────
|
||||
# 2️⃣ Chosen-token trimming (build quotas *before* we cut)
|
||||
# ────────────────────────────────────────────────────────────────
|
||||
chosen_cts_orig = Counter(tok
|
||||
for r in rows
|
||||
for tok in (r["multi_chosen_decoded"] or []))
|
||||
|
||||
_log_top(chosen_cts_orig, "ORIGINAL CHOSEN TOKENS")
|
||||
|
||||
# Trim the peak: cap top tokens to match the 10th highest count
|
||||
if len(chosen_cts_orig) >= 10:
|
||||
top_counts = sorted(chosen_cts_orig.values(), reverse=True)
|
||||
cap_value = top_counts[9] # 10th highest count
|
||||
chosen_cts_capped = Counter()
|
||||
for tok, cnt in chosen_cts_orig.items():
|
||||
chosen_cts_capped[tok] = min(cnt, cap_value)
|
||||
else:
|
||||
chosen_cts_capped = chosen_cts_orig.copy()
|
||||
|
||||
# Now calculate regularization on the capped distribution
|
||||
med_chosen = float(np.median(list(chosen_cts_capped.values())))
|
||||
w_chosen = {tok: 1.0 if c <= med_chosen
|
||||
else (med_chosen / c) ** chosen_reg_strength
|
||||
for tok, c in chosen_cts_capped.items()}
|
||||
|
||||
tgt_chosen = {tok: int(round(c * w_chosen.get(tok, 1.0)))
|
||||
for tok, c in chosen_cts_capped.items()}
|
||||
|
||||
# Log the target quotas
|
||||
quota_items = sorted(tgt_chosen.items(), key=lambda x: x[1], reverse=True)[:20]
|
||||
quota_str = ", ".join(f"{tok!r}:{quota}" for tok, quota in quota_items)
|
||||
logger.info(f"[ftpo-loader] CHOSEN TARGET QUOTAS top-20 → {quota_str}")
|
||||
logger.info(
|
||||
" ↳ watch quotas «%s»: %s (was %s) «%s»: %s (was %s)",
|
||||
_WATCH[0], tgt_chosen.get(_WATCH[0], 0), chosen_cts_orig.get(_WATCH[0], 0),
|
||||
_WATCH[1], tgt_chosen.get(_WATCH[1], 0), chosen_cts_orig.get(_WATCH[1], 0),
|
||||
)
|
||||
|
||||
_log_top(Counter(r["rejected_decoded"] for r in rows), "POST-CHOSEN")
|
||||
|
||||
# ────────────────────────────────────────────────────────────────
|
||||
# 3️⃣ Apply min_chosen_tokens row filter
|
||||
# ────────────────────────────────────────────────────────────────
|
||||
rows = [r for r in rows if len(r["multi_chosen_decoded"]) >= min_chosen_tokens]
|
||||
|
||||
_log_top(Counter(r["rejected_decoded"] for r in rows), "POST-MIN-FILTER")
|
||||
|
||||
# ────────────────────────────────────────────────────────────────
|
||||
# 4️⃣ Row-level quota sampling **now** that trimming & filtering
|
||||
# are done. Scale the original ratios to the remaining size.
|
||||
# ────────────────────────────────────────────────────────────────
|
||||
N_final = max_train_examples or len(rows)
|
||||
target_rows = {tok: int(round(ratio_rej[tok] * N_final))
|
||||
for tok in ratio_rej}
|
||||
|
||||
rng.shuffle(rows)
|
||||
selected, seen = [], defaultdict(int)
|
||||
selected_indices = set() # Track indices instead of row objects
|
||||
|
||||
# First pass: try to fill quotas
|
||||
for i, r in enumerate(rows):
|
||||
tok = r["rejected_decoded"]
|
||||
if seen[tok] < target_rows.get(tok, 0):
|
||||
selected.append(r)
|
||||
selected_indices.add(i)
|
||||
seen[tok] += 1
|
||||
if len(selected) >= N_final:
|
||||
break
|
||||
|
||||
# Second pass: if still short, keep adding while maintaining proportions
|
||||
if len(selected) < N_final:
|
||||
# Build index of remaining rows by token
|
||||
remaining_by_token = defaultdict(list)
|
||||
for i, r in enumerate(rows):
|
||||
if i not in selected_indices:
|
||||
remaining_by_token[r["rejected_decoded"]].append((i, r))
|
||||
|
||||
# Keep adding until we reach N_final
|
||||
while len(selected) < N_final:
|
||||
# Find token that's furthest below its target ratio AND has rows available
|
||||
best_tok = None
|
||||
best_ratio_diff = -1
|
||||
|
||||
for tok, available_rows in remaining_by_token.items():
|
||||
if not available_rows: # Skip tokens with no remaining rows
|
||||
continue
|
||||
|
||||
current_ratio = seen[tok] / len(selected) if len(selected) > 0 else 0
|
||||
target_ratio = ratio_rej.get(tok, 0)
|
||||
ratio_diff = target_ratio - current_ratio
|
||||
|
||||
if ratio_diff > best_ratio_diff:
|
||||
best_ratio_diff = ratio_diff
|
||||
best_tok = tok
|
||||
|
||||
# If no tokens have available rows, we're done
|
||||
if best_tok is None:
|
||||
break
|
||||
|
||||
# Add one row for the most underrepresented token
|
||||
idx, r = remaining_by_token[best_tok].pop()
|
||||
selected.append(r)
|
||||
selected_indices.add(idx)
|
||||
seen[best_tok] += 1
|
||||
|
||||
rows = selected
|
||||
|
||||
# ── Dump the final row subset exactly as it was read (no tokenisation) ──
|
||||
if experiment_run_dir is not None:
|
||||
ts = datetime.now(timezone.utc).astimezone()\
|
||||
.strftime("%Y-%m-%d_%H-%M-%S")
|
||||
dump_file = experiment_run_dir / f"ftpo_training_set_used_{ts}.jsonl"
|
||||
try:
|
||||
with open(dump_file, "w", encoding="utf-8") as fh:
|
||||
for r in rows:
|
||||
fh.write(json.dumps(r, ensure_ascii=False) + "\n")
|
||||
logger.info("[ftpo-loader] dumped %d rows → %s", len(rows), dump_file)
|
||||
except Exception as e:
|
||||
logger.warning("[ftpo-loader] failed to dump training rows: %s", e)
|
||||
|
||||
|
||||
|
||||
_log_top(Counter(r["rejected_decoded"] for r in rows), "AFTER-SAMPLING")
|
||||
logger.info("[ftpo-loader] kept %d rows after quota sampling", len(rows))
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 5️⃣ tokenisation (unchanged section)
|
||||
# ------------------------------------------------------------------
|
||||
from datasets import Dataset
|
||||
ds = Dataset.from_list(rows)
|
||||
|
||||
tokenizer.truncation_side = "left"
|
||||
num_proc = num_proc or max(1, int(os.cpu_count() / 4))
|
||||
|
||||
def _tok(batch):
|
||||
out_prompt, out_chosen, out_rej, out_valid = [], [], [], []
|
||||
|
||||
prompt_tok = tokenizer(
|
||||
batch["context_with_chat_template"],
|
||||
add_special_tokens=False,
|
||||
truncation=False,
|
||||
return_attention_mask=False,
|
||||
).input_ids
|
||||
|
||||
for p_ids, chosen_surf, rej_surf in zip(
|
||||
prompt_tok, batch["multi_chosen_decoded"], batch["rejected_decoded"]
|
||||
):
|
||||
chosen_surf = chosen_surf or []
|
||||
chosen_tok_ids = [tokenizer(t, add_special_tokens=False).input_ids
|
||||
for t in chosen_surf]
|
||||
rej_tok_ids = tokenizer(rej_surf, add_special_tokens=False).input_ids
|
||||
|
||||
valid = (
|
||||
chosen_tok_ids
|
||||
and all(len(t) == 1 for t in chosen_tok_ids)
|
||||
and len(rej_tok_ids) == 1
|
||||
and rej_surf.strip().lower() not in stop_words
|
||||
and len(p_ids) + 1 <= max_seq_len
|
||||
)
|
||||
if valid and rej_tok_ids[0] in [t[0] for t in chosen_tok_ids]:
|
||||
valid = False
|
||||
|
||||
out_valid.append(valid)
|
||||
if valid:
|
||||
out_prompt.append(p_ids)
|
||||
out_chosen.append([t[0] for t in chosen_tok_ids])
|
||||
out_rej.append(rej_tok_ids[0])
|
||||
else:
|
||||
out_prompt.append([0]); out_chosen.append([0]); out_rej.append(0)
|
||||
|
||||
return {
|
||||
"prompt_ids": out_prompt,
|
||||
"chosen_ids": out_chosen,
|
||||
"rejected_token_id": out_rej,
|
||||
"__valid": out_valid,
|
||||
}
|
||||
|
||||
ds = ds.map(
|
||||
_tok, batched=True, batch_size=batch_size,
|
||||
remove_columns=ds.column_names,
|
||||
num_proc=num_proc, desc="tokenising",
|
||||
)
|
||||
ds = ds.filter(lambda ex: ex["__valid"], num_proc=num_proc, desc="filter")
|
||||
ds = ds.remove_columns("__valid")
|
||||
if len(ds) == 0:
|
||||
raise ValueError("no ftpo samples survived length / sanity checks")
|
||||
|
||||
return ds.shuffle(seed=3407)
|
||||
116
utils/fix_gemma.py
Normal file
116
utils/fix_gemma.py
Normal file
@@ -0,0 +1,116 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Gemma-3 key repair utility
|
||||
==========================
|
||||
|
||||
Repairs checkpoints whose weight names are in either of the two incorrect
|
||||
forms:
|
||||
|
||||
1. model.language_model.embed_tokens.weight (# leading "model.")
|
||||
2. language_model.embed_tokens.weight (# missing ".model.")
|
||||
|
||||
to the correct form:
|
||||
|
||||
language_model.model.embed_tokens.weight
|
||||
|
||||
Usage:
|
||||
python repair_gemma3_keys.py /path/to/checkpoint_dir
|
||||
"""
|
||||
|
||||
import sys, json, shutil
|
||||
from pathlib import Path
|
||||
from safetensors.torch import safe_open, save_file
|
||||
|
||||
BAD_LEADING = "model." # variant 1
|
||||
GOOD_PREFIX = "language_model."
|
||||
GOOD_FULL = "language_model.model."
|
||||
OUT_SUFFIX = "_repaired"
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# key transformation ----------------------------------------------------
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
def fix_key(key: str) -> str:
|
||||
"""
|
||||
1) strip a leading "model." if present
|
||||
2) ensure "language_model." is followed by "model."
|
||||
"""
|
||||
# step 1 – drop wrapper prefix once
|
||||
if key.startswith(BAD_LEADING):
|
||||
key = key[len(BAD_LEADING):]
|
||||
|
||||
# step 2 – insert ".model." if missing
|
||||
if key.startswith(GOOD_PREFIX) and not key.startswith(GOOD_FULL):
|
||||
key = GOOD_FULL + key[len(GOOD_PREFIX):]
|
||||
|
||||
return key
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# shard processing ------------------------------------------------------
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
def repair_shard(src: Path, dst: Path) -> None:
|
||||
"""
|
||||
Re-write a .safetensors shard with corrected keys.
|
||||
"""
|
||||
corrected = {}
|
||||
|
||||
with safe_open(src, framework="pt", device="cpu") as f:
|
||||
for old_key in f.keys():
|
||||
corrected[fix_key(old_key)] = f.get_tensor(old_key)
|
||||
|
||||
save_file(corrected, dst, metadata={"format": "pt"})
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# driver ----------------------------------------------------------------
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
def main():
|
||||
if len(sys.argv) != 2:
|
||||
sys.exit("Usage: repair_gemma3_keys.py <checkpoint_dir>")
|
||||
|
||||
src_dir = Path(sys.argv[1]).expanduser().resolve()
|
||||
if not src_dir.is_dir():
|
||||
sys.exit(f"Directory not found: {src_dir}")
|
||||
|
||||
out_dir = src_dir.with_name(src_dir.name + OUT_SUFFIX)
|
||||
out_dir.mkdir(exist_ok=True)
|
||||
|
||||
index_path = src_dir / "model.safetensors.index.json"
|
||||
if not index_path.is_file():
|
||||
sys.exit("model.safetensors.index.json not found in checkpoint dir.")
|
||||
|
||||
# ---- load index ----------------------------------------------------
|
||||
with open(index_path, "r") as f:
|
||||
index = json.load(f)
|
||||
|
||||
new_weight_map = {}
|
||||
processed_shards = set()
|
||||
|
||||
# ---- process every tensor key -------------------------------------
|
||||
for old_key, shard_name in index["weight_map"].items():
|
||||
new_key = fix_key(old_key)
|
||||
new_weight_map[new_key] = shard_name
|
||||
|
||||
if shard_name in processed_shards:
|
||||
continue
|
||||
processed_shards.add(shard_name)
|
||||
repair_shard(src_dir / shard_name, out_dir / shard_name)
|
||||
|
||||
index["weight_map"] = new_weight_map
|
||||
|
||||
# ---- write new index ----------------------------------------------
|
||||
with open(out_dir / "model.safetensors.index.json", "w") as f:
|
||||
json.dump(index, f, indent=2)
|
||||
|
||||
# ---- copy auxiliary files -----------------------------------------
|
||||
for fp in src_dir.iterdir():
|
||||
if fp.name == "model.safetensors.index.json" or fp.suffix == ".safetensors":
|
||||
continue
|
||||
shutil.copy2(fp, out_dir / fp.name)
|
||||
|
||||
print(f"✓ Repaired checkpoint written to {out_dir}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
150
utils/fs_helpers.py
Normal file
150
utils/fs_helpers.py
Normal file
@@ -0,0 +1,150 @@
|
||||
import nltk
|
||||
import datetime
|
||||
from pathlib import Path
|
||||
import logging
|
||||
import sys
|
||||
import json
|
||||
import shutil
|
||||
from typing import List
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
def set_from_json(path: Path) -> set[str]:
|
||||
"""Return a hashable set of strings from a ban-list file that may be
|
||||
either ["foo bar", …] or [["foo bar", 1], …]."""
|
||||
if not path or not path.is_file():
|
||||
return set()
|
||||
try:
|
||||
raw = json.loads(path.read_text("utf-8"))
|
||||
except Exception:
|
||||
return set()
|
||||
|
||||
if not isinstance(raw, list):
|
||||
return set()
|
||||
|
||||
out = set()
|
||||
for item in raw:
|
||||
if isinstance(item, list) and item: # slop format
|
||||
out.add(str(item[0]))
|
||||
else: # plain string
|
||||
out.add(str(item))
|
||||
return out
|
||||
|
||||
|
||||
def merge_custom_bans_into_file(path: Path, extra_items: List[str]) -> None:
|
||||
"""Merge extra_items into `path`, preserving original on-disk format."""
|
||||
|
||||
# 1) read whatever is already there
|
||||
try:
|
||||
current_raw = json.loads(path.read_text("utf-8")) if path.exists() else []
|
||||
except Exception:
|
||||
current_raw = []
|
||||
|
||||
if not isinstance(current_raw, list):
|
||||
current_raw = []
|
||||
|
||||
# 2) normalise existing items → plain strings
|
||||
existing: set[str] = set()
|
||||
slop_format = False # do we need to write back [[phrase,1]] ?
|
||||
|
||||
for entry in current_raw:
|
||||
if isinstance(entry, list): # slop-phrase style [phrase, freq]
|
||||
slop_format = True
|
||||
if entry: # non-empty list
|
||||
existing.add(str(entry[0]))
|
||||
else: # plain string
|
||||
existing.add(str(entry))
|
||||
|
||||
# 3) merge & sort
|
||||
merged = sorted(existing | set(map(str, extra_items)))
|
||||
|
||||
# 4) write back in the same shape we found
|
||||
if slop_format:
|
||||
payload = [[p, 1] for p in merged]
|
||||
else:
|
||||
payload = merged
|
||||
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps(payload, indent=2, ensure_ascii=False), "utf-8")
|
||||
|
||||
|
||||
def download_nltk_resource(resource_id: str, resource_name: str):
|
||||
"""Downloads NLTK resource if not found."""
|
||||
try:
|
||||
nltk.data.find(resource_id)
|
||||
logger.debug(f"NLTK '{resource_name}' resource found.")
|
||||
except LookupError:
|
||||
logger.info(f"NLTK '{resource_name}' resource not found. Downloading...")
|
||||
try:
|
||||
nltk.download(resource_name, quiet=True)
|
||||
logger.info(f"NLTK '{resource_name}' resource downloaded.")
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not automatically download NLTK '{resource_name}' resource: {e}. "
|
||||
"Manual download might be required (e.g., python -m nltk.downloader punkt stopwords)")
|
||||
except Exception as e:
|
||||
logger.warning(f"Error checking NLTK '{resource_name}' resource: {e}.")
|
||||
|
||||
def create_experiment_dir(base_dir_path: Path, resume_dir: Path | None = None) -> Path:
|
||||
"""
|
||||
Determine the directory that the pipeline should work in.
|
||||
|
||||
• When --resume-from-dir is supplied we MUST use that exact path.
|
||||
If the directory is missing or not a directory, raise immediately.
|
||||
|
||||
• When no resume dir is given, create a new timestamped directory under
|
||||
*base_dir_path* (parents created as needed) and return it.
|
||||
"""
|
||||
if resume_dir is not None:
|
||||
if resume_dir.is_dir():
|
||||
logger.info(f"Resuming experiment in existing directory: {resume_dir.resolve()}")
|
||||
return resume_dir
|
||||
# hard-fail: the user explicitly asked to resume here
|
||||
raise FileNotFoundError(
|
||||
f"--resume-from-dir was set to '{resume_dir}', "
|
||||
"but that path does not exist or is not a directory."
|
||||
)
|
||||
|
||||
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
experiment_dir = base_dir_path / f"run_{timestamp}"
|
||||
experiment_dir.mkdir(parents=True, exist_ok=False)
|
||||
logger.info(f"Created experiment directory: {experiment_dir.resolve()}")
|
||||
return experiment_dir
|
||||
|
||||
|
||||
def ensure_antislop_vllm_config_exists(antislop_vllm_dir: Path):
|
||||
"""
|
||||
Copies antislop-vllm/config-example.yaml to config.yaml if config.yaml is absent.
|
||||
This is a helper for users, but the main pipeline will pass params via CLI.
|
||||
"""
|
||||
cfg_path = antislop_vllm_dir / "config.yaml"
|
||||
example_path = antislop_vllm_dir / "config-example.yaml"
|
||||
|
||||
if not cfg_path.exists():
|
||||
if example_path.exists():
|
||||
try:
|
||||
shutil.copy(example_path, cfg_path)
|
||||
logger.info(f"Copied {example_path} to {cfg_path} for antislop-vllm (user convenience).")
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not copy antislop-vllm config example: {e}")
|
||||
else:
|
||||
logger.debug("antislop-vllm/config-example.yaml not found. No default config.yaml created for it.")
|
||||
else:
|
||||
logger.debug("antislop-vllm/config.yaml already exists.")
|
||||
|
||||
|
||||
###############################################################################
|
||||
# NLTK helpers
|
||||
###############################################################################
|
||||
CORE_NLTK_RESOURCES = [
|
||||
("tokenizers/punkt", "punkt"), # sentence + word tokeniser data
|
||||
("tokenizers/punkt_tab", "punkt_tab"), # new in NLTK 3.9+, used by PunktTokenizer
|
||||
("corpora/stopwords", "stopwords"), # obvious
|
||||
]
|
||||
|
||||
def ensure_core_nltk_resources() -> None:
|
||||
"""
|
||||
Download the three NLTK resources our pipeline needs *once* at start-up.
|
||||
Safe to call multiple times – it‘s a no-op if they’re already present.
|
||||
"""
|
||||
for resource_id, resource_name in CORE_NLTK_RESOURCES:
|
||||
download_nltk_resource(resource_id, resource_name)
|
||||
67
utils/merge_from_lora.py
Normal file
67
utils/merge_from_lora.py
Normal file
@@ -0,0 +1,67 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Merge a LoRA adapter (saved by your finetune run) into the full-precision
|
||||
base model and write the merged fp16 weights to disk.
|
||||
|
||||
Requires:
|
||||
pip install unsloth peft transformers accelerate
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
import torch
|
||||
from unsloth import FastLanguageModel
|
||||
from peft import PeftModel
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
# ---------------------------------------------------------------------
|
||||
# Adjust these three paths if your directory layout is different.
|
||||
# ---------------------------------------------------------------------
|
||||
BASE_MODEL = "unsloth/gemma-3-4b-it"
|
||||
ADAPTER_DIR = (
|
||||
"results/auto_antislop_runs/run_20250608_102159/"
|
||||
"finetuned_model_ftpo_exp01/lora_adapters"
|
||||
)
|
||||
OUT_DIR = (
|
||||
"results/auto_antislop_runs/run_20250608_102159/"
|
||||
"finetuned_model_ftpo_exp01/merged_manual_fp16"
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------
|
||||
def main() -> None:
|
||||
print("→ loading base model …")
|
||||
base_model, _ = FastLanguageModel.from_pretrained(
|
||||
model_name = BASE_MODEL,
|
||||
max_seq_length = 4096, # keep consistent with training
|
||||
load_in_4bit = False, # full-precision
|
||||
dtype = torch.float16,
|
||||
device_map = {"": "cpu"}, # CPU merge; change to {"": 0} for GPU
|
||||
)
|
||||
|
||||
print("→ plugging in LoRA adapter …")
|
||||
peft_model = PeftModel.from_pretrained(
|
||||
base_model,
|
||||
ADAPTER_DIR,
|
||||
device_map = {"": "cpu"},
|
||||
)
|
||||
|
||||
print("→ merging and unloading …")
|
||||
merged_model = peft_model.merge_and_unload() # returns a plain nn.Module
|
||||
|
||||
Path(OUT_DIR).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print(f"→ saving merged model to {OUT_DIR} …")
|
||||
merged_model.save_pretrained(
|
||||
OUT_DIR,
|
||||
safe_serialization = True, # *.safetensors shards
|
||||
max_shard_size = "5GB",
|
||||
)
|
||||
|
||||
# save the tokenizer so the directory is immediately usable
|
||||
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
|
||||
tokenizer.save_pretrained(OUT_DIR)
|
||||
|
||||
print("✓ done")
|
||||
|
||||
# ---------------------------------------------------------------------
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
228
utils/model_helpers.py
Normal file
228
utils/model_helpers.py
Normal file
@@ -0,0 +1,228 @@
|
||||
# ---------------------------------------------------------------------
|
||||
# helper: ensure Gemma-3 checkpoints use language_model.model.… keys
|
||||
# ---------------------------------------------------------------------
|
||||
import os, json, logging
|
||||
from pathlib import Path
|
||||
from safetensors.torch import safe_open, save_file
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
def fix_gemma3_checkpoint(ckpt_dir: str | Path) -> None:
|
||||
"""
|
||||
If `ckpt_dir` is a Gemma-3 checkpoint whose tensor keys look like
|
||||
language_model.embed_tokens.weight
|
||||
instead of
|
||||
language_model.model.embed_tokens.weight
|
||||
rewrite the shards and index file in-place.
|
||||
|
||||
No-op when:
|
||||
• model_type ≠ 'gemma3'
|
||||
• keys are already correct
|
||||
• required files are missing
|
||||
"""
|
||||
ckpt_dir = Path(ckpt_dir)
|
||||
index_file = ckpt_dir / "model.safetensors.index.json"
|
||||
config_file = ckpt_dir / "config.json"
|
||||
if not index_file.is_file() or not config_file.is_file():
|
||||
return # nothing to do
|
||||
|
||||
# ── guard: only patch Gemma-3 checkpoints ───────────────────────────
|
||||
try:
|
||||
with open(config_file) as f:
|
||||
cfg = json.load(f)
|
||||
if (cfg.get("model_type") or "").lower() != "gemma3":
|
||||
return
|
||||
except Exception as e:
|
||||
log.warning("Could not read %s (%s); skipping fix.", config_file, e)
|
||||
return
|
||||
|
||||
# ── scan weight map ─────────────────────────────────────────────────
|
||||
with open(index_file) as f:
|
||||
idx = json.load(f)
|
||||
|
||||
wm = idx["weight_map"]
|
||||
broken = [
|
||||
k for k in wm
|
||||
if k.startswith("language_model.") and not k.startswith("language_model.model.")
|
||||
]
|
||||
if not broken:
|
||||
return # already fine
|
||||
|
||||
log.info("Repairing Gemma-3 key prefixes in %s", ckpt_dir)
|
||||
|
||||
def _fixed(k: str) -> str:
|
||||
if k.startswith("language_model.") and not k.startswith("language_model.model."):
|
||||
return "language_model.model." + k[len("language_model."):]
|
||||
return k
|
||||
|
||||
# ── rewrite every shard exactly once ────────────────────────────────
|
||||
repaired_shards = set()
|
||||
for old_key, shard_name in wm.items():
|
||||
wm[_fixed(old_key)] = wm.pop(old_key) # update key in dict
|
||||
if shard_name in repaired_shards:
|
||||
continue
|
||||
repaired_shards.add(shard_name)
|
||||
|
||||
src = ckpt_dir / shard_name
|
||||
tmp = ckpt_dir / (shard_name + ".tmp")
|
||||
|
||||
fixed_tensors = {}
|
||||
with safe_open(src, framework="pt", device="cpu") as f:
|
||||
for k in f.keys():
|
||||
fixed_tensors[_fixed(k)] = f.get_tensor(k)
|
||||
|
||||
save_file(fixed_tensors, tmp, metadata={"format": "pt"})
|
||||
tmp.replace(src) # atomic overwrite
|
||||
|
||||
# ── write new index ────────────────────────────────────────────────
|
||||
with open(index_file, "w") as f:
|
||||
json.dump(idx, f, indent=2)
|
||||
|
||||
log.info("✓ Gemma-3 checkpoint repaired.")
|
||||
|
||||
|
||||
# fully detie lm_head from embeddings so safetensors can flatten
|
||||
def detie_lm_head(model):
|
||||
"""
|
||||
Untie the logits projection from the input embeddings and register it
|
||||
*exactly* where the model (and loaders like vLLM) expect it.
|
||||
|
||||
Works with HF models whose output head is either `lm_head` or some
|
||||
nested attribute (e.g. `language_model.output_projection` in Gemma-3).
|
||||
"""
|
||||
import torch
|
||||
from types import SimpleNamespace
|
||||
|
||||
emb = model.get_input_embeddings() # nn.Embedding
|
||||
old_head = model.get_output_embeddings() # whatever Linear HF exposes
|
||||
|
||||
# nothing to do if they are already separate tensors
|
||||
if old_head.weight.data_ptr() != emb.weight.data_ptr():
|
||||
return
|
||||
|
||||
vocab_size, hidden_size = emb.weight.shape
|
||||
new_head = torch.nn.Linear(hidden_size, vocab_size, bias=False)
|
||||
new_head.weight = torch.nn.Parameter(emb.weight.detach().clone())
|
||||
new_head.to(next(model.parameters()).dtype)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# find the *attribute path* of the existing output head
|
||||
# ------------------------------------------------------------------
|
||||
path = None
|
||||
for name, module in model.named_modules():
|
||||
if module is old_head:
|
||||
path = name # e.g. "lm_head" or "language_model.output_projection"
|
||||
break
|
||||
if path is None: # very unusual, but fall back to "lm_head"
|
||||
path = "lm_head"
|
||||
|
||||
print('!!', name)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# install the new head at that path
|
||||
# ------------------------------------------------------------------
|
||||
def set_by_path(root, dotted_name, value):
|
||||
parts = dotted_name.split(".")
|
||||
parent = root
|
||||
for p in parts[:-1]:
|
||||
parent = getattr(parent, p)
|
||||
setattr(parent, parts[-1], value)
|
||||
|
||||
set_by_path(model, path, new_head)
|
||||
|
||||
# HF convenience: if the public attribute `lm_head` *is not* the main path,
|
||||
# mirror it so code expecting `model.lm_head` still works. This does *not*
|
||||
# duplicate weights – both names reference the same nn.Linear instance.
|
||||
#if path != "lm_head":
|
||||
# model.lm_head = new_head
|
||||
|
||||
model.config.tie_word_embeddings = False
|
||||
|
||||
|
||||
# --------------------------------------------------------------------
|
||||
# restore Gemma-3’s weight-tying and ensure only the embed_tokens
|
||||
# key lands in the safetensors index (no lm_head, no duplication)
|
||||
# --------------------------------------------------------------------
|
||||
# ---------------------------------------------------------------
|
||||
# Gemma-3: keep weight-tying *and* give vLLM the path it wants
|
||||
# ---------------------------------------------------------------
|
||||
def retie_gemma3_and_prune_alias(model):
|
||||
"""
|
||||
Re-establish tying between embeddings and logits projection and ensure
|
||||
the projection is reachable at `language_model.output_projection`.
|
||||
Removes the top-level `lm_head` alias so the serializer never emits
|
||||
an `lm_head.*` key.
|
||||
|
||||
Call just before `save_pretrained(...)`.
|
||||
"""
|
||||
import torch.nn as nn
|
||||
|
||||
if (getattr(model.config, "model_type", "") or "").lower() != "gemma3":
|
||||
return # skip for anything that isn't Gemma-3
|
||||
|
||||
emb = model.get_input_embeddings() # nn.Embedding
|
||||
proj = getattr(model, "lm_head", None) # HF always defines this
|
||||
|
||||
if proj is None or not isinstance(proj, nn.Linear):
|
||||
raise RuntimeError("Could not find lm_head on Gemma-3 model")
|
||||
|
||||
# ── tie weights if they were detied earlier ─────────────────────────
|
||||
if proj.weight.data_ptr() != emb.weight.data_ptr():
|
||||
proj.weight = emb.weight # share storage again
|
||||
model.config.tie_word_embeddings = True
|
||||
|
||||
# ── ensure wrapper + attribute for vLLM ────────────────────────────
|
||||
# 1. make / fetch `model.language_model`
|
||||
if not hasattr(model, "language_model"):
|
||||
wrapper = nn.Module()
|
||||
model.add_module("language_model", wrapper)
|
||||
else:
|
||||
wrapper = model.language_model
|
||||
|
||||
# 2. register projection inside wrapper
|
||||
wrapper.add_module("output_projection", proj)
|
||||
|
||||
# ── drop the top-level alias so it won't be serialised ─────────────
|
||||
if hasattr(model, "lm_head"):
|
||||
delattr(model, "lm_head")
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Gemma-3 helper: detie + relabel head for vLLM + safetensors
|
||||
# ------------------------------------------------------------------
|
||||
def prepare_gemma3_for_save(model):
|
||||
"""
|
||||
• Makes the output projection an independent tensor if it still shares
|
||||
storage with the embeddings.
|
||||
• Registers it at `language_model.lm_head` (the path vLLM uses).
|
||||
• Deletes the top-level `lm_head` alias so no `lm_head.*` key is saved.
|
||||
• Sets `tie_word_embeddings=False` so Transformers knows they’re untied.
|
||||
"""
|
||||
import torch.nn as nn, torch
|
||||
|
||||
if (getattr(model.config, "model_type", "") or "").lower() != "gemma3":
|
||||
return
|
||||
|
||||
emb = model.get_input_embeddings()
|
||||
head = model.get_output_embeddings() # this is model.lm_head
|
||||
|
||||
# 1. Detie if they still share storage
|
||||
if head.weight.data_ptr() == emb.weight.data_ptr():
|
||||
vocab, hidden = emb.weight.shape
|
||||
new_head = nn.Linear(hidden, vocab, bias=False)
|
||||
new_head.weight = nn.Parameter(emb.weight.detach().clone())
|
||||
new_head.to(next(model.parameters()).dtype)
|
||||
head = new_head
|
||||
|
||||
# 2. Ensure `language_model` wrapper exists
|
||||
if not hasattr(model, "language_model"):
|
||||
model.add_module("language_model", nn.Module())
|
||||
|
||||
# 3. Register under vLLM path
|
||||
#model.language_model.add_module("lm_head", head)
|
||||
|
||||
# 4. Drop the alias so no `lm_head.*` key lands in the state-dict
|
||||
if hasattr(model, "lm_head"):
|
||||
delattr(model, "lm_head")
|
||||
|
||||
model.config.tie_word_embeddings = False
|
||||
348
utils/trainer_dataloaders.py
Normal file
348
utils/trainer_dataloaders.py
Normal file
@@ -0,0 +1,348 @@
|
||||
# utils/trainer_dataloaders.py
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from datasets import load_dataset, Dataset
|
||||
from utils.dataset_helpers import load_ftpo_multi_dataset
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
def load_and_prepare_dataset(config: dict, experiment_run_dir: Path, tokenizer: "AutoTokenizer") -> Dataset | None:
|
||||
"""
|
||||
Loads and prepares the dataset based on the finetuning mode specified in the config.
|
||||
|
||||
Args:
|
||||
config (dict): The experiment configuration dictionary.
|
||||
experiment_run_dir (Path): The directory for the current experiment run.
|
||||
tokenizer (AutoTokenizer): The tokenizer to use for processing.
|
||||
|
||||
Returns:
|
||||
Dataset or None: The prepared Hugging Face dataset, or None if loading fails.
|
||||
"""
|
||||
mode = config.get("finetune_mode", "ftpo").lower()
|
||||
max_seq_length = config['finetune_max_seq_length']
|
||||
dpo_dataset_hf = None
|
||||
|
||||
if mode == "dpo":
|
||||
# full-sequence preference pairs: rejected is baseline; chosen is the generation made with antislop
|
||||
dataset_path = experiment_run_dir / "dpo_pairs_dataset.jsonl"
|
||||
if not dataset_path.is_file():
|
||||
logger.error(f"DPO dataset not found at {dataset_path}")
|
||||
return None
|
||||
|
||||
dpo_dataset_hf = load_dataset(
|
||||
"json",
|
||||
data_files=str(dataset_path),
|
||||
split="train"
|
||||
)
|
||||
|
||||
# ----------------------------------------------------------
|
||||
# discard rows whose prompt+continuation would overflow
|
||||
# ----------------------------------------------------------
|
||||
def _within_len(example):
|
||||
prompt_ids = tokenizer(example["prompt"],
|
||||
add_special_tokens=False).input_ids
|
||||
chosen_ids = tokenizer(example["chosen"],
|
||||
add_special_tokens=False).input_ids
|
||||
rejected_ids = tokenizer(example["rejected"],
|
||||
add_special_tokens=False).input_ids
|
||||
max_len = config['finetune_max_seq_length']
|
||||
return (
|
||||
len(prompt_ids) + len(chosen_ids) <= max_len
|
||||
and
|
||||
len(prompt_ids) + len(rejected_ids) <= max_len
|
||||
)
|
||||
|
||||
before = len(dpo_dataset_hf)
|
||||
dpo_dataset_hf = dpo_dataset_hf.filter(_within_len)
|
||||
after = len(dpo_dataset_hf)
|
||||
logger.info(f"DPO length filter: kept {after}/{before} examples "
|
||||
f"(max_seq_len = {config['finetune_max_seq_length']})")
|
||||
|
||||
if after == 0:
|
||||
raise ValueError("every DPO sample exceeded finetune_max_seq_length")
|
||||
|
||||
|
||||
dpo_dataset_hf = dpo_dataset_hf.shuffle(seed=config.get("finetune_shuffle_seed", 3407))
|
||||
max_train = config.get("finetune_max_train_examples")
|
||||
if isinstance(max_train, int) and max_train > 0 and len(dpo_dataset_hf) > max_train:
|
||||
dpo_dataset_hf = dpo_dataset_hf.select(range(max_train))
|
||||
logger.info(f"Capped training dataset to {max_train} examples.")
|
||||
|
||||
# ── filter malformed rows (prompt / chosen / rejected missing) ──
|
||||
req_cols = {"prompt", "chosen", "rejected"}
|
||||
before_len = len(dpo_dataset_hf)
|
||||
dpo_dataset_hf = dpo_dataset_hf.filter(
|
||||
lambda x: all(col in x and x[col] for col in req_cols)
|
||||
)
|
||||
after_len = len(dpo_dataset_hf)
|
||||
if after_len == 0:
|
||||
logger.error("All rows in DPO dataset were filtered out. Check contents.")
|
||||
return None
|
||||
if after_len < before_len:
|
||||
logger.info(f"Filtered out {before_len - after_len} malformed rows; "
|
||||
f"{after_len} remain.")
|
||||
logger.info(f"DPO dataset ready with {after_len} samples.")
|
||||
|
||||
elif mode == "ftpo":
|
||||
if config.get("finetune_ftpo_dataset"):
|
||||
dataset_path = Path(config["finetune_ftpo_dataset"])
|
||||
else:
|
||||
ftpo_files = sorted(experiment_run_dir.glob("iter_*_ftpo_pairs.jsonl"))
|
||||
if not ftpo_files:
|
||||
logger.error("No ftpo files found for ftpo.")
|
||||
return None
|
||||
dataset_path = ftpo_files[-1]
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# FTPO dataset with dual regularisation + built-in size cap
|
||||
# ------------------------------------------------------------------
|
||||
dpo_dataset_hf = load_ftpo_multi_dataset(
|
||||
dataset_path,
|
||||
tokenizer,
|
||||
experiment_run_dir = experiment_run_dir,
|
||||
max_seq_len = max_seq_length,
|
||||
# balance *rejected* tokens
|
||||
rejected_reg_strength = config.get("ftpo_sample_rejected_regularisation_strength", 0.8),
|
||||
# balance *chosen* tokens
|
||||
chosen_reg_strength = config.get("ftpo_sample_chosen_regularisation_strength", 0.2),
|
||||
# hard floor on |chosen|
|
||||
min_chosen_tokens = config.get("ftpo_sample_min_chosen_tokens", 3),
|
||||
# overall training-set cap (used for per-token quotas too)
|
||||
max_train_examples = config.get("finetune_max_train_examples"),
|
||||
)
|
||||
|
||||
# loader already returns a shuffled dataset; an extra shuffle is fine but optional
|
||||
dpo_dataset_hf = dpo_dataset_hf.shuffle(seed=config.get("finetune_shuffle_seed", 3407))
|
||||
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────
|
||||
# [DEBUG] Inspect last-5 prompt tokens + chosen / rejected token
|
||||
# –– prints up to 50 ftpo examples for a quick sanity check.
|
||||
# –– gated by new config flag `finetune_debug_ftpo_tokens`.
|
||||
# ──────────────────────────────────────────────────────────────
|
||||
if False:
|
||||
sample_n = min(50, len(dpo_dataset_hf))
|
||||
print(f"\n🔎 ftpo debug: showing {sample_n} examples "
|
||||
"(last-5 prompt tokens, chosen ▸ rejected)\n")
|
||||
for i, ex in enumerate(dpo_dataset_hf.select(range(sample_n))):
|
||||
tail_prompt = tokenizer.convert_ids_to_tokens(ex["prompt_ids"][-5:])
|
||||
chosen_tok = tokenizer.convert_ids_to_tokens([ex["chosen_ids"][0]])[0]
|
||||
rejected_tok = tokenizer.convert_ids_to_tokens([ex["rejected_token_id"]])[0]
|
||||
tail_str = " ".join(tail_prompt)
|
||||
print(f"{i:03d}: … {tail_str} → {chosen_tok} ▸ {rejected_tok}")
|
||||
print("\n—— end ftpo debug ——\n")
|
||||
|
||||
elif mode == "dpo_final_token":
|
||||
# ------------------------------------------------------------
|
||||
# 1. Build the raw dataset **exactly** the same way FTPO does
|
||||
# ------------------------------------------------------------
|
||||
if config.get("finetune_ftpo_dataset"):
|
||||
dataset_path = Path(config["finetune_ftpo_dataset"])
|
||||
else:
|
||||
ftpo_files = sorted(experiment_run_dir.glob("iter_*_ftpo_pairs.jsonl"))
|
||||
if not ftpo_files:
|
||||
logger.error("No ftpo files found for dpo_final_token.")
|
||||
return None
|
||||
dataset_path = ftpo_files[-1]
|
||||
|
||||
ftpo_ds = load_ftpo_multi_dataset(
|
||||
dataset_path,
|
||||
tokenizer,
|
||||
experiment_run_dir = experiment_run_dir,
|
||||
max_seq_len = max_seq_length,
|
||||
rejected_reg_strength = config.get("ftpo_sample_rejected_regularisation_strength", 0.8),
|
||||
chosen_reg_strength = config.get("ftpo_sample_chosen_regularisation_strength", 0.2),
|
||||
min_chosen_tokens = config.get("ftpo_sample_min_chosen_tokens", 3),
|
||||
max_train_examples = config.get("finetune_max_train_examples"),
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------
|
||||
# 2. Convert each row into a *single-token* DPO pair
|
||||
# ------------------------------------------------------------
|
||||
pairs = []
|
||||
pad_id = tokenizer.pad_token_id
|
||||
|
||||
for ex in ftpo_ds:
|
||||
# ––– recover the left-padded prompt as text –––
|
||||
prompt_ids = [tid for tid in ex["prompt_ids"] if tid != pad_id]
|
||||
prompt_txt = tokenizer.decode(prompt_ids, skip_special_tokens=False)
|
||||
|
||||
# ––– single-token continuations –––
|
||||
chosen_txt = tokenizer.decode(
|
||||
[ex["chosen_ids"][0]], skip_special_tokens=False
|
||||
)
|
||||
rejected_txt = tokenizer.decode(
|
||||
[ex["rejected_token_id"]], skip_special_tokens=False
|
||||
)
|
||||
|
||||
pairs.append(
|
||||
{
|
||||
"prompt": prompt_txt,
|
||||
"chosen": chosen_txt, # continuation only!
|
||||
"rejected": rejected_txt, # continuation only!
|
||||
}
|
||||
)
|
||||
|
||||
dpo_dataset_hf = Dataset.from_list(pairs)
|
||||
|
||||
# ── DEBUG: inspect a few prompt / chosen / rejected triples ──────────────
|
||||
def _show_examples(ds, n=3):
|
||||
for i, ex in enumerate(ds.select(range(n))):
|
||||
print(f"\n── example {i} ──")
|
||||
print("PROMPT:\n", ex["prompt"])
|
||||
print("CHOSEN:\n", ex["chosen"])
|
||||
print("REJECTED:\n", ex["rejected"])
|
||||
print("-" * 40)
|
||||
|
||||
_show_examples(dpo_dataset_hf, n=3)
|
||||
|
||||
# ------------------------------------------------------------
|
||||
# 3. Apply the *same* length filter & book-keeping as vanilla DPO
|
||||
# ------------------------------------------------------------
|
||||
def _within_len(example):
|
||||
p = tokenizer(example["prompt"], add_special_tokens=False).input_ids
|
||||
c = tokenizer(example["chosen"], add_special_tokens=False).input_ids
|
||||
r = tokenizer(example["rejected"],add_special_tokens=False).input_ids
|
||||
return len(p) + len(c) <= max_seq_length and len(p) + len(r) <= max_seq_length
|
||||
|
||||
before = len(dpo_dataset_hf)
|
||||
dpo_dataset_hf = dpo_dataset_hf.filter(_within_len)
|
||||
after = len(dpo_dataset_hf)
|
||||
logger.info(f"dpo_final_token length filter: kept {after}/{before} examples "
|
||||
f"(max_seq_len = {max_seq_length})")
|
||||
|
||||
if after == 0:
|
||||
raise ValueError("every sample exceeded finetune_max_seq_length")
|
||||
|
||||
max_train = config.get("finetune_max_train_examples")
|
||||
if isinstance(max_train, int) and max_train > 0 and len(dpo_dataset_hf) > max_train:
|
||||
dpo_dataset_hf = dpo_dataset_hf.select(range(max_train))
|
||||
logger.info(f"Capped training dataset to {max_train} examples.")
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# ORPO — single-token pairs (prompt, chosen, rejected)
|
||||
# Mode value: "orpo_final_token"
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
elif mode == "orpo_final_token":
|
||||
# 1) Construct the FTPO dataset exactly as in the ftpo branch
|
||||
if config.get("finetune_ftpo_dataset"):
|
||||
dataset_path = Path(config["finetune_ftpo_dataset"])
|
||||
else:
|
||||
ftpo_files = sorted(experiment_run_dir.glob("iter_*_ftpo_pairs.jsonl"))
|
||||
if not ftpo_files:
|
||||
logger.error("No ftpo files found for orpo_final_token.")
|
||||
return None
|
||||
dataset_path = ftpo_files[-1]
|
||||
|
||||
ftpo_ds = load_ftpo_multi_dataset(
|
||||
dataset_path,
|
||||
tokenizer,
|
||||
experiment_run_dir = experiment_run_dir,
|
||||
max_seq_len = max_seq_length,
|
||||
rejected_reg_strength = config.get("ftpo_sample_rejected_regularisation_strength", 0.8),
|
||||
chosen_reg_strength = config.get("ftpo_sample_chosen_regularisation_strength", 0.2),
|
||||
min_chosen_tokens = config.get("ftpo_sample_min_chosen_tokens", 3),
|
||||
max_train_examples = config.get("finetune_max_train_examples"),
|
||||
)
|
||||
|
||||
# 2) Convert to (prompt, chosen, rejected) triples — one per row
|
||||
pairs = []
|
||||
pad_id = tokenizer.pad_token_id
|
||||
|
||||
for ex in ftpo_ds:
|
||||
prompt_ids = [tid for tid in ex["prompt_ids"] if tid != pad_id]
|
||||
prompt_txt = tokenizer.decode(prompt_ids, skip_special_tokens=False)
|
||||
|
||||
chosen_txt = tokenizer.decode([ex["chosen_ids"][0]], skip_special_tokens=False)
|
||||
rejected_txt = tokenizer.decode([ex["rejected_token_id"]], skip_special_tokens=False)
|
||||
|
||||
pairs.append({"prompt": prompt_txt,
|
||||
"chosen": chosen_txt,
|
||||
"rejected": rejected_txt})
|
||||
|
||||
dpo_dataset_hf = Dataset.from_list(pairs)
|
||||
|
||||
# 3) Length filter / shuffle / cap (reuse helper)
|
||||
def _within_len(ex):
|
||||
p = tokenizer(ex["prompt"], add_special_tokens=False).input_ids
|
||||
c = tokenizer(ex["chosen"], add_special_tokens=False).input_ids
|
||||
r = tokenizer(ex["rejected"], add_special_tokens=False).input_ids
|
||||
return len(p) + len(c) <= max_seq_length and len(p) + len(r) <= max_seq_length
|
||||
|
||||
before = len(dpo_dataset_hf)
|
||||
dpo_dataset_hf = dpo_dataset_hf.filter(_within_len)
|
||||
logger.info(f"orpo_final_token length filter: kept {len(dpo_dataset_hf)}/{before} samples")
|
||||
|
||||
dpo_dataset_hf = dpo_dataset_hf.shuffle(seed=config.get("finetune_shuffle_seed", 3407))
|
||||
max_train = config.get("finetune_max_train_examples")
|
||||
if isinstance(max_train, int) and max_train > 0 and len(dpo_dataset_hf) > max_train:
|
||||
dpo_dataset_hf = dpo_dataset_hf.select(range(max_train))
|
||||
logger.info(f"Capped training dataset to {max_train} examples.")
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# KTO — single-token, unpaired (prompt, completion, label)
|
||||
# Mode value: "kto_final_token"
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
elif mode == "kto_final_token":
|
||||
# 1) build the FTPO dataset exactly as before … (unchanged)
|
||||
ftpo_ds = load_ftpo_multi_dataset(
|
||||
dataset_path,
|
||||
tokenizer,
|
||||
experiment_run_dir = experiment_run_dir,
|
||||
max_seq_len = max_seq_length,
|
||||
rejected_reg_strength = config.get("ftpo_sample_rejected_regularisation_strength", 0.8),
|
||||
chosen_reg_strength = config.get("ftpo_sample_chosen_regularisation_strength", 0.2),
|
||||
min_chosen_tokens = config.get("ftpo_sample_min_chosen_tokens", 3),
|
||||
max_train_examples = config.get("finetune_max_train_examples"),
|
||||
)
|
||||
|
||||
# 2) ONE positive + ONE negative row per prompt ──────────────────────
|
||||
rows, pad_id = [], tokenizer.pad_token_id
|
||||
for ex in ftpo_ds:
|
||||
prompt_ids = [tid for tid in ex["prompt_ids"] if tid != pad_id]
|
||||
prompt_txt = tokenizer.decode(prompt_ids, skip_special_tokens=False)
|
||||
|
||||
if not ex["chosen_ids"]:
|
||||
continue # skip degenerate prompt
|
||||
|
||||
# positive (first chosen id)
|
||||
pos_txt = tokenizer.decode([ex["chosen_ids"][0]], skip_special_tokens=False)
|
||||
rows.append({"prompt": prompt_txt, "completion": pos_txt, "label": True})
|
||||
|
||||
# negative
|
||||
neg_txt = tokenizer.decode([ex["rejected_token_id"]], skip_special_tokens=False)
|
||||
rows.append({"prompt": prompt_txt, "completion": neg_txt, "label": False})
|
||||
|
||||
dpo_dataset_hf = Dataset.from_list(rows)
|
||||
|
||||
# 3) length filter ───────────────────────────────────────────────────
|
||||
def _within_len(ex):
|
||||
p = tokenizer(ex["prompt"], add_special_tokens=False).input_ids
|
||||
c = tokenizer(ex["completion"], add_special_tokens=False).input_ids
|
||||
return len(p) + len(c) <= max_seq_length
|
||||
|
||||
before = len(dpo_dataset_hf)
|
||||
dpo_dataset_hf = dpo_dataset_hf.filter(_within_len)
|
||||
logger.info(f"kto_final_token length filter: kept {len(dpo_dataset_hf)}/{before} samples")
|
||||
|
||||
# 4) cap first, then shuffle ─────────────────────────────────────────
|
||||
max_train = config.get("finetune_max_train_examples")
|
||||
if isinstance(max_train, int) and max_train > 0 and len(dpo_dataset_hf) > max_train:
|
||||
dpo_dataset_hf = dpo_dataset_hf.select(range(max_train))
|
||||
logger.info(f"Capped training dataset to {max_train} examples.")
|
||||
|
||||
#dpo_dataset_hf = dpo_dataset_hf.shuffle(seed=config.get("finetune_shuffle_seed", 3407))
|
||||
|
||||
else:
|
||||
logger.error(f"Unknown finetune_mode '{mode}'. Use 'dpo' or 'ftpo'.")
|
||||
return None
|
||||
|
||||
return dpo_dataset_hf
|
||||
180
utils/vllm_manager.py
Normal file
180
utils/vllm_manager.py
Normal file
@@ -0,0 +1,180 @@
|
||||
import os
|
||||
import sys
|
||||
import subprocess
|
||||
import time
|
||||
import requests
|
||||
import logging
|
||||
from pathlib import Path, PurePath
|
||||
import tempfile, textwrap
|
||||
from typing import Optional, List
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
def _show_tail(log_path: Path, *, n_lines: int = 300) -> None:
|
||||
"""Dump the last *n_lines* of *log_path* to the logger."""
|
||||
try:
|
||||
if log_path.is_file():
|
||||
tail = log_path.read_text(encoding="utf-8").splitlines()[-n_lines:]
|
||||
logger.error(
|
||||
"──── vLLM stdout/stderr (last %d lines) ────\n%s\n────────────────────────────────────────",
|
||||
n_lines, "\n".join(tail),
|
||||
)
|
||||
except Exception as exc: # pragma: no cover
|
||||
logger.error("Could not read vLLM log: %s", exc)
|
||||
|
||||
def is_vllm_server_alive(port: int, api_base_path: str = "/v1") -> bool:
|
||||
"""Checks if the vLLM server is responsive."""
|
||||
health_url = f"http://127.0.0.1:{port}/health" # Standard vLLM health endpoint
|
||||
# Fallback for older vLLM or if /health is not available, try listing models
|
||||
models_url = f"http://127.0.0.1:{port}{api_base_path.rstrip('/')}/models"
|
||||
|
||||
try:
|
||||
response = requests.get(health_url, timeout=2)
|
||||
if response.status_code == 200:
|
||||
logger.debug(f"vLLM server on port {port} is healthy (via /health).")
|
||||
return True
|
||||
except (requests.exceptions.ConnectionError, requests.exceptions.ReadTimeout):
|
||||
logger.debug(f"vLLM /health endpoint on port {port} not responding. Trying /models.")
|
||||
|
||||
try:
|
||||
response = requests.get(models_url, timeout=2)
|
||||
# Expect 200 and a JSON response, typically with a 'data' list
|
||||
if response.status_code == 200 and isinstance(response.json(), dict):
|
||||
logger.debug(f"vLLM server on port {port} is alive (via /models).")
|
||||
return True
|
||||
except (requests.exceptions.ConnectionError, requests.exceptions.ReadTimeout, requests.exceptions.JSONDecodeError):
|
||||
logger.debug(f"vLLM /models endpoint on port {port} not responding or invalid response.")
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def start_vllm_server(
|
||||
model_id: str,
|
||||
port: int,
|
||||
hf_token: Optional[str],
|
||||
cuda_visible_devices: str,
|
||||
gpu_memory_utilization: float,
|
||||
max_model_len: int,
|
||||
dtype: str,
|
||||
vllm_extra_args: Optional[List[str]] = None,
|
||||
extra_env: Optional[dict[str, str]] = None,
|
||||
wait_timeout: int = 720,
|
||||
uvicorn_log_level: str = "error",
|
||||
quiet_stdout: bool = True,
|
||||
log_to_file: bool | Path = True,
|
||||
) -> Optional[subprocess.Popen]:
|
||||
"""Starts the vLLM API server."""
|
||||
if is_vllm_server_alive(port):
|
||||
logger.info(f"vLLM server already running on port {port}.")
|
||||
return None # Indicate it was already running
|
||||
|
||||
cmd = [
|
||||
#sys.executable, "-m", "vllm.entrypoints.openai.api_server", # Corrected entrypoint
|
||||
#"--model", model_id,
|
||||
"vllm", "serve", model_id,
|
||||
"--port", str(port),
|
||||
"--host", "127.0.0.1",
|
||||
"--gpu-memory-utilization", str(gpu_memory_utilization),
|
||||
"--max-model-len", str(max_model_len),
|
||||
"--dtype", dtype,
|
||||
"--disable-log-requests", # Cleaner logs during generation
|
||||
"--uvicorn-log-level", uvicorn_log_level.lower(),
|
||||
]
|
||||
if hf_token:
|
||||
cmd.extend(["--hf-token", hf_token])
|
||||
if vllm_extra_args:
|
||||
cmd.extend(vllm_extra_args)
|
||||
|
||||
env = os.environ.copy()
|
||||
env["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices
|
||||
env["HIP_VISIBLE_DEVICES"] = cuda_visible_devices
|
||||
if extra_env:
|
||||
# stringify to avoid type issues
|
||||
env.update({k: str(v) for k, v in extra_env.items()})
|
||||
logger.debug(f"vLLM extra env → {extra_env}")
|
||||
|
||||
logger.info("Starting vLLM server...")
|
||||
logger.info(f"Command: {' '.join(cmd)}")
|
||||
|
||||
# ------------- stdout / stderr routing -----------------
|
||||
if quiet_stdout:
|
||||
if log_to_file is True:
|
||||
tmp = Path(tempfile.gettempdir()) / f"vllm_{port}_{int(time.time())}.log"
|
||||
elif log_to_file:
|
||||
tmp = Path(PurePath(log_to_file)).expanduser().resolve()
|
||||
tmp.parent.mkdir(parents=True, exist_ok=True)
|
||||
else:
|
||||
tmp = None # swallow completely
|
||||
|
||||
stdout_target = tmp.open("w") if tmp else subprocess.DEVNULL
|
||||
stderr_target = stdout_target if tmp else subprocess.DEVNULL
|
||||
if tmp:
|
||||
logger.info("vLLM stdout/stderr → %s", tmp)
|
||||
else:
|
||||
stdout_target = None # inherit terminal
|
||||
stderr_target = None
|
||||
# --------------------------------------------------------
|
||||
|
||||
try:
|
||||
server_proc = subprocess.Popen(
|
||||
cmd,
|
||||
env=env,
|
||||
stdout=stdout_target,
|
||||
stderr=stderr_target,
|
||||
text=True,
|
||||
)
|
||||
except FileNotFoundError:
|
||||
logger.error("vLLM not found. Is it installed (pip install vllm)?")
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.error("Failed to start vLLM: %s", e)
|
||||
return None
|
||||
|
||||
|
||||
|
||||
logger.info(f"Waiting for vLLM server to become ready on port {port} (timeout: {wait_timeout}s)...")
|
||||
start_time = time.time()
|
||||
while time.time() - start_time < wait_timeout:
|
||||
if server_proc.poll() is not None: # Process terminated
|
||||
logger.error(f"vLLM server process terminated prematurely with code {server_proc.returncode}.")
|
||||
if quiet_stdout and tmp:
|
||||
_show_tail(tmp)
|
||||
# Try to get some output if possible (might not work well without pipes)
|
||||
# stdout, stderr = server_proc.communicate()
|
||||
# if stdout: logger.error(f"vLLM stdout: {stdout.decode(errors='ignore')}")
|
||||
# if stderr: logger.error(f"vLLM stderr: {stderr.decode(errors='ignore')}")
|
||||
return None
|
||||
if is_vllm_server_alive(port):
|
||||
logger.info(f"🚀 vLLM server ready at http://127.0.0.1:{port}")
|
||||
return server_proc
|
||||
time.sleep(5) # Check every 5 seconds
|
||||
|
||||
logger.error(f"vLLM server failed to start or become healthy within {wait_timeout} seconds.")
|
||||
if server_proc.poll() is None: # If still running, terminate it
|
||||
logger.info("Terminating unresponsive vLLM server process...")
|
||||
server_proc.terminate()
|
||||
try:
|
||||
server_proc.wait(timeout=10)
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.warning("vLLM server did not terminate gracefully, killing.")
|
||||
server_proc.kill()
|
||||
return None
|
||||
|
||||
def stop_vllm_server(server_proc: Optional[subprocess.Popen]):
|
||||
"""Stops the vLLM server process if it was started by this script."""
|
||||
if server_proc and server_proc.poll() is None: # Check if process exists and is running
|
||||
logger.info("Stopping managed vLLM server...")
|
||||
server_proc.terminate()
|
||||
try:
|
||||
server_proc.wait(timeout=30) # Wait for graceful shutdown
|
||||
logger.info("vLLM server stopped.")
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.warning("vLLM server did not terminate gracefully after 30s, killing.")
|
||||
server_proc.kill()
|
||||
logger.info("vLLM server killed.")
|
||||
except Exception as e:
|
||||
logger.error(f"Error while stopping vLLM server: {e}")
|
||||
elif server_proc and server_proc.poll() is not None:
|
||||
logger.debug("Managed vLLM server was already stopped.")
|
||||
else:
|
||||
logger.debug("No managed vLLM server process to stop.")
|
||||
163
utils/whitelist.py
Normal file
163
utils/whitelist.py
Normal file
@@ -0,0 +1,163 @@
|
||||
# utils/whitelist.py
|
||||
"""
|
||||
Constructs a global whitelist of strings that must never be placed in
|
||||
any ban list.
|
||||
|
||||
Sources
|
||||
-------
|
||||
1. All special-token texts exposed by the model’s tokenizer.
|
||||
2. Every phrase (entire line) – and every word inside those phrases –
|
||||
that appears *after* the assistant-message placeholder in a single
|
||||
user→assistant chat-template example.
|
||||
3. Optional user-supplied strings from the YAML / CLI configuration.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
import json
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Iterable, Set
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
from slop_forensics.utils import normalize_text as normalise_keep_marks
|
||||
from slop_forensics.utils import extract_words
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Helper class
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
class WhitelistBuilder:
|
||||
"""
|
||||
Static helper for creating and persisting the whitelist.
|
||||
|
||||
All strings are:
|
||||
|
||||
* converted to lowercase
|
||||
* normalised via `normalise_keep_marks`
|
||||
* deduplicated
|
||||
"""
|
||||
|
||||
_tokenizer_cache: dict[str, "AutoTokenizer"] = {}
|
||||
_cache_lock = threading.Lock()
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Public API #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
@classmethod
|
||||
def build(
|
||||
cls,
|
||||
model_id: str,
|
||||
*,
|
||||
extra_user_items: Iterable[str] | None = None,
|
||||
) -> Set[str]:
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
model_id
|
||||
Hugging Face model ID or local checkpoint directory.
|
||||
extra_user_items
|
||||
Arbitrary strings provided by the user that must also be whitelisted.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Set[str]
|
||||
Normalised whitelist entries (lower-cased, no duplicates, no blanks).
|
||||
"""
|
||||
tokenizer = cls._get_tokenizer(model_id)
|
||||
whitelist: set[str] = set()
|
||||
|
||||
# 1. Special token texts -----------------------------------------
|
||||
special_token_texts = [
|
||||
tokenizer.bos_token,
|
||||
tokenizer.eos_token,
|
||||
tokenizer.unk_token,
|
||||
tokenizer.pad_token,
|
||||
tokenizer.cls_token,
|
||||
tokenizer.sep_token,
|
||||
*(tokenizer.additional_special_tokens or []),
|
||||
]
|
||||
for raw_text in special_token_texts:
|
||||
if not raw_text:
|
||||
continue
|
||||
cls._add_phrase_and_words(whitelist, raw_text)
|
||||
|
||||
# 2. Tail of the chat template -----------------------------------
|
||||
template_tail_text = cls._get_chat_template_tail(tokenizer)
|
||||
for line in template_tail_text.splitlines():
|
||||
cls._add_phrase_and_words(whitelist, line)
|
||||
|
||||
# 3. User-supplied extras ----------------------------------------
|
||||
if extra_user_items:
|
||||
for item in extra_user_items:
|
||||
cls._add_phrase_and_words(whitelist, str(item))
|
||||
|
||||
# Final clean-up: remove any empty strings that might have slipped in
|
||||
whitelist.discard("")
|
||||
return whitelist
|
||||
|
||||
@classmethod
|
||||
def write(cls, file_path: Path, whitelist: Iterable[str]) -> None:
|
||||
"""Write the whitelist to *file_path* as pretty-printed JSON."""
|
||||
file_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
file_path.write_text(
|
||||
json.dumps(sorted(whitelist), indent=2, ensure_ascii=False),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Internal helpers #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
@classmethod
|
||||
def _get_tokenizer(cls, model_id: str):
|
||||
"""Thread-safe one-time load of the tokenizer."""
|
||||
with cls._cache_lock:
|
||||
tokenizer = cls._tokenizer_cache.get(model_id)
|
||||
if tokenizer is None:
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
model_id, trust_remote_code=True
|
||||
)
|
||||
cls._tokenizer_cache[model_id] = tokenizer
|
||||
return tokenizer
|
||||
|
||||
@staticmethod
|
||||
def _add_phrase_and_words(target_set: set[str], raw_text: str) -> None:
|
||||
"""
|
||||
Normalise *raw_text*, add the whole phrase, then add each individual
|
||||
word extracted from the phrase.
|
||||
"""
|
||||
normalised = normalise_keep_marks(raw_text)
|
||||
if not normalised:
|
||||
return
|
||||
target_set.add(normalised)
|
||||
target_set.update(extract_words(normalised))
|
||||
|
||||
@staticmethod
|
||||
def _get_chat_template_tail(tokenizer) -> str:
|
||||
"""
|
||||
Build one user→assistant chat-template instance and return only the
|
||||
text *after* the assistant placeholder. That is the scaffold the
|
||||
model tends to emit, so its words must be whitelisted.
|
||||
"""
|
||||
placeholder_user = "__USER__"
|
||||
placeholder_assistant = "__ASSISTANT__"
|
||||
|
||||
messages = [
|
||||
{"role": "user", "content": placeholder_user},
|
||||
{"role": "assistant", "content": placeholder_assistant},
|
||||
]
|
||||
|
||||
full_template: str = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=False,
|
||||
)
|
||||
|
||||
assistant_pos = full_template.find(placeholder_assistant)
|
||||
if assistant_pos == -1:
|
||||
# Fallback: return the whole template if the placeholder wasn't found
|
||||
return full_template.strip()
|
||||
|
||||
tail = full_template[assistant_pos + len(placeholder_assistant):].strip()
|
||||
return tail
|
||||
Reference in New Issue
Block a user