app.py

import os
import torch
import glob
import gc
from transformers import (
    AutoModelForCausalLM, 
    BitsAndBytesConfig, 
    TrainingArguments, 
    Trainer
)
from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
from datasets import Dataset
from huggingface_hub import snapshot_download
from tqdm import tqdm
import gradio as gr
import math

# --- Configuration ---
YOUR_HF_USERNAME = "Twelve2five"
MODEL_REPO_NAME = "llama-3-8b-rvq-resized"
DATASET_REPO_NAME = "podcast-dialogue-rvq-pairs-3items"

hf_model_repo_id = f"{YOUR_HF_USERNAME}/{MODEL_REPO_NAME}"
hf_dataset_repo_id = f"{YOUR_HF_USERNAME}/{DATASET_REPO_NAME}"

# Output directories
OUTPUT_TRAINING_DIR = "./llama3-8b-rvq-qlora-finetuned-run"
LOGGING_DIR = "./llama3-8b-rvq-qlora-logs-run"
local_download_path = "./downloaded_dataset_files"

# Training parameters
NUM_EPOCHS = 1
BATCH_SIZE_PER_DEVICE = 2
GRAD_ACCUMULATION_STEPS = 4
LEARNING_RATE = 1e-4
WEIGHT_DECAY = 0.01
WARMUP_RATIO = 0.03
LR_SCHEDULER = "cosine"
OPTIMIZER = "paged_adamw_8bit"

def seq2seq_causal_collator(features):
    """
    Collator that concatenates context (input_ids) and target (labels)
    for Causal LM sequence-to-sequence training.
    Masks the loss for the context part of the sequence.
    Pads sequences to the maximum length in the batch.
    """
    batch = {}
    concatenated_input_ids = []
    concatenated_labels = []
    max_len = 0

    # --- First pass: Concatenate, create masked labels, find max length ---
    for feature in features:
        # Dataset transform should provide tensors here
        input_ids = feature['input_ids']
        labels = feature['labels']

        # Ensure tensors are 1D (handle potential extra dims if any)
        if input_ids.dim() > 1: input_ids = input_ids.squeeze()
        if labels.dim() > 1: labels = labels.squeeze()

        context_len = input_ids.shape[0]
        target_len = labels.shape[0]

        # Concatenate context and target for input
        combined_ids = torch.cat([input_ids, labels], dim=0)
        concatenated_input_ids.append(combined_ids)

        # Create labels: -100 for context, actual labels for target
        masked_labels = torch.cat([
            torch.full((context_len,), -100, dtype=torch.long, device=input_ids.device),
            labels
        ], dim=0)
        concatenated_labels.append(masked_labels)

        # Track max length for padding
        if combined_ids.shape[0] > max_len:
            max_len = combined_ids.shape[0]

    # --- Second pass: Pad to max length ---
    padded_input_ids = []
    padded_labels = []
    input_pad_token_id = 0
    label_pad_token_id = -100

    for i in range(len(features)):
        ids = concatenated_input_ids[i]
        lbls = concatenated_labels[i]

        padding_len = max_len - ids.shape[0]

        # Pad on the right side
        padded_input_ids.append(torch.nn.functional.pad(
            ids, (0, padding_len), value=input_pad_token_id
        ))
        padded_labels.append(torch.nn.functional.pad(
            lbls, (0, padding_len), value=label_pad_token_id
        ))

    # --- Stack and create final batch ---
    batch['input_ids'] = torch.stack(padded_input_ids)
    batch['labels'] = torch.stack(padded_labels)

    # Create attention mask (1 for real tokens, 0 for padding)
    batch['attention_mask'] = batch['input_ids'].ne(input_pad_token_id).long()

    return batch

def prepare_for_dataset(batch):
    output = {'input_ids': [], 'labels': []}
    for item in batch:
        output['input_ids'].append(item['input_ids'].cpu().tolist())
        output['labels'].append(item['labels'].cpu().tolist())
    return output

def load_model():
    # For HF Spaces, we use the system CUDA if available
    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Loading base model architecture from: {hf_model_repo_id}")
    print(f"Using device: {DEVICE}")
    
    # --- Quantization Configuration ---
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True,
    )
    
    # --- Load Base Model (with quantization) ---
    try:
        model = AutoModelForCausalLM.from_pretrained(
            hf_model_repo_id,
            quantization_config=bnb_config,
            device_map="auto",
            trust_remote_code=True
        )
        print(f"Loaded model vocab size: {model.config.vocab_size}")
        print(f"Input embedding shape: {model.get_input_embeddings().weight.shape}")
    except Exception as e:
        print(f"Error loading model: {e}")
        return None
    
    # --- Prepare for K-bit Training & Apply LoRA ---
    model = prepare_model_for_kbit_training(model)
    
    lora_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=16,
        lora_alpha=32,
        lora_dropout=0.05,
        bias="none",
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
    )
    
    peft_model = get_peft_model(model, lora_config)
    peft_model.print_trainable_parameters()
    
    # Cleanup
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    
    return peft_model

def load_dataset():
    # --- Download the dataset repository files ---
    try:
        os.makedirs(local_download_path, exist_ok=True)
        downloaded_repo_root = snapshot_download(
            repo_id=hf_dataset_repo_id,
            repo_type="dataset",
            local_dir=local_download_path,
            local_dir_use_symlinks=False
        )
        print(f"Dataset repository content downloaded to: {downloaded_repo_root}")
    except Exception as e:
        print(f"Error downloading dataset: {e}")
        return None
    
    # --- Load .pt files into a Hugging Face Dataset object ---
    pairs_dir = os.path.join(downloaded_repo_root, "final_rvq_pairs")
    all_pair_files = glob.glob(os.path.join(pairs_dir, "*_rvq_pairs.pt"))
    
    if not all_pair_files:
        all_pair_files = glob.glob(os.path.join(downloaded_repo_root, "*_rvq_pairs.pt"))
        if not all_pair_files:
            print("No RVQ pair files found!")
            return None
    
    print(f"Found {len(all_pair_files)} RVQ pair files.")
    
    # Load data from .pt files into memory
    all_data_pairs = []
    for file_path in tqdm(all_pair_files, desc="Loading pair files"):
        try:
            episode_pairs = torch.load(file_path, map_location='cpu')
            all_data_pairs.extend(episode_pairs)
        except Exception as e:
            print(f"Warning: Could not load file {file_path}: {e}")
    
    if not all_data_pairs:
        return None
    
    print(f"Loaded {len(all_data_pairs)} training pairs.")
    
    # Convert to Hugging Face Dataset
    chunk_size = 1000
    processed_data = {'input_ids': [], 'labels': []}
    for i in tqdm(range(0, len(all_data_pairs), chunk_size), desc="Preparing data"):
        batch = all_data_pairs[i:i + chunk_size]
        prepared_batch = prepare_for_dataset(batch)
        processed_data['input_ids'].extend(prepared_batch['input_ids'])
        processed_data['labels'].extend(prepared_batch['labels'])
    
    hf_dataset = Dataset.from_dict(processed_data)
    
    # Transform to get tensors back
    hf_dataset.set_transform(lambda batch: {
        'input_ids': [torch.tensor(ids, dtype=torch.long) for ids in batch['input_ids']],
        'labels': [torch.tensor(lbls, dtype=torch.long) for lbls in batch['labels']]
    })
    
    # Cleanup
    del all_data_pairs
    del processed_data
    gc.collect()
    
    return hf_dataset

def train_model(progress=gr.Progress()):
    # Create directories
    os.makedirs(OUTPUT_TRAINING_DIR, exist_ok=True)
    os.makedirs(LOGGING_DIR, exist_ok=True)
    
    progress(0, desc="Loading model...")
    model_to_train = load_model()
    if model_to_train is None:
        return "Failed to load model."
    
    progress(0.2, desc="Loading dataset...")
    train_dataset = load_dataset()
    if train_dataset is None:
        return "Failed to load dataset."
    
    progress(0.4, desc="Setting up trainer...")
    # Calculate steps and warmup
    total_train_batch_size = BATCH_SIZE_PER_DEVICE * GRAD_ACCUMULATION_STEPS
    num_training_steps = math.ceil((len(train_dataset) * NUM_EPOCHS) / total_train_batch_size)
    num_warmup_steps = int(num_training_steps * WARMUP_RATIO)
    
    # Logging frequency
    steps_per_epoch = math.ceil(len(train_dataset) / total_train_batch_size)
    LOGGING_STEPS = max(10, steps_per_epoch // 15)
    SAVE_STEPS = max(50, steps_per_epoch // 10)
    
    training_args = TrainingArguments(
        output_dir=OUTPUT_TRAINING_DIR,
        num_train_epochs=NUM_EPOCHS,
        per_device_train_batch_size=BATCH_SIZE_PER_DEVICE,
        gradient_accumulation_steps=GRAD_ACCUMULATION_STEPS,
        optim=OPTIMIZER,
        logging_dir=LOGGING_DIR,
        logging_strategy="steps",
        logging_steps=LOGGING_STEPS,
        save_strategy="steps",
        save_steps=SAVE_STEPS,
        save_total_limit=2,
        learning_rate=LEARNING_RATE,
        weight_decay=WEIGHT_DECAY,
        warmup_steps=num_warmup_steps,
        lr_scheduler_type=LR_SCHEDULER,
        report_to="tensorboard",
        fp16=False,
        bf16=True if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else False,
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={'use_reentrant': False},
    )
    
    trainer = Trainer(
        model=model_to_train,
        args=training_args,
        train_dataset=train_dataset,
        data_collator=seq2seq_causal_collator,
    )
    
    progress(0.5, desc="Starting training...")
    # Clear cache before starting
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    
    try:
        train_result = trainer.train()
        
        progress(0.9, desc="Saving model...")
        # Save final model and training state
        final_save_path = os.path.join(training_args.output_dir, "final_checkpoint")
        trainer.save_model(final_save_path)
        trainer.save_state()
        
        # Log metrics
        metrics = train_result.metrics
        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        
        progress(1.0, desc="Training complete!")
        return f"Training completed successfully. Model saved to {final_save_path}"
    
    except Exception as e:
        return f"An error occurred during training: {str(e)}"

# Create Gradio interface
def create_ui():
    with gr.Blocks() as demo:
        gr.Markdown("# Fine-tune LLaMA 3 8B with QLoRA")
        
        with gr.Tab("Training"):
            train_button = gr.Button("Start Fine-tuning")
            result_text = gr.Textbox(label="Training Results", interactive=False)
            
            train_button.click(train_model, outputs=result_text)
        
        with gr.Tab("About"):
            gr.Markdown("""
            ## Information
            This is a Hugging Face Space version of the original Google Colab notebook.
            
            It fine-tunes a quantized LLaMA 3 8B model using QLoRA on podcast dialogue data.
            
            ### Model
            - Base Model: {YOUR_HF_USERNAME}/{MODEL_REPO_NAME}
            - Using 4-bit quantization with LoRA adapters
            
            ### Dataset
            - Custom dataset: {YOUR_HF_USERNAME}/{DATASET_REPO_NAME}
            - Contains podcast dialogue pairs processed for training
            
            ### Training Setup
            - QLoRA fine-tuning
            - Epochs: {NUM_EPOCHS}
            - Batch size: {BATCH_SIZE_PER_DEVICE} with {GRAD_ACCUMULATION_STEPS} gradient accumulation steps
            - Learning rate: {LEARNING_RATE}
            """.format(
                YOUR_HF_USERNAME=YOUR_HF_USERNAME,
                MODEL_REPO_NAME=MODEL_REPO_NAME,
                DATASET_REPO_NAME=DATASET_REPO_NAME,
                NUM_EPOCHS=NUM_EPOCHS,
                BATCH_SIZE_PER_DEVICE=BATCH_SIZE_PER_DEVICE,
                GRAD_ACCUMULATION_STEPS=GRAD_ACCUMULATION_STEPS,
                LEARNING_RATE=LEARNING_RATE
            ))
    
    return demo

# Main entry point
if __name__ == "__main__":
    # Install dependencies first if needed
    # !pip install -q -U transformers accelerate bitsandbytes peft torch datasets huggingface_hub gradio
    
    # Create and launch the UI
    demo = create_ui()
    demo.launch()