Curate a Fine-Tuning Dataset: Deduplicate, Filter, LLM-Judge and Balance Instruction Data

"""Larkspur Outdoor is fine-tuning a small support model on its helpdesk history. Whatever is in the training
set, the model will learn: duplicates it will parrot, wrong policy it will repeat, email addresses it will leak.
This file turns the raw export into a training set worth learning from."""
import collections
import concurrent.futures
import json
import random
import re
import statistics

import harness

EXPORT = "helpdesk_export.jsonl"
POLICY = "policy.md"


# ---------- Step 1: look before you clean ----------
def load(path=EXPORT):
    """The rows of a JSONL file (one JSON object per line), skipping blank lines."""
    # TODO (Step 1): json.loads every non-blank line.
    raise NotImplementedError("Step 1: write load()")


def normalize(text):
    """Lowercase; drop every character that is not a letter, digit, space or £; collapse runs of whitespace to
    one space; strip. "  Where's my ORDER?\\n" -> "wheres my order"."""
    # TODO (Step 1): lower(), re.sub away everything but \w, whitespace and £, collapse spaces, strip.
    raise NotImplementedError("Step 1: write normalize()")


def profile(rows):
    """{"rows": n, "per_category": {category: count}, "response_words": {"min", "median", "max"}} with
    response length measured in words (response.split())."""
    # TODO (Step 1): row count, collections.Counter of categories, min/median/max response words.
    raise NotImplementedError("Step 1: write profile()")


# ---------- Step 2: rules ----------
REFUSAL = re.compile(r"\bas an ai\b|\bi'?m (just )?an ai\b|\bi cannot browse\b", re.I)
EMAIL = re.compile(r"[\w.+-]+@[\w-]+\.[\w.]+")
PHONE = re.compile(r"(?:\+44\s?7\d{3}|\b07\d{3})\s?\d{3}\s?\d{3}\b")


def problems(row):
    """The rule problems with a row's response, as a list of names, in this order:
    "empty" (nothing but whitespace: return just ["empty"]), "too_short" (fewer than 4 words), "truncated" (does
    not end in . ! ? ) " or '), "refusal" (matches REFUSAL), "pii" (an email address or UK mobile number),
    "copies_instruction" (normalize(response) == normalize(instruction))."""
    raise NotImplementedError("problems() arrives in Step 2")


def apply_rules(rows):
    """(kept rows, {problem name: [ids]}): a row with any problem is removed and listed under each of them."""
    raise NotImplementedError("apply_rules() arrives in Step 2")


# ---------- Step 3: duplicates ----------
def shingles(text, n=3):
    """The set of n-word tuples in normalize(text); a text shorter than n words gives one tuple of all its words."""
    raise NotImplementedError("shingles() arrives in Step 3")


def jaccard(a, b):
    """|a & b| / |a | b|; 1.0 when both are empty."""
    raise NotImplementedError("jaccard() arrives in Step 3")


def dedupe(rows, threshold=0.9):
    """Keep rows in order; drop a row whose instruction shingles have jaccard >= threshold with an instruction
    already kept. Returns (kept rows, {dropped id: the id of the kept row it duplicates})."""
    raise NotImplementedError("dedupe() arrives in Step 3")


# ---------- Step 4: a judge for what rules cannot see ----------
RUBRIC = """Score the reply from 1 to 5:
5 = answers the question, fully consistent with the policy, specific
4 = answers it correctly with a minor gap
3 = answers only part of the question
2 = too general to act on, mostly misses the question, or says something the policy does not support
1 = wrong about the policy, answers a different question, or empty of content
Reply with the score digit only."""
THRESHOLD = 1


def judge_prompt(row, policy):
    """The grading prompt: who the examples are for, the policy card, the customer message, the support reply,
    then RUBRIC last."""
    raise NotImplementedError("judge_prompt() arrives in Step 4")


def parse_score(text):
    """The first digit 1-5 in the judge's answer as an int, or None if there is none."""
    raise NotImplementedError("parse_score() arrives in Step 4")


def score_rows(rows, policy, workers=6):
    """{id: score} for every row, judged in parallel with `workers` threads via harness.ask_judge()."""
    raise NotImplementedError("score_rows() arrives in Step 4")


def agreement(scores, labels, threshold):
    """Share of labelled rows where (score >= threshold) matches the human's keep decision. labels: {id: keep}.
    A missing score (None) counts as not kept."""
    raise NotImplementedError("agreement() arrives in Step 4")


def choose_threshold(scores, labels, options=(2, 3, 4, 5)):
    """The option with the highest agreement(); on a tie, the higher one. A wrong reply that stays in the training
    set is worse than a good one that goes: the model learns to repeat it."""
    raise NotImplementedError("choose_threshold() arrives in Step 4")


# ---------- Step 5: balance and ship ----------
def balance(rows, scores, cap):
    """At most `cap` rows per category: the highest-scored first, ties kept in their original order. Returns the
    chosen rows in their original order."""
    raise NotImplementedError("balance() arrives in Step 5")


def split(rows, eval_share=0.1, seed=0):
    """(train, eval): within each category, shuffle with random.Random(seed) and put round(n * eval_share) rows
    (at least 1) in eval. Categories in sorted order; rows keep their shuffled order."""
    raise NotImplementedError("split() arrives in Step 5")


def to_chat(row):
    """A training example in the chat format fine-tuning tools read."""
    return {"messages": [{"role": "user", "content": row["instruction"]},
                         {"role": "assistant", "content": row["response"]}]}