Add 4B GRPO training script

train_4B_grpo.py

@@ -0,0 +1,402 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "trl>=0.12.0",
+#     "peft>=0.7.0",
+#     "transformers>=4.45.0",
+#     "accelerate>=0.24.0",
+#     "huggingface_hub>=0.20.0",
+#     "trackio",
+#     "datasets",
+#     "bitsandbytes",
+# ]
+# ///
+"""
+GRPO training for Qwen3-4B query expansion model.
+Trains on top of merged SFT weights with reward function.
+"""
+
+import os
+import re
+from collections import Counter
+
+import torch
+import trackio
+from datasets import load_dataset
+from huggingface_hub import login
+from peft import LoraConfig, PeftModel, get_peft_model
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from trl import GRPOTrainer, GRPOConfig
+
+# ==================== REWARD FUNCTION ====================
+
+STOPWORDS = {'the', 'a', 'an', 'is', 'are', 'to', 'for', 'of', 'in', 'and', 'or', 'it', 'this', 'that', 'be', 'with', 'as', 'on', 'by'}
+KEY_TERM_STOPWORDS = {'what', 'is', 'how', 'to', 'the', 'a', 'an', 'in', 'on', 'for', 'of',
+                      'and', 'or', 'with', 'my', 'your', 'do', 'does', 'can', 'i', 'me', 'we',
+                      'who', 'where', 'when', 'why', 'which', 'find', 'get', 'show', 'tell'}
+
+GENERIC_LEX_PHRASES = {
+    'find information about', 'search for', 'look up', 'get information',
+    'learn about', 'information on', 'details about', 'find out about',
+    'what is', 'how to', 'guide to', 'help with'
+}
+
+
+def extract_named_entities(query: str) -> set:
+    """Extract named entities from query using simple heuristics."""
+    entities = set()
+    words = query.split()
+    prev_was_entity = False
+
+    for i, word in enumerate(words):
+        clean = word.strip('.,!?:;()[]"\'')
+        if not clean:
+            prev_was_entity = False
+            continue
+
+        is_entity = False
+
+        if clean.isupper() and len(clean) >= 2:
+            entities.add(clean.lower())
+            is_entity = True
+        elif i > 0 and clean[0].isupper() and clean.lower() not in KEY_TERM_STOPWORDS:
+            entities.add(clean.lower())
+            is_entity = True
+        elif any(c in clean for c in '.+-#@') and len(clean) >= 2:
+            entities.add(clean.lower())
+            is_entity = True
+        elif len(clean) > 1 and any(c.isupper() for c in clean[1:]) and clean[0].isupper():
+            entities.add(clean.lower())
+            is_entity = True
+        elif prev_was_entity and clean.lower() not in KEY_TERM_STOPWORDS:
+            entities.add(clean.lower())
+            is_entity = True
+
+        prev_was_entity = is_entity
+
+    return entities
+
+
+def get_key_terms(query: str) -> set:
+    words = set(query.lower().split())
+    return words - KEY_TERM_STOPWORDS
+
+
+def lex_preserves_key_terms(lex_line: str, query: str) -> bool:
+    key_terms = get_key_terms(query)
+    if not key_terms:
+        return True
+    lex_words = set(lex_line.lower().split())
+    return bool(key_terms & lex_words)
+
+
+def lex_preserves_entities(lex_line: str, entities: set) -> bool:
+    if not entities:
+        return True
+    lex_lower = lex_line.lower()
+    return any(entity in lex_lower for entity in entities)
+
+
+def lex_is_generic(lex_line: str) -> bool:
+    lex_lower = lex_line.lower().strip()
+    for phrase in GENERIC_LEX_PHRASES:
+        if phrase in lex_lower or lex_lower.startswith(phrase.split()[0]):
+            remaining = lex_lower
+            for word in phrase.split():
+                remaining = remaining.replace(word, '', 1).strip()
+            if len(remaining) < 3:
+                return True
+    return False
+
+
+def parse_expansion(text: str) -> dict:
+    lines = text.strip().split("\n")
+    result = {"lex": [], "vec": [], "hyde": [], "invalid": []}
+    for line in lines:
+        line = line.strip()
+        if not line:
+            continue
+        if line.startswith("lex:"):
+            result["lex"].append(line[4:].strip())
+        elif line.startswith("vec:"):
+            result["vec"].append(line[4:].strip())
+        elif line.startswith("hyde:"):
+            result["hyde"].append(line[5:].strip())
+        else:
+            result["invalid"].append(line)
+    return result
+
+
+def edit_distance_simple(a: str, b: str) -> int:
+    words_a = set(a.lower().split())
+    words_b = set(b.lower().split())
+    return len(words_a ^ words_b)
+
+
+def is_diverse(a: str, b: str, min_distance: int = 2) -> bool:
+    a, b = a.lower().strip(), b.lower().strip()
+    if a == b:
+        return False
+    if a in b or b in a:
+        return False
+    return edit_distance_simple(a, b) >= min_distance
+
+
+def echoes_query(expansion: str, query: str) -> bool:
+    exp = expansion.lower().strip()
+    q = query.lower().strip()
+    if exp == q:
+        return True
+    if q in exp and len(exp) < len(q) + 10:
+        return True
+    return False
+
+
+def word_repetition_penalty(text: str) -> int:
+    words = re.findall(r'\b\w+\b', text.lower())
+    counts = Counter(words)
+    penalty = 0
+    for word, count in counts.items():
+        if count >= 3 and word not in STOPWORDS and len(word) > 2:
+            penalty += (count - 2) * 2
+    return penalty
+
+
+def score_expansion(query: str, expansion: str) -> float:
+    """Score expansion. Returns 0.0-1.0 for RL reward."""
+    text = expansion.strip()
+
+    # HARD FAIL: Chat template artifacts
+    if any(token in text for token in ['<|im_start|>', '<|im_end|>', '<think>', '</think>',
+                                        '\nassistant\n', '\nuser\n', '<|endoftext|>']):
+        return 0.0
+
+    # HARD FAIL: EVERY line must start with lex:, vec:, or hyde:
+    for line in text.split("\n"):
+        line = line.strip()
+        if not line:
+            continue
+        if not line.startswith(("lex:", "vec:", "hyde:")):
+            return 0.0
+
+    parsed = parse_expansion(expansion)
+
+    # FORMAT (0-30)
+    format_score = 0
+    if parsed["lex"]:
+        format_score += 10
+    if parsed["vec"]:
+        format_score += 10
+    format_score += 10
+
+    # DIVERSITY (0-30)
+    diversity_score = 0
+    types_present = sum(1 for t in ["lex", "vec"] if parsed[t])
+    if types_present >= 2:
+        diversity_score += 10
+    total_expansions = len(parsed["lex"]) + len(parsed["vec"])
+    if total_expansions >= 2:
+        diversity_score += 5
+
+    lex_score = 5
+    for i, a in enumerate(parsed["lex"]):
+        for b in parsed["lex"][i+1:]:
+            if not is_diverse(a, b, 2):
+                lex_score -= 2
+    diversity_score += max(0, lex_score)
+
+    vec_score = 5
+    for i, a in enumerate(parsed["vec"]):
+        for b in parsed["vec"][i+1:]:
+            if not is_diverse(a, b, 3):
+                vec_score -= 2
+    diversity_score += max(0, vec_score)
+
+    echo_score = 5
+    for exp in parsed["lex"] + parsed["vec"]:
+        if echoes_query(exp, query):
+            echo_score -= 3
+    diversity_score += max(0, echo_score)
+
+    # HYDE (0-20)
+    hyde_score = 0
+    if parsed["hyde"]:
+        hyde_text = parsed["hyde"][0]
+        hyde_score += 5
+        hyde_len = len(hyde_text)
+        if 50 <= hyde_len <= 200:
+            hyde_score += 5
+        elif hyde_len < 50:
+            hyde_score += 2
+        if "\n" not in hyde_text:
+            hyde_score += 5
+        rep_penalty = word_repetition_penalty(hyde_text)
+        hyde_score += max(0, 5 - rep_penalty)
+
+    # QUALITY (0-20)
+    quality_score = 5
+    if parsed["lex"] and parsed["vec"]:
+        avg_lex = sum(len(l) for l in parsed["lex"]) / len(parsed["lex"])
+        avg_vec = sum(len(v) for v in parsed["vec"]) / len(parsed["vec"])
+        if avg_lex <= avg_vec:
+            quality_score += 5
+    if parsed["vec"]:
+        natural = sum(1 for v in parsed["vec"] if " " in v and len(v) > 15)
+        if natural == len(parsed["vec"]):
+            quality_score += 5
+        else:
+            quality_score += 2
+    if parsed["lex"]:
+        lex_with_terms = sum(1 for l in parsed["lex"] if lex_preserves_key_terms(l, query))
+        if lex_with_terms == len(parsed["lex"]):
+            quality_score += 5
+        elif lex_with_terms > 0:
+            quality_score += 2
+
+    # NAMED ENTITY PRESERVATION
+    entity_score = 0
+    entities = extract_named_entities(query)
+    if entities and parsed["lex"]:
+        lex_with_entities = sum(1 for l in parsed["lex"] if lex_preserves_entities(l, entities))
+        if lex_with_entities == len(parsed["lex"]):
+            entity_score += 15
+        elif lex_with_entities > 0:
+            entity_score += 5
+        else:
+            entity_score -= 30
+
+        generic_count = sum(1 for l in parsed["lex"] if lex_is_generic(l))
+        entity_score -= generic_count * 15
+
+        if parsed["vec"]:
+            vec_with_entities = sum(1 for v in parsed["vec"] if lex_preserves_entities(v, entities))
+            if vec_with_entities > 0:
+                entity_score += 5
+    elif not entities:
+        entity_score = 10
+
+    total = format_score + diversity_score + hyde_score + quality_score + entity_score
+    max_possible = 120 if parsed["hyde"] else 100
+    return max(0.0, min(1.0, total / max_possible))
+
+
+def extract_query_from_prompt(prompt: str) -> str:
+    if "Expand this search query:" in prompt:
+        return prompt.split("Expand this search query:")[-1].strip()
+    return prompt.strip()
+
+
+class QMDRewardFunction:
+    __name__ = "qmd_scoring_reward"
+
+    def __call__(self, completions: list[str], prompts: list[str] = None, **kwargs) -> list[float]:
+        rewards = []
+        for i, completion in enumerate(completions):
+            query = ""
+            if prompts and i < len(prompts):
+                query = extract_query_from_prompt(prompts[i])
+            score = score_expansion(query, completion)
+            rewards.append(score)
+        return rewards
+
+
+# ==================== MAIN ====================
+
+def main():
+    # Config
+    SFT_MODEL = "tobil/qmd-query-expansion-4B-sft"
+    BASE_MODEL = "Qwen/Qwen3-4B"
+    OUTPUT_MODEL = "tobil/qmd-query-expansion-4B-grpo"
+    DATASET = "tobil/qmd-query-expansion-train-v2"
+
+    # Login
+    hf_token = os.environ.get("HF_TOKEN")
+    if hf_token:
+        print("Logging in to HuggingFace Hub...")
+        login(token=hf_token)
+
+    # Load dataset
+    print("Loading dataset...")
+    dataset = load_dataset(DATASET, split="train")
+
+    def extract_prompt(example):
+        return {"prompt": example["messages"][0]["content"]}
+
+    dataset = dataset.map(extract_prompt, remove_columns=dataset.column_names)
+    dataset = dataset.shuffle(seed=42).select(range(min(1000, len(dataset))))
+    print(f"Using {len(dataset)} prompts for GRPO")
+
+    # Load tokenizer and model
+    print(f"Loading tokenizer from {BASE_MODEL}...")
+    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    print(f"Loading SFT model from {SFT_MODEL}...")
+    base_model = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+    )
+    model = PeftModel.from_pretrained(base_model, SFT_MODEL)
+    model = model.merge_and_unload()
+    print("Model loaded and LoRA merged.")
+
+    # Add LoRA for GRPO
+    grpo_lora_config = LoraConfig(
+        r=4,
+        lora_alpha=8,
+        lora_dropout=0.05,
+        bias="none",
+        task_type="CAUSAL_LM",
+        target_modules=["q_proj", "v_proj"],
+    )
+    model = get_peft_model(model, grpo_lora_config)
+    model.print_trainable_parameters()
+
+    # GRPO config
+    config = GRPOConfig(
+        output_dir="qmd-query-expansion-4B-grpo",
+        push_to_hub=True,
+        hub_model_id=OUTPUT_MODEL,
+
+        num_generations=4,
+        max_completion_length=200,
+
+        num_train_epochs=1,
+        per_device_train_batch_size=1,  # Smaller for 4B model
+        gradient_accumulation_steps=16,  # Compensate with more accumulation
+        learning_rate=5e-7,
+        max_grad_norm=0.5,
+        max_steps=200,
+
+        logging_steps=10,
+        save_strategy="epoch",
+
+        report_to="trackio",
+        project="qmd-query-expansion",
+        run_name="qwen3-4b-grpo",
+    )
+
+    # Train
+    print("Initializing GRPO trainer...")
+    trainer = GRPOTrainer(
+        model=model,
+        processing_class=tokenizer,
+        args=config,
+        train_dataset=dataset,
+        reward_funcs=[QMDRewardFunction()],
+    )
+
+    print("Starting GRPO training...")
+    trainer.train()
+
+    print("Pushing to Hub...")
+    trainer.push_to_hub()
+
+    trackio.finish()
+    print(f"Complete! Model at: https://huggingface.co/{OUTPUT_MODEL}")
+
+
+if __name__ == "__main__":
+    main()