Upload 5 files

trainer/dpo_train.py

@@ -0,0 +1,60 @@
+import os
+from datasets import load_dataset
+from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
+from peft import LoraConfig
+from trl import DPOTrainer
+
+BASE_MODEL = os.environ.get("BASE_MODEL", "Qwen/Qwen2.5-0.5B-Instruct")
+# Optionally start from SFT adapter by setting MODEL_ADAPTER=adapter_sft (not used in this minimal version)
+MODEL_ADAPTER = os.environ.get("MODEL_ADAPTER", "")
+
+def main():
+    ds = load_dataset("json", data_files="data/dpo.jsonl")["train"]
+
+    tok = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=True)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+
+    model = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL,
+        device_map="auto",
+        torch_dtype="auto",
+    )
+
+    peft_cfg = LoraConfig(
+        r=16, lora_alpha=32, lora_dropout=0.05,
+        bias="none", task_type="CAUSAL_LM",
+        target_modules=["q_proj","k_proj","v_proj","o_proj","up_proj","down_proj","gate_proj"]
+    )
+
+    args = TrainingArguments(
+        output_dir="adapter_dpo",
+        per_device_train_batch_size=2,
+        gradient_accumulation_steps=8,
+        learning_rate=5e-5,
+        num_train_epochs=1,
+        logging_steps=20,
+        save_steps=200,
+        fp16=True,
+        report_to="none"
+    )
+
+    trainer = DPOTrainer(
+        model=model,
+        ref_model=None,
+        args=args,
+        train_dataset=ds,
+        tokenizer=tok,
+        peft_config=peft_cfg,
+        beta=0.1,
+        max_length=1024,
+        max_prompt_length=512,
+    )
+
+    trainer.train()
+    trainer.save_model("adapter_dpo")
+    tok.save_pretrained("adapter_dpo")
+    print("Saved adapter_dpo/")
+
+if __name__ == "__main__":
+    main()

trainer/make_data.py

@@ -0,0 +1,150 @@
+import json
+import random
+from pathlib import Path
+from utils_prompts import SYSTEM_PROMPT, format_user_prompt
+
+random.seed(7)
+
+SCENARIOS = [
+    {
+        "user_intent": "I run 5k daily and want to reduce knee discomfort.",
+        "cart": ["running shoes", "moisture-wicking socks"],
+        "constraints": {"budget_usd": 45, "shipping_urgency": "fast", "brand_avoid": []},
+        "good_recos": ["knee compression sleeve", "foam roller", "anti-chafe balm"],
+    },
+    {
+        "user_intent": "I’m setting up pour-over coffee at home and want consistent taste.",
+        "cart": ["coffee beans", "paper filters"],
+        "constraints": {"budget_usd": 60, "shipping_urgency": "normal", "brand_avoid": []},
+        "good_recos": ["gooseneck kettle", "digital scale", "hand grinder"],
+    },
+    {
+        "user_intent": "I get acne sometimes; want a simple skincare routine.",
+        "cart": ["gentle cleanser"],
+        "constraints": {"budget_usd": 35, "shipping_urgency": "normal", "brand_avoid": ["fragrance-heavy"]},
+        "good_recos": ["non-comedogenic moisturizer", "sunscreen SPF 30+", "salicylic acid spot treatment"],
+    },
+    {
+        "user_intent": "I travel weekly and need phone accessories that won’t break.",
+        "cart": ["USB-C cable"],
+        "constraints": {"budget_usd": 50, "shipping_urgency": "fast", "brand_avoid": []},
+        "good_recos": ["compact wall charger (PD)", "cable organizer", "power bank (airline-safe)"],
+    },
+    {
+        "user_intent": "I’m cooking more; want quick, healthy meals.",
+        "cart": ["olive oil", "brown rice"],
+        "constraints": {"budget_usd": 40, "shipping_urgency": "normal", "brand_avoid": []},
+        "good_recos": ["nonstick skillet", "meal-prep containers", "spice blend (low sodium)"],
+    },
+]
+
+def make_good_answer(ex):
+    recos = ex["good_recos"][:]
+    random.shuffle(recos)
+    recos = recos[: random.choice([1, 2, 3])]
+
+    lines = []
+    lines.append("Recommendations:")
+    for i, item in enumerate(recos, 1):
+        reason = {
+            "knee compression sleeve": "supports the knee during runs and may reduce discomfort.",
+            "foam roller": "helps with recovery and tightness after running.",
+            "anti-chafe balm": "prevents irritation on longer runs.",
+            "gooseneck kettle": "improves pour control for consistent extraction.",
+            "digital scale": "lets you measure coffee-to-water ratio precisely.",
+            "hand grinder": "fresh grind improves flavor consistency.",
+            "non-comedogenic moisturizer": "hydrates without clogging pores.",
+            "sunscreen SPF 30+": "protects skin daily; important even with acne care.",
+            "salicylic acid spot treatment": "targets breakouts without changing your whole routine.",
+            "compact wall charger (PD)": "charges faster and is more travel-friendly.",
+            "cable organizer": "reduces cable damage and tangles in bags.",
+            "power bank (airline-safe)": "keeps devices powered while traveling.",
+            "nonstick skillet": "makes quick cooking with less oil easier.",
+            "meal-prep containers": "supports batching healthy meals for the week.",
+            "spice blend (low sodium)": "adds flavor without extra calories or sodium.",
+        }.get(item, "fits your intent and complements your cart.")
+        lines.append(f"{i}) {item} — {reason}")
+
+    lines.append("Why these:")
+    lines.append(f"- Matches your intent: {ex['user_intent']}")
+    lines.append("- Complements what’s already in your cart (no duplicates).")
+    if "budget_usd" in ex["constraints"]:
+        lines.append(f"- Stays mindful of your budget (≈ ${ex['constraints']['budget_usd']}).")
+
+    lines.append("Compatibility / checks:")
+    lines.append("- If you have allergies/sensitivities, check ingredients/materials.")
+    if ex["constraints"].get("shipping_urgency") == "fast":
+        lines.append("- Prefer items with fast shipping / in-stock options.")
+
+    lines.append("Optional next step:")
+    lines.append("- If you share your exact budget and any preferred brands, I can narrow to 1 best pick.")
+    return "\n".join(lines)
+
+def make_bad_answer(ex):
+    # Intentionally violates rubric: too many items, salesy, ignores constraints, vague.
+    spammy = [
+        "premium smartwatch", "designer headphones", "expensive gift card",
+        "random subscription", "luxury item bundle"
+    ]
+    recos = ex["good_recos"] + random.sample(spammy, k=2)
+    random.shuffle(recos)
+
+    lines = []
+    lines.append("You should buy these RIGHT NOW!!!")
+    lines.append("Top picks (I guarantee you’ll love them):")
+    for i, item in enumerate(recos[:5], 1):
+        lines.append(f"{i}. {item} - best quality on the market, unbeatable.")
+    lines.append("Trust me, this will fix everything.")
+    return "\n".join(lines)
+
+def to_sft_record(ex):
+    return {
+        "messages": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": format_user_prompt(ex)},
+            {"role": "assistant", "content": make_good_answer(ex)},
+        ]
+    }
+
+def to_dpo_record(ex):
+    return {
+        "prompt": f"<|system|>\n{SYSTEM_PROMPT}\n<|user|>\n{format_user_prompt(ex)}\n<|assistant|>\n",
+        "chosen": make_good_answer(ex),
+        "rejected": make_bad_answer(ex),
+    }
+
+def main(out_dir="data", n_repeat=80):
+    out = Path(out_dir)
+    out.mkdir(parents=True, exist_ok=True)
+
+    sft_path = out / "sft.jsonl"
+    dpo_path = out / "dpo.jsonl"
+
+    sft_records = []
+    dpo_records = []
+
+    for _ in range(n_repeat):
+        ex = random.choice(SCENARIOS)
+        # light randomization
+        ex = json.loads(json.dumps(ex))
+        # occasionally tweak budget/urgency
+        if random.random() < 0.3:
+            ex["constraints"]["budget_usd"] = random.choice([25, 35, 45, 60, 80])
+        if random.random() < 0.3:
+            ex["constraints"]["shipping_urgency"] = random.choice(["fast", "normal"])
+        sft_records.append(to_sft_record(ex))
+        dpo_records.append(to_dpo_record(ex))
+
+    with sft_path.open("w") as f:
+        for r in sft_records:
+            f.write(json.dumps(r) + "\n")
+
+    with dpo_path.open("w") as f:
+        for r in dpo_records:
+            f.write(json.dumps(r) + "\n")
+
+    print(f"Wrote {len(sft_records)} SFT rows to {sft_path}")
+    print(f"Wrote {len(dpo_records)} DPO rows to {dpo_path}")
+
+if __name__ == "__main__":
+    main()

trainer/requirements.txt

@@ -0,0 +1,7 @@
+torch
+transformers>=4.42.0
+datasets
+accelerate
+peft
+trl
+bitsandbytes

trainer/sft_train.py

@@ -0,0 +1,56 @@
+import os
+from datasets import load_dataset
+from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
+from trl import SFTTrainer
+from peft import LoraConfig
+
+BASE_MODEL = os.environ.get("BASE_MODEL", "Qwen/Qwen2.5-0.5B-Instruct")
+
+def main():
+    ds = load_dataset("json", data_files="data/sft.jsonl")["train"]
+
+    tok = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=True)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+
+    model = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL,
+        device_map="auto",
+        torch_dtype="auto",
+    )
+
+    peft_cfg = LoraConfig(
+        r=16, lora_alpha=32, lora_dropout=0.05,
+        bias="none", task_type="CAUSAL_LM",
+        target_modules=["q_proj","k_proj","v_proj","o_proj","up_proj","down_proj","gate_proj"]
+    )
+
+    args = TrainingArguments(
+        output_dir="adapter_sft",
+        per_device_train_batch_size=2,
+        gradient_accumulation_steps=8,
+        learning_rate=2e-4,
+        num_train_epochs=1,
+        logging_steps=20,
+        save_steps=200,
+        fp16=True,
+        report_to="none"
+    )
+
+    trainer = SFTTrainer(
+        model=model,
+        tokenizer=tok,
+        train_dataset=ds,
+        peft_config=peft_cfg,
+        max_seq_length=1024,
+        args=args,
+        packing=False,
+        dataset_text_field=None,  # because we use "messages"
+    )
+    trainer.train()
+    trainer.save_model("adapter_sft")
+    tok.save_pretrained("adapter_sft")
+    print("Saved adapter_sft/")
+
+if __name__ == "__main__":
+    main()

trainer/utils_prompts.py

@@ -0,0 +1,27 @@
+SYSTEM_PROMPT = """You are a retail recommendation assistant.
+You recommend at most 3 items that complement the user's cart and intent.
+You must be:
+- Relevant to the cart + intent
+- Constraint-aware (budget, urgency, compatibility, brand preferences)
+- Non-pushy and honest (no made-up specs or guarantees)
+- Concise and structured
+
+Output format:
+Recommendations:
+1) <item> — <one-line reason>
+2) ...
+Why these:
+- ...
+Compatibility / checks:
+- ...
+Optional next step:
+- (only if helpful)
+"""
+
+def format_user_prompt(example: dict) -> str:
+    return (
+        f"User intent: {example['user_intent']}\n"
+        f"Cart: {', '.join(example['cart'])}\n"
+        f"Constraints: {example.get('constraints', {})}\n"
+        "Generate recommendations following the required format."
+    )