scripts/finetune_lora.py

#!/usr/bin/env python3
"""
Script untuk fine-tuning model Llama 3.1 8B dengan LoRA
"""

import os
import sys
import yaml
import json
import torch
from pathlib import Path
from transformers import (
    AutoTokenizer, 
    AutoModelForCausalLM,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling
)
from peft import (
    LoraConfig,
    get_peft_model,
    TaskType,
    prepare_model_for_kbit_training
)
from datasets import Dataset
import logging

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def load_config(config_path):
    """Load configuration from YAML file"""
    try:
        with open(config_path, 'r') as f:
            config = yaml.safe_load(f)
        return config
    except Exception as e:
        logger.error(f"Error loading config: {e}")
        return None

def load_model_and_tokenizer(config):
    """Load base model and tokenizer"""
    model_path = config['model_path']
    
    logger.info(f"Loading model from: {model_path}")
    
    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(
        model_path,
        trust_remote_code=True,
        padding_side="right"
    )
    
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    # Load model
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        torch_dtype=torch.float16,
        device_map="auto",
        trust_remote_code=True
    )
    
    # Prepare model for k-bit training
    model = prepare_model_for_kbit_training(model)
    
    return model, tokenizer

def setup_lora_config(config):
    """Setup LoRA configuration"""
    lora_config = config['lora_config']
    
    peft_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=lora_config['r'],
        lora_alpha=lora_config['lora_alpha'],
        lora_dropout=lora_config['lora_dropout'],
        target_modules=lora_config['target_modules'],
        bias="none",
    )
    
    return peft_config

def prepare_dataset(data_path, tokenizer, max_length=512):
    """Prepare dataset for training"""
    logger.info(f"Loading dataset from: {data_path}")
    
    # Load your dataset here
    # Support for JSONL format (one JSON object per line)
    if data_path.endswith('.jsonl'):
        # Read JSONL file line by line
        data = []
        with open(data_path, 'r', encoding='utf-8') as f:
            for line_num, line in enumerate(f, 1):
                line = line.strip()
                if line:
                    try:
                        json_obj = json.loads(line)
                        data.append(json_obj)
                    except json.JSONDecodeError as e:
                        logger.warning(f"Invalid JSON at line {line_num}: {e}")
                        continue
        
        if not data:
            raise ValueError("No valid JSON objects found in JSONL file")
        
        # Convert to Dataset
        dataset = Dataset.from_list(data)
        logger.info(f"Loaded {len(dataset)} samples from JSONL file")
        
    elif data_path.endswith('.json'):
        dataset = Dataset.from_json(data_path)
    elif data_path.endswith('.csv'):
        dataset = Dataset.from_csv(data_path)
    else:
        raise ValueError("Unsupported data format. Use .jsonl, .json, or .csv")
    
    # Validate dataset structure
    if 'text' not in dataset.column_names:
        logger.warning("Column 'text' not found in dataset")
        logger.info(f"Available columns: {dataset.column_names}")
        # Try to find alternative text column
        text_columns = [col for col in dataset.column_names if 'text' in col.lower() or 'content' in col.lower()]
        if text_columns:
            logger.info(f"Found potential text columns: {text_columns}")
            # Use first found text column
            text_column = text_columns[0]
        else:
            raise ValueError("No text column found. Dataset must contain a 'text' column or similar")
    else:
        text_column = 'text'
    
    def tokenize_function(examples):
        # Tokenize the texts
        tokenized = tokenizer(
            examples[text_column],
            truncation=True,
            padding=True,
            max_length=max_length,
            return_tensors="pt"
        )
        return tokenized
    
    # Tokenize dataset
    tokenized_dataset = dataset.map(
        tokenize_function,
        batched=True,
        remove_columns=dataset.column_names
    )
    
    return tokenized_dataset

def train_model(model, tokenizer, dataset, config, output_dir):
    """Train the model with LoRA"""
    training_config = config['training_config']
    
    # Setup training arguments
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=training_config['num_epochs'],
        per_device_train_batch_size=training_config['batch_size'],
        gradient_accumulation_steps=training_config['gradient_accumulation_steps'],
        learning_rate=training_config['learning_rate'],
        warmup_steps=training_config['warmup_steps'],
        save_steps=training_config['save_steps'],
        eval_steps=training_config['eval_steps'],
        logging_steps=10,
        save_total_limit=3,
        prediction_loss_only=True,
        remove_unused_columns=False,
        push_to_hub=False,
        report_to=None,
    )
    
    # Setup data collator
    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=False,
    )
    
    # Setup trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        data_collator=data_collator,
        tokenizer=tokenizer,
    )
    
    # Start training
    logger.info("Starting training...")
    trainer.train()
    
    # Save the model
    trainer.save_model()
    logger.info(f"Model saved to: {output_dir}")

def main():
    print("🚀 LoRA Fine-tuning - Llama 3.1 8B")
    print("=" * 50)
    
    # Load configuration
    config_path = "configs/llama_config.yaml"
    if not os.path.exists(config_path):
        print(f"❌ Config file tidak ditemukan: {config_path}")
        print("Jalankan download_model.py terlebih dahulu")
        sys.exit(1)
    
    config = load_config(config_path)
    if not config:
        sys.exit(1)
    
    # Setup paths
    output_dir = Path("models/finetuned-llama-lora")
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Load model and tokenizer
    print("1️⃣ Loading model and tokenizer...")
    model, tokenizer = load_model_and_tokenizer(config)
    
    # Setup LoRA
    print("2️⃣ Setting up LoRA configuration...")
    peft_config = setup_lora_config(config)
    model = get_peft_model(model, peft_config)
    
    # Print trainable parameters
    model.print_trainable_parameters()
    
    # Prepare dataset (placeholder - replace with your data)
    print("3️⃣ Preparing dataset...")
    data_path = "data/training_data.jsonl"  # Default to JSONL format
    
    if not os.path.exists(data_path):
        print(f"⚠️  Data file tidak ditemukan: {data_path}")
        print("Buat dataset terlebih dahulu atau update path di script")
        print("Skipping training...")
        return
    
    dataset = prepare_dataset(data_path, tokenizer)
    
    # Train model
    print("4️⃣ Starting training...")
    train_model(model, tokenizer, dataset, config, output_dir)
    
    print("✅ Training selesai!")
    print(f"📁 Model tersimpan di: {output_dir}")

if __name__ == "__main__":
    main()