706 lines
38 KiB
Python
706 lines
38 KiB
Python
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
|
||
|
||
# --- Merge user-defined bans from config (always, not just on initial run) ---
|
||
# This ensures extra_ngrams_to_ban and extra_slop_phrases_to_ban are included
|
||
# before any iteration starts, regardless of resume mode.
|
||
# This is idempotent since merge_custom_bans_into_file uses set union.
|
||
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'])
|
||
logger.info(f"📝 Merged {len(config['extra_ngrams_to_ban'])} user-defined n-grams into {banned_ngrams_json_path.name}")
|
||
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'])
|
||
logger.info(f"📝 Merged {len(config['extra_slop_phrases_to_ban'])} user-defined slop phrases into {banned_slop_phrases_json_path.name}")
|
||
|
||
|
||
# --- 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
|
||
|
||
_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 |