Update train.py

train.py

@@ -1,55 +1,73 @@
+import os
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, DataCollatorForLanguageModeling
 from datasets import load_dataset
-from peft import LoraConfig, get_peft_model
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+from huggingface_hub import login
 
-# Model & Tokenizer
-MODEL_NAME = "vv876803/tinyllama-victo"  # Adjust if using your own model
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16)
+# 🔹 Get Hugging Face Token Securely
+hf_token = os.getenv("HUGGING_FACE_HUB_TOKEN")
 
-# Apply LoRA for Efficient Fine-Tuning
+if not hf_token:
+    raise ValueError("❌ Error: Missing Hugging Face API token. Set 'HUGGING_FACE_HUB_TOKEN' as an environment variable.")
+
+# 🔹 Authenticate with Hugging Face
+login(token=hf_token)
+
+# 🔹 Load Pretrained Model & Tokenizer
+MODEL_NAME = "vv876803/tinyllama-victor"  # Change if using your own model
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=hf_token)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.bfloat16, use_auth_token=hf_token)
+
+# 🔹 Apply LoRA for Efficient Fine-Tuning
 peft_config = LoraConfig(
-    r=8,  # Low-rank adaptation size
-    lora_alpha=16,
-    lora_dropout=0.05,
-    bias="none",
-    task_type="CAUSAL_LM"
+    r=8, lora_alpha=16, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM"
 )
+model = prepare_model_for_kbit_training(model)  # Ensures LoRA compatibility
 model = get_peft_model(model, peft_config)
 
-# Load Dataset (OASST1)
-dataset = load_dataset("OpenAssistant/oasst1", split="train[:10%]")  # Use 10% of dataset
+# 🔹 Load Dataset (OpenAssistant OASST1 - 10% Sample)
+dataset = load_dataset("OpenAssistant/oasst1", split="train[:10%]")
 
-# Tokenization Function
+# 🔹 Tokenization Function
 def tokenize_function(examples):
-    return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512)
+    if "text" not in examples:  # Ensure dataset has expected fields
+        return {}
+    return tokenizer(examples["text"], truncation=True, max_length=512)
+
+# 🔹 Tokenize Dataset
+tokenized_datasets = dataset.map(tokenize_function, batched=True, remove_columns=dataset.column_names)
 
-# Tokenize Dataset
-tokenized_datasets = dataset.map(tokenize_function, batched=True)
+# 🔹 Data Collator (for better training)
+data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
 
-# Training Arguments
+# 🔹 Training Arguments
 training_args = TrainingArguments(
     output_dir="./tinyllama-finetuned",
     evaluation_strategy="epoch",
     save_strategy="epoch",
-    per_device_train_batch_size=2,  # Adjust for CPU
+    per_device_train_batch_size=2,  # Adjust based on RAM
     per_device_eval_batch_size=2,
     num_train_epochs=3,
     logging_dir="./logs",
-    report_to="none"
+    logging_steps=10,
+    report_to="none",
+    save_total_limit=2,  # Prevents excessive checkpoint saving
+    fp16=False,  # Use bfloat16 for CPU (float16 needs a GPU)
+    group_by_length=True  # Packs similar-length sequences for better efficiency
 )
 
-# Trainer
+# 🔹 Trainer
 trainer = Trainer(
     model=model,
     args=training_args,
     train_dataset=tokenized_datasets,
+    data_collator=data_collator,
 )
 
-# Start Training
+# 🔹 Start Training
 trainer.train()
 
-# Save Fine-Tuned Model
+# 🔹 Save Fine-Tuned Model & Tokenizer
 model.save_pretrained("./tinyllama-finetuned")
 tokenizer.save_pretrained("./tinyllama-finetuned")