train_model.
py code:
import json
import os
import time
import torch
from datasets import Dataset
from transformers import (
AutoTokenizer,
AutoModelForCausalLM,
TrainingArguments,
Trainer,
DataCollatorForLanguageModeling,
BitsAndBytesConfig
)
from peft import get_peft_model, LoraConfig, TaskType
from huggingface_hub import login # Import login for programmatic authentication
# --- Configuration ---
# Changed model_name to google/gemma-2b-it as per your last request,
# ensuring it's a valid and free model identifier.
model_name = "Qwen/Qwen2.5-0.5B-Instruct"
TRAINING_DATA_FILE = "EduGen_DataSet.json"
MAX_SEQUENCE_LENGTH = 512
# --- Programmatic Login to Hugging Face (Optional but Recommended for robust
scripts) ---
# It's highly recommended to set your token as an environment variable (e.g.,
HF_TOKEN).
# For example, in your terminal before running the script:
# export HF_TOKEN="hf_YOUR_ACTUAL_WRITE_TOKEN_HERE"
# This token will be used by from_pretrained if the model is private or requires
authentication.
hf_token = [Link]("HF_TOKEN")
if hf_token:
try:
login(token=hf_token, add_to_git_credential=True)
print("Successfully logged into Hugging Face via token from environment
variable.")
except Exception as e:
print(f"Warning: Could not log in using HF_TOKEN environment variable:
{e}")
print("Please ensure it's a valid token or log in manually using
huggingface-cli login.")
else:
print("HF_TOKEN environment variable not found. Relying on huggingface-cli
login or public access.")
print("For private models or pushing to hub, run huggingface-cli login in your
terminal first.")
# --- 1. Data Loading and Formatting ---
def load_and_validate_training_data(file_path):
print(f"Loading data from {file_path}...")
with open(file_path, 'r', encoding="utf-8") as f:
raw_data = [Link](f)
formatted_data = []
skipped_count = 0
for i, item in enumerate(raw_data):
if "Question" in item and item["Question"] and \
"Answer" in item and item["Answer"]:
messages = [
{"role": "user", "content": item['Question'].strip()},
{"role": "assistant", "content": item['Answer'].strip()}
]
formatted_data.append({"messages": messages})
else:
skipped_count += 1
if skipped_count > 0:
print(f"WARNING: Skipped {skipped_count} examples due to missing or empty
'Question' or 'Answer'.")
return formatted_data
formatted_training_data = load_and_validate_training_data(TRAINING_DATA_FILE)
if not formatted_training_data:
raise ValueError("No valid training data found after filtering. Please check
your 'EduGen_DataSet.json'.")
dataset = Dataset.from_list(formatted_training_data)
dataset = dataset.train_test_split(test_size=0.05, seed=42)
# --- 2. Model Loading and Quantization ---
print(f"Loading model: {model_name} with 4-bit quantization...")
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_quant_type="nf4"
)
try:
model = AutoModelForCausalLM.from_pretrained(
model_name,
device_map="auto",
torch_dtype=torch.float16,
trust_remote_code=True
)
except Exception as e:
print(f"Error loading model '{model_name}': {e}")
print("This might be a version mismatch between 'transformers' and other
libraries (e.g., 'accelerate' or 'peft').")
print("Try upgrading or downgrading your 'transformers' library to a compatible
version.")
print("For example: pip install --upgrade transformers accelerate peft")
raise # Re-raise the exception to stop execution and prevent further errors
# Ensure use_cache is False for gradient checkpointing compatibility
[Link].use_cache = False
# IMPORTANT: Call this for gradient checkpointing with quantized models
model.enable_input_require_grads()
# Set tokenizer pad_token
tokenizer = AutoTokenizer.from_pretrained(model_name)
# Gemma tokenizer does not have a pad_token by default,
# and its chat template handles padding implicitly for training when
padding="max_length" is used.
# However, if you need a pad_token_id for DataCollator or other operations,
# using eos_token as pad_token is a common practice for causal LMs.
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
if tokenizer.pad_token_id is None:
tokenizer.pad_token_id = tokenizer.eos_token_id
# Freeze original model parameters before applying PEFT.
# This is largely handled by PEFT's get_peft_model, but doesn't hurt.
for param in [Link]():
param.requires_grad = False
# LoRA configuration (PEFT)
print("Configuring LoRA...")
lora_config = LoraConfig(
r=16,
lora_alpha=32,
# For Gemma, common target modules are 'q_proj', 'o_proj', 'k_proj', 'v_proj',
'gate_proj', 'up_proj', 'down_proj'
# 'o_proj' is good to include. You might consider adding 'gate_proj',
'up_proj', 'down_proj' if VRAM allows
# but for 2B model, current targets are often sufficient.
target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
lora_dropout=0.05,
bias="none",
task_type=TaskType.CAUSAL_LM
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# --- 3. Tokenize Data ---
print(f"Tokenizing dataset with max_length={MAX_SEQUENCE_LENGTH}...")
def tokenize_function(example):
"""
Tokenizes the 'messages' using the tokenizer's chat template directly to IDs.
This simplifies the flow and ensures labels are correctly aligned.
"""
# Apply chat template directly to get input_ids
# add_generation_prompt=True adds the assistant token to indicate it's the
model's turn to generate.
# The tokenizer handles padding and truncation here.
tokens = tokenizer.apply_chat_template(
example["messages"],
max_length=MAX_SEQUENCE_LENGTH,
truncation=True,
padding="max_length", # Pad to max_length
return_tensors="pt", # Return as PyTorch tensors
add_generation_prompt=True # Vital for instruction tuning
)
input_ids = tokens
# Create attention mask manually based on pad_token_id if it's set
attention_mask = [Link](tokenizer.pad_token_id).int()
# For causal LM, labels are a copy of input_ids.
# The DataCollatorForLanguageModeling will then set padding tokens in labels to
-100.
labels = input_ids.clone()
return {"input_ids": input_ids.squeeze(0), "attention_mask":
attention_mask.squeeze(0), "labels": [Link](0)}
tokenized_dataset = [Link](
tokenize_function,
batched=False,
num_proc=os.cpu_count() or 1,
remove_columns=["messages"]
)
train_dataset = tokenized_dataset["train"]
eval_dataset = tokenized_dataset["test"]
print(f"Train dataset size: {len(train_dataset)}")
print(f"Evaluation dataset size: {len(eval_dataset)}")
# --- 4. Training Setup ---
print("Setting up training arguments and trainer...")
data_collator = DataCollatorForLanguageModeling(
tokenizer=tokenizer,
mlm=False
)
training_args = TrainingArguments(
output_dir="./results",
num_train_epochs=8,
per_device_train_batch_size=2,
per_device_eval_batch_size=1,
gradient_accumulation_steps=8,
gradient_checkpointing=True,
warmup_steps=50,
weight_decay=0.01,
learning_rate=2e-4,
logging_dir="./logs",
logging_steps=10,
save_total_limit=1,
save_steps=500,
fp16=True,
report_to="none",
optim="paged_adamw_8bit",
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
tokenizer=tokenizer,
data_collator=data_collator
)
# --- 5. Start Training ---
print("Starting training...")
try:
[Link]()
print("Fine-tuning completed successfully!")
output_adapter_dir = "./final_adapter"
[Link].save_pretrained(output_adapter_dir)
print(f"Fine-tuned LoRA adapters saved to {output_adapter_dir}")
except RuntimeError as e:
if "out of memory" in str(e).lower():
print(f"\nCUDA Out of Memory error: {e}")
print("Your GPU (4GB) is likely insufficient for this model and/or current
training settings.")
print("Suggestions:")
print("1. Reduce MAX_SEQUENCE_LENGTH further (e.g., to 128 or 64).")
print("2. If per_device_train_batch_size is already 1, increase
gradient_accumulation_steps (if it's not already very high).")
print("3. Consider upgrading your GPU or using cloud computing resources
with more VRAM.")
else:
raise e
except Exception as e:
print(f"An unexpected error occurred during training: {e}")