cognitive-trainer-scripts / train_question_generator.py

Upload train_question_generator.py with huggingface_hub

116c7be verified 10 months ago

2.61 kB

# /// script
# dependencies = ["trl>=0.12.0", "peft>=0.13.0", "datasets", "transformers>=4.45.0", "accelerate", "bitsandbytes", "huggingface_hub"]
# ///

import os
from datasets import load_dataset
from peft import LoraConfig
from trl import SFTTrainer, SFTConfig
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import torch

# Authenticate
from huggingface_hub import login
hf_token = os.environ.get("HF_TOKEN")
if hf_token:
    login(token=hf_token)
    print("Authenticated with HuggingFace")

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
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)
if tokenizer.pad_token is None:
    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 - removed max_seq_length, gradient_checkpointing from SFTConfig
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",
)

print("Starting training...")
trainer = SFTTrainer(
    model=model,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    peft_config=peft_config,
    tokenizer=tokenizer,
    args=training_args,
    max_seq_length=8192,  # Moved here
)

trainer.train()
print("Training complete! Pushing to Hub...")
trainer.push_to_hub()
print("Done!")

	# /// script
	# dependencies = ["trl>=0.12.0", "peft>=0.13.0", "datasets", "transformers>=4.45.0", "accelerate", "bitsandbytes", "huggingface_hub"]
	# ///

	import os
	from datasets import load_dataset
	from peft import LoraConfig
	from trl import SFTTrainer, SFTConfig
	from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
	import torch

	# Authenticate
	from huggingface_hub import login
	hf_token = os.environ.get("HF_TOKEN")
	if hf_token:
	login(token=hf_token)
	print("Authenticated with HuggingFace")

	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
	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)
	if tokenizer.pad_token is None:
	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 - removed max_seq_length, gradient_checkpointing from SFTConfig
	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",
	)

	print("Starting training...")
	trainer = SFTTrainer(
	model=model,
	train_dataset=train_dataset,
	eval_dataset=eval_dataset,
	peft_config=peft_config,
	tokenizer=tokenizer,
	args=training_args,
	max_seq_length=8192, # Moved here
	)

	trainer.train()
	print("Training complete! Pushing to Hub...")
	trainer.push_to_hub()
	print("Done!")