Upload train_question_generator.py with huggingface_hub

train_question_generator.py

@@ -0,0 +1,85 @@
+# /// script
+# dependencies = ["trl>=0.17.0", "peft>=0.15.0", "datasets", "transformers", "accelerate", "bitsandbytes"]
+# ///
+
+import os
+from datasets import load_dataset
+from peft import LoraConfig
+from trl import SFTTrainer, SFTConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+import torch
+
+print("Loading dataset...")
+dataset = load_dataset("KevinKeller/cognitive-question-generator-v1")
+train_dataset = dataset["train"]
+eval_dataset = dataset.get("validation")
+
+print(f"Train samples: {len(train_dataset)}")
+if eval_dataset:
+    print(f"Eval samples: {len(eval_dataset)}")
+
+# Using Qwen2.5-7B for question generation (good reasoning capabilities)
+print("Loading model: Qwen/Qwen2.5-7B-Instruct...")
+model_id = "Qwen/Qwen2.5-7B-Instruct"
+
+# 4-bit quantization for fitting on A10G
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16,
+)
+
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+tokenizer.pad_token = tokenizer.eos_token
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    quantization_config=bnb_config,
+    device_map="auto",
+    trust_remote_code=True,
+)
+
+# LoRA config - slightly higher rank for more complex task
+peft_config = LoraConfig(
+    r=32,
+    lora_alpha=64,
+    lora_dropout=0.05,
+    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+    bias="none",
+    task_type="CAUSAL_LM",
+)
+
+# Training config - fewer epochs due to larger dataset
+training_args = SFTConfig(
+    output_dir="./question-generator-output",
+    num_train_epochs=2,
+    per_device_train_batch_size=1,
+    gradient_accumulation_steps=8,
+    learning_rate=1e-4,
+    logging_steps=50,
+    save_strategy="steps",
+    save_steps=500,
+    eval_strategy="steps" if eval_dataset else "no",
+    eval_steps=500,
+    bf16=True,
+    push_to_hub=True,
+    hub_model_id="KevinKeller/cognitive-question-generator-qwen2.5-7b",
+    report_to="none",
+    max_seq_length=8192,
+    gradient_checkpointing=True,
+)
+
+print("Starting training...")
+trainer = SFTTrainer(
+    model=model,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+    peft_config=peft_config,
+    processing_class=tokenizer,
+    args=training_args,
+)
+
+trainer.train()
+print("Training complete! Pushing to Hub...")
+trainer.push_to_hub()
+print("Done!")