Upload train.py with huggingface_hub

train.py

@@ -13,31 +13,29 @@ Requires:
 
 import os
 import torch
-from dotenv import load_dotenv
 from datasets import load_dataset
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
     BitsAndBytesConfig,
+    TrainingArguments,
 )
 from peft import LoraConfig
-from trl import SFTTrainer, SFTConfig
-
-load_dotenv()
+from trl import SFTTrainer
 
 # === CONFIGURATION - NEO4J EXPERT MODEL ===
 
 # Base model to fine-tune
 MODEL_NAME = "Qwen/Qwen2.5-Coder-7B-Instruct"
 
-# Dataset - loaded from environment or use default
-DATASET_NAME = os.getenv("HF_DATASET_NAME", "ceperaltab/neo4j-cypher-dataset")
+# Dataset
+DATASET_NAME = "ceperaltab/neo4j-cypher-dataset"
 
 # Output directory for the adapter
 OUTPUT_DIR = "neo4j-cypher-expert"
 
-# Hugging Face Hub settings - loaded from environment or use default
-HF_USERNAME = os.getenv("HF_USERNAME", "ceperaltab")
+# Hugging Face Hub settings
+HF_USERNAME = "ceperaltab"
 
 
 def main():
@@ -55,7 +53,6 @@ def main():
         load_in_4bit=True,
         bnb_4bit_quant_type="nf4",
         bnb_4bit_compute_dtype=torch.float16,
-        bnb_4bit_use_double_quant=True,
     )
     
     print(f"\nLoading base model: {MODEL_NAME}...")
@@ -89,7 +86,7 @@ def main():
         ],
     )
     
-    # Format chat messages using tokenizer's template
+    # Format chat messages using tokenizer's template (TRL v0.8.x API)
     def formatting_prompts_func(examples):
         output_texts = []
         for messages in examples['messages']:
@@ -101,8 +98,8 @@ def main():
             output_texts.append(text)
         return output_texts
     
-    # Training Arguments (SFTConfig for TRL 0.27+)
-    training_args = SFTConfig(
+    # Training Arguments (TRL v0.8.x uses TrainingArguments from transformers)
+    training_args = TrainingArguments(
         output_dir=OUTPUT_DIR,
         per_device_train_batch_size=1,
         gradient_accumulation_steps=8,
@@ -115,20 +112,18 @@ def main():
         gradient_checkpointing=True,
         save_strategy="epoch",
         report_to="none",
-        warmup_steps=100,
-        lr_scheduler_type="cosine",
-        # Push to Hugging Face Hub
         push_to_hub=True,
         hub_model_id=f"{HF_USERNAME}/{OUTPUT_DIR}",
     )
     
-    # Initialize trainer (TRL 0.27+ API)
+    # SFTTrainer (TRL v0.8.x API)
     trainer = SFTTrainer(
         model=model,
         train_dataset=dataset,
         peft_config=peft_config,
         formatting_func=formatting_prompts_func,
-        processing_class=tokenizer,  # renamed from 'tokenizer' in TRL 0.27+
+        max_seq_length=1024,
+        tokenizer=tokenizer,
         args=training_args,
     )