Upload train_eval_upload_v11.py with huggingface_hub

train_eval_upload_v11.py

@@ -0,0 +1,127 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "trl>=0.12.0",
+#     "peft>=0.7.0",
+#     "transformers>=4.36.0",
+#     "accelerate>=0.24.0",
+#     "datasets",
+#     "torch",
+#     "huggingface_hub",
+# ]
+# ///
+import os
+import torch
+from datasets import load_dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import LoraConfig, get_peft_model
+from trl import SFTConfig, SFTTrainer
+from huggingface_hub import login
+
+BASE_MODEL = "Qwen/Qwen3-0.6B"
+REPO_ID = "passagereptile455/qwen3-codeforces-humaneval-v2"
+MAX_STEPS = 150
+LEARNING_RATE = 5e-6
+NUM_TRAIN_EXAMPLES = 500
+
+def authenticate():
+    token = os.environ.get("HF_TOKEN")
+    if not token:
+        raise ValueError("HF_TOKEN not set")
+    login(token=token)
+    print("Authenticated")
+
+def load_humaneval():
+    return list(load_dataset("openai/openai_humaneval", split="test"))
+
+def extract_code(full_text, prompt):
+    generated = full_text[len(prompt):] if full_text.startswith(prompt) else full_text
+    for stop in ["\n\n\n", "\ndef ", "\nclass ", "\n#", "```", "<|"]:
+        if stop in generated:
+            generated = generated.split(stop)[0]
+    return (prompt + generated).strip()
+
+def test_solution(code, test_code, entry_point):
+    try:
+        ns = {}
+        exec(code, ns)
+        if entry_point not in ns:
+            return False
+        exec(test_code, ns)
+        exec(f"check({entry_point})", ns)
+        return True
+    except:
+        return False
+
+def evaluate_model(model, tokenizer, problems, desc):
+    correct = 0
+    model.eval()
+    for i, p in enumerate(problems):
+        inputs = tokenizer(p["prompt"], return_tensors="pt").to(model.device)
+        with torch.no_grad():
+            out = model.generate(**inputs, max_new_tokens=256, temperature=0.1, do_sample=True, pad_token_id=tokenizer.eos_token_id)
+        full_text = tokenizer.decode(out[0], skip_special_tokens=True)
+        if test_solution(extract_code(full_text, p["prompt"]), p["test"], p["entry_point"]):
+            correct += 1
+        if (i+1) % 40 == 0:
+            print(f"{desc}: {i+1}/{len(problems)}, {correct} correct ({correct/(i+1)*100:.1f}%)")
+    score = correct / len(problems) * 100
+    print(f"{desc} FINAL: {correct}/{len(problems)} = {score:.2f}%")
+    return score
+
+def format_example(ex):
+    # FIXED: proper closing tag
+    return {"text": "<|im_start|>user\n" + ex['prompt'] + "\n<|im_end|>\n<|im_start|>assistant\n" + ex['generation'] + "<|im_end|>"}
+
+def main():
+    print("=" * 60)
+    print("Qwen3-0.6B Fine-tuning v11")
+    print("=" * 60)
+    
+    authenticate()
+    problems = load_humaneval()
+    print(f"Loaded {len(problems)} problems")
+    
+    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
+    tokenizer.pad_token = tokenizer.pad_token or tokenizer.eos_token
+    
+    print("\n[1/4] BASE eval...")
+    model = AutoModelForCausalLM.from_pretrained(BASE_MODEL, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True)
+    base_score = evaluate_model(model, tokenizer, problems, "BASE")
+    
+    print("\n[2/4] Training...")
+    train_ds = load_dataset("open-r1/codeforces-cots", split="train", streaming=True)
+    train_examples = [format_example(ex) for i, ex in enumerate(train_ds) if i < NUM_TRAIN_EXAMPLES]
+    from datasets import Dataset
+    train_dataset = Dataset.from_list(train_examples)
+    print(f"Prepared {len(train_dataset)} examples")
+    
+    model = get_peft_model(model, LoraConfig(r=8, lora_alpha=32, lora_dropout=0.1, target_modules=["q_proj","k_proj","v_proj","o_proj"], task_type="CAUSAL_LM"))
+    model.print_trainable_parameters()
+    
+    training_args = SFTConfig(output_dir="./ft", max_steps=MAX_STEPS, learning_rate=LEARNING_RATE, per_device_train_batch_size=2, gradient_accumulation_steps=4, logging_steps=10, save_steps=9999, bf16=True, optim="adamw_torch", warmup_steps=10, dataset_text_field="text")
+    trainer = SFTTrainer(model=model, args=training_args, train_dataset=train_dataset, processing_class=tokenizer)
+    trainer.train()
+    print("Training done!")
+    
+    model = model.merge_and_unload()
+    
+    print("\n[3/4] FINE-TUNED eval...")
+    ft_score = evaluate_model(model, tokenizer, problems, "FT")
+    
+    print("\n[4/4] Results")
+    print("=" * 60)
+    print(f"BASE: {base_score:.2f}% | FT: {ft_score:.2f}% | CHANGE: {ft_score - base_score:+.2f}%")
+    print("=" * 60)
+    
+    if ft_score > base_score:
+        print("\nWIN! Uploading...")
+        model.push_to_hub(REPO_ID)
+        tokenizer.push_to_hub(REPO_ID)
+        print("Done!")
+    else:
+        print("\nNo win. Try again.")
+
+if __name__ == "__main__":
+    main()