initial upload
This commit is contained in:
sam-paech
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
|
||||
Reference in New Issue
Block a user