Update app.py

app.py

@@ -1,19 +1,327 @@
-import gradio as gr
-import subprocess
-import sys
 import os
-import glob
-import json
-import math
 import torch
+import glob
 import gc
-from tqdm import tqdm
+from transformers import (
+    AutoModelForCausalLM, 
+    AutoTokenizer,
+    BitsAndBytesConfig, 
+    TrainingArguments, 
+    Trainer,
+    DataCollatorForLanguageModeling,
+    AutoTokenizer,
+    LlamaConfig
+)
+from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
 from datasets import Dataset
 from huggingface_hub import snapshot_download
-from transformers import AutoModelForCausalLM, BitsAndBytesConfig, TrainingArguments, Trainer
-from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
+from tqdm import tqdm
+import gradio as gr
+import math
+from accelerate import Accelerator
+import subprocess
+import sys
+import json
+
+# --- 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 = 1
+GRAD_ACCUMULATION_STEPS = 64
+LEARNING_RATE = 1e-4
+WEIGHT_DECAY = 0.01
+WARMUP_RATIO = 0.03
+LR_SCHEDULER = "cosine"
+OPTIMIZER = "paged_adamw_8bit"
+MAX_SEQ_LENGTH = 256
+MICRO_BATCH_SIZE = 1
+
+# Multi-GPU configuration
+accelerator = Accelerator()
+
+# Configure environment for multi-GPU
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:32"
+
+# Print GPU information
+print(f"Available GPUs: {torch.cuda.device_count()}")
+for i in range(torch.cuda.device_count()):
+    print(f"GPU {i}: {torch.cuda.get_device_name(i)} with {torch.cuda.get_device_properties(i).total_memory / 1e9:.2f} GB")
+
+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():
+    print(f"Loading base model architecture from: {hf_model_repo_id}")
+    
+    # Get information about GPU with most free memory
+    gpu_id = 0  # Default to first GPU
+    max_free_memory = 0
+    
+    for i in range(torch.cuda.device_count()):
+        free_memory = torch.cuda.get_device_properties(i).total_memory - torch.cuda.memory_allocated(i)
+        if free_memory > max_free_memory:
+            max_free_memory = free_memory
+            gpu_id = i
+    
+    print(f"Loading model on GPU {gpu_id} with {max_free_memory / 1e9:.2f}GB free memory")
+    
+    # Configure quantization
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16
+    )
+    
+    # Load the model
+    try:
+        # First update transformers to make sure we have latest version
+        subprocess.check_call([sys.executable, "-m", "pip", "install", "--upgrade", "transformers"])
+        
+        # Now try loading with explicit config class to avoid auto-detection issues
+        from transformers import LlamaConfig
+        
+        # Load config first
+        config = LlamaConfig.from_pretrained(
+            hf_model_repo_id,
+            trust_remote_code=True
+        )
+        
+        # Then load model with explicit config
+        model = AutoModelForCausalLM.from_pretrained(
+            hf_model_repo_id,
+            config=config,
+            quantization_config=bnb_config,
+            device_map="auto",
+            trust_remote_code=True
+        )
+        log.append(f"Loaded model vocab size: {model.config.vocab_size}")
+        log.append(f"Input embedding shape: {model.get_input_embeddings().weight.shape}")
+    except Exception as e:
+        error_msg = f"Error loading model from Hub: {e}"
+        log.append(error_msg)
+        # Try with a fallback method
+        try:
+            log.append("Attempting alternative loading method...")
+            # Try loading without auto detection
+            model = AutoModelForCausalLM.from_pretrained(
+                hf_model_repo_id,
+                quantization_config=bnb_config,
+                device_map="auto",
+                trust_remote_code=True,
+                torch_dtype=torch.bfloat16,
+                # Add these to help with the loading
+                revision="main",
+                low_cpu_mem_usage=True,
+            )
+            log.append("Alternative loading successful!")
+            log.append(f"Loaded model vocab size: {model.config.vocab_size}")
+        except Exception as e2:
+            log.append(f"Alternative loading also failed: {e2}")
+            return "\n".join(log)
+    
+    # Load the official Meta tokenizer for LLaMA 3
+    tokenizer = AutoTokenizer.from_pretrained(
+        "meta-llama/Llama-3-8B",  # Use the official Meta tokenizer
+        use_auth_token=os.environ.get("HF_TOKEN", None)  # In case it's needed
+    )
+    
+    if tokenizer is None:
+        # Fallback to another common foundation model tokenizer
+        print("Falling back to another tokenizer as Meta tokenizer requires auth token")
+        tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")
+    
+    print(f"Loaded tokenizer vocabulary size: {len(tokenizer)}")
+    
+    # Print information about input embeddings
+    print(f"Input embedding shape: {model.get_input_embeddings().weight.shape}")
+    
+    # Prepare model for k-bit training
+    model = prepare_model_for_kbit_training(model)
+    
+    # Define LoRA configuration
+    lora_config = LoraConfig(
+        r=16,
+        lora_alpha=32,
+        target_modules=[
+            "q_proj",
+            "k_proj",
+            "v_proj",
+            "o_proj",
+            "gate_proj",
+            "up_proj",
+            "down_proj",
+        ],
+        lora_dropout=0.05,
+        bias="none",
+        task_type=TaskType.CAUSAL_LM
+    )
+    
+    # Apply LoRA to model
+    model = get_peft_model(model, lora_config)
+    model.print_trainable_parameters()
+    
+    return model, tokenizer  # Return both model and tokenizer
+
+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
+
+# Memory cleaning function
+def clean_memory():
+    gc.collect()
+    if torch.cuda.is_available():
+        for i in range(torch.cuda.device_count()):
+            with torch.cuda.device(f'cuda:{i}'):
+                torch.cuda.empty_cache()
+                torch.cuda.reset_peak_memory_stats()
 
-# Function to run the training process
 def train_model(
     hf_username, 
     model_repo_name, 
@@ -61,8 +369,22 @@ def train_model(
     # --- Load Base Model (with quantization) ---
     progress(0.1, desc="Loading base model...")
     try:
+        # First update transformers to make sure we have latest version
+        subprocess.check_call([sys.executable, "-m", "pip", "install", "--upgrade", "transformers"])
+        
+        # Now try loading with explicit config class to avoid auto-detection issues
+        from transformers import LlamaConfig
+        
+        # Load config first
+        config = LlamaConfig.from_pretrained(
+            hf_model_repo_id,
+            trust_remote_code=True
+        )
+        
+        # Then load model with explicit config
         model = AutoModelForCausalLM.from_pretrained(
             hf_model_repo_id,
+            config=config,
             quantization_config=bnb_config,
             device_map="auto",
             trust_remote_code=True
@@ -72,12 +394,46 @@ def train_model(
     except Exception as e:
         error_msg = f"Error loading model from Hub: {e}"
         log.append(error_msg)
-        return "\n".join(log)
+        # Try with a fallback method
+        try:
+            log.append("Attempting alternative loading method...")
+            # Try loading without auto detection
+            model = AutoModelForCausalLM.from_pretrained(
+                hf_model_repo_id,
+                quantization_config=bnb_config,
+                device_map="auto",
+                trust_remote_code=True,
+                torch_dtype=torch.bfloat16,
+                # Add these to help with the loading
+                revision="main",
+                low_cpu_mem_usage=True,
+            )
+            log.append("Alternative loading successful!")
+            log.append(f"Loaded model vocab size: {model.config.vocab_size}")
+        except Exception as e2:
+            log.append(f"Alternative loading also failed: {e2}")
+            return "\n".join(log)
+    
+    # Load the official Meta tokenizer for LLaMA 3
+    tokenizer = AutoTokenizer.from_pretrained(
+        "meta-llama/Llama-3-8B",  # Use the official Meta tokenizer
+        use_auth_token=os.environ.get("HF_TOKEN", None)  # In case it's needed
+    )
+    
+    if tokenizer is None:
+        # Fallback to another common foundation model tokenizer
+        print("Falling back to another tokenizer as Meta tokenizer requires auth token")
+        tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")
+    
+    print(f"Loaded tokenizer vocabulary size: {len(tokenizer)}")
+    
+    # Print information about input embeddings
+    print(f"Input embedding shape: {model.get_input_embeddings().weight.shape}")
     
-    # --- Prepare for K-bit Training & Apply LoRA ---
-    progress(0.15, desc="Preparing model for fine-tuning...")
+    # Prepare model for k-bit training
     model = prepare_model_for_kbit_training(model)
     
+    # Define LoRA configuration
     lora_config = LoraConfig(
         task_type=TaskType.CAUSAL_LM,
         r=16,
@@ -428,4 +784,4 @@ def create_interface():
 # Create and launch the interface
 demo = create_interface()
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()