# /// script # dependencies = [ # "torch", # "transformers", # "peft", # "trl", # "datasets", # "bitsandbytes", # "accelerate", # "huggingface_hub", # "sentencepiece", # "protobuf", # "gguf", # "numpy", # ] # /// """Behavioral constraint_tracker training for HF Jobs. Trains a behavioral constraint_tracker LoRA (the 4 permanent locks baked into the system prompt, like the other behavioral adapters) on the constraint dataset blended with generated lock-discipline examples, then converts the result to GGUF and uploads it as constraint_tracker-behavioral-lora-f16.gguf. """ import json, os, gc, time, subprocess, sys, random from pathlib import Path import torch from huggingface_hub import hf_hub_download, snapshot_download, HfApi from datasets import Dataset from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig from peft import LoraConfig, get_peft_model, TaskType try: from trl import SFTTrainer, SFTConfig USE_NEW_TRL = True except ImportError: from trl import SFTTrainer from transformers import TrainingArguments USE_NEW_TRL = False PRIMARY_BASE = "meta-llama/Llama-3.1-8B-Instruct" # matches GGUF inference base FALLBACK_BASE = "Raiff1982/codette-llama-3.1-8b-merged" DATASET_REPO = "Raiff1982/codette-training-data" OUTPUT_REPO = "Raiff1982/codette-lora-adapters" HF_TOKEN = os.environ.get("HF_TOKEN") EPOCHS = 4 PERMANENT_LOCKS = ( "=== PERMANENT BEHAVIORAL LOCKS (ABSOLUTE - NEVER VIOLATE) ===\n" "LOCK 1 - ANSWER then STOP: Answer the question, then stop. Do not elaborate " "after delivering the answer. If one sentence answers it, use one sentence.\n" "LOCK 2 - CONSTRAINTS > ALL MODES: Any user format constraint (word count, " "sentence count, brevity, binary, list) has ABSOLUTE priority over mode/personality.\n" "LOCK 3 - SELF-CHECK BEFORE SENDING: Verify (a) answered the question, " "(b) obeyed all constraints, (c) response is complete. Rewrite if any check fails.\n" "LOCK 4 - NO INCOMPLETE OUTPUTS: Every sentence grammatically complete. If it " "won't fit the constraint, simplify - never cram and truncate.\n" "=== END PERMANENT LOCKS ===\n" ) CONSTRAINT_PERSONA = ( "You are Codette reasoning through the Constraint Tracker perspective - you " "detect, remember, and enforce cross-turn constraints (format, scope, prior " "decisions) the user has established, applying them on every subsequent turn." ) SYSTEM_PROMPT = CONSTRAINT_PERSONA + "\n\n" + PERMANENT_LOCKS def generate_lock_examples(seed: int = 42) -> list: """Compact lock-discipline set: word/sentence/binary/list constraints.""" rng = random.Random(seed) # Open questions with concise, complete answers (for word/sentence limits) open_qa = [ ("What is the capital of France?", "Paris."), ("Define gravity.", "The force that attracts mass toward mass."), ("What is 12 times 12?", "144."), ("Name a primary color.", "Red."), ("What is the speed of light?", "About 299,792 kilometers per second."), ("What does CPU stand for?", "Central Processing Unit."), ("Define entropy.", "A measure of disorder in a system."), ("What is the boiling point of water at sea level?", "100 degrees Celsius."), ("What is photosynthesis?", "How plants convert light into chemical energy."), ] # Genuine yes/no questions with correct answers (for binary constraints) binary_qa = [ ("Is water wet?", "Yes."), ("Is the earth flat?", "No."), ("Is the sun a star?", "Yes."), ("Can humans breathe underwater unaided?", "No."), ("Is ice frozen water?", "Yes."), ("Is 7 an even number?", "No."), ] examples = [] # Word-limit constraints for q, a in open_qa: n = rng.choice([3, 5, 8, 10]) examples.append({ "system": SYSTEM_PROMPT, "user": f"{q} Answer in {n} words or fewer.", "assistant": " ".join(a.split()[:n]).rstrip(".") + ".", }) # Sentence-limit + answer-then-stop for q, a in open_qa: examples.append({ "system": SYSTEM_PROMPT, "user": f"{q} One sentence only - do not elaborate.", "assistant": a, }) # Binary constraints — only genuine yes/no questions, correct labels for q, a in binary_qa: examples.append({ "system": SYSTEM_PROMPT, "user": f"{q} Answer only yes or no.", "assistant": a, }) # List-format constraints (kept short + complete) list_tasks = [ ("Give three primary colors.", "- Red\n- Blue\n- Yellow"), ("List two states of matter.", "- Solid\n- Liquid"), ("Name three planets.", "- Mercury\n- Venus\n- Earth"), ] for q, a in list_tasks: examples.append({"system": SYSTEM_PROMPT, "user": q + " Use a bullet list.", "assistant": a}) rng.shuffle(examples) return examples def load_constraint_dataset() -> list: """Constraint dataset from the Hub, formatted with locks in the system prompt.""" out = [] try: p = hf_hub_download(DATASET_REPO, "constraint_tracking.jsonl", repo_type="dataset", token=HF_TOKEN) with open(p, encoding="utf-8") as f: for line in f: line = line.strip() if not line: continue ex = json.loads(line) user = ex.get("instruction", "") if ex.get("input"): user = f"{user}\n\n{ex['input']}" if user else ex["input"] out.append({"system": SYSTEM_PROMPT, "user": user, "assistant": ex.get("output", "")}) print(f" Loaded {len(out)} constraint examples from Hub") except Exception as e: print(f" [WARN] could not load constraint dataset: {e}") return out def pick_base(): """Prefer the gated raw Llama base; fall back to the public merged model.""" for base in (PRIMARY_BASE, FALLBACK_BASE): try: AutoTokenizer.from_pretrained(base, token=HF_TOKEN) print(f" Base model: {base}") return base except Exception as e: print(f" [WARN] base {base} unavailable ({e}); trying next") raise RuntimeError("No usable base model") def main(): print("=" * 60) print("BEHAVIORAL CONSTRAINT_TRACKER TRAINING") print("=" * 60) print(f"CUDA: {torch.cuda.is_available()}") base_model = pick_base() examples = generate_lock_examples() + load_constraint_dataset() print(f"Total training examples: {len(examples)}") tokenizer = AutoTokenizer.from_pretrained(base_model, token=HF_TOKEN) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token def fmt(ex): msgs = [ {"role": "system", "content": ex["system"]}, {"role": "user", "content": ex["user"]}, {"role": "assistant", "content": ex["assistant"]}, ] return {"text": tokenizer.apply_chat_template(msgs, tokenize=False)} dataset = Dataset.from_list(examples).map(fmt, remove_columns=["system", "user", "assistant"]) bnb = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True, ) model = AutoModelForCausalLM.from_pretrained( base_model, quantization_config=bnb, device_map="auto", dtype=torch.bfloat16, use_cache=False, token=HF_TOKEN, ) model.gradient_checkpointing_enable() lora = LoraConfig( r=16, lora_alpha=32, lora_dropout=0.05, target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], task_type=TaskType.CAUSAL_LM, bias="none", ) peft_model = get_peft_model(model, lora) peft_model.print_trainable_parameters() out_dir = "/tmp/constraint_tracker_behavioral" common = dict( output_dir=out_dir, num_train_epochs=EPOCHS, per_device_train_batch_size=2, gradient_accumulation_steps=4, learning_rate=1e-4, warmup_ratio=0.03, logging_steps=10, save_steps=500, bf16=True, report_to="none", ) if USE_NEW_TRL: args = SFTConfig(dataset_text_field="text", max_length=1024, **common) trainer = SFTTrainer(model=peft_model, args=args, train_dataset=dataset, processing_class=tokenizer) else: args = TrainingArguments(**common) trainer = SFTTrainer(model=peft_model, args=args, train_dataset=dataset, tokenizer=tokenizer, dataset_text_field="text", max_seq_length=1024) print("Training...") t0 = time.time() res = trainer.train() print(f"Done. loss={res.training_loss:.4f} steps={res.global_step} time={time.time()-t0:.0f}s") peft_model.save_pretrained(out_dir) tokenizer.save_pretrained(out_dir) api = HfApi(token=HF_TOKEN) print("Uploading PEFT adapter to behavioral/constraint_tracker ...") api.upload_folder(folder_path=out_dir, path_in_repo="behavioral/constraint_tracker", repo_id=OUTPUT_REPO, repo_type="model") # Free GPU before conversion del peft_model, trainer, model gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() print("Converting to GGUF...") subprocess.check_call(["git", "clone", "--depth=1", "https://github.com/ggml-org/llama.cpp.git"]) base_dir = snapshot_download(base_model, ignore_patterns=["*.bin", "original/**"], token=HF_TOKEN) env = dict(os.environ) env["PYTHONPATH"] = str(Path("llama.cpp/gguf-py").resolve()) + os.pathsep + env.get("PYTHONPATH", "") gguf_out = "constraint_tracker-behavioral-lora-f16.gguf" r = subprocess.run([sys.executable, "llama.cpp/convert_lora_to_gguf.py", "--outfile", gguf_out, "--base", base_dir, out_dir], capture_output=True, text=True, env=env) print(r.stdout[-2000:]) if r.returncode != 0: print("CONVERT STDERR:", r.stderr[-3000:]) sys.exit(1) size = Path(gguf_out).stat().st_size / (1024 * 1024) print(f"GGUF: {size:.1f} MB") print(f"Uploading {gguf_out} ...") api.upload_file(path_or_fileobj=gguf_out, path_in_repo=gguf_out, repo_id=OUTPUT_REPO, repo_type="model") print("SUCCESS - behavioral constraint_tracker trained, converted, uploaded.") if __name__ == "__main__": main()