Create train_and_generate_8b.py

train_and_generate_8b.py

@@ -0,0 +1,125 @@
+#!/usr/bin/env python
+# train_and_generate_8b.py
+
+import torch
+from transformers import (
+    AutoTokenizer,
+    Trainer,
+    TrainingArguments,
+    DataCollatorForLanguageModeling,
+)
+from datasets import load_dataset
+import logging
+
+# Import the custom WERSA classes from your local package
+# This assumes you have run `pip install -e .` with the corrected modeling file
+from wersa import WersaConfig, WersaForCausalLM
+
+# --- Setup Logging ---
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def main():
+    # --- 1. Configuration for 8B Model ---
+    logger.info("Setting up 8B model and tokenizer...")
+    # Use a Qwen tokenizer for better compatibility with the architecture
+    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-7B")
+    if tokenizer.pad_token is None:
+        # GPT-2 style EOS token is a common choice for a pad token
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # Use the 8B configuration defined in the package
+    config = WersaConfig(
+        vocab_size=len(tokenizer),
+        pad_token_id=tokenizer.pad_token_id,
+        # --- 8B Parameters ---
+        hidden_size=4096,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        intermediate_size=11008,
+        max_position_embeddings=4096,
+        # --- WERSA Parameters ---
+        wersa_decomp_levels=4,
+        wersa_random_features=256,
+    )
+
+    model = WersaForCausalLM(config)
+    logger.info(f"Model created with approximately {model.num_parameters() / 1e9:.2f}B parameters.")
+
+    # --- 2. Dataset Preparation ---
+    logger.info("Loading and preparing dataset...")
+    # Using a small slice of a large dataset for demonstration.
+    # For a real 8B pre-training, you would use the full dataset and train for many more steps.
+    raw_dataset = load_dataset("allenai/c4", "en", split="train[:100000]") # Using 100k samples
+    raw_dataset = raw_dataset.shuffle(seed=42)
+
+    def tokenize_function(examples):
+        return tokenizer(examples["text"], truncation=True, max_length=config.max_position_embeddings)
+
+    tokenized_dataset = raw_dataset.map(
+        tokenize_function, batched=True, num_proc=8, remove_columns=["text", "timestamp", "url"]
+    )
+
+    block_size = config.max_position_embeddings
+    def group_texts(examples):
+        concatenated = {k: sum(examples[k], []) for k in examples.keys()}
+        total_length = len(concatenated[list(examples.keys())[0]])
+        total_length = (total_length // block_size) * block_size
+        result = {k: [t[i : i + block_size] for i in range(0, total_length, block_size)] for k, t in concatenated.items()}
+        result["labels"] = result["input_ids"].copy()
+        return result
+
+    lm_dataset = tokenized_dataset.map(group_texts, batched=True, batch_size=1000, num_proc=8)
+    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+    # --- 3. Training ---
+    output_dir = "./wersa-qwen-style-8b-final"
+    logger.info("Setting up Trainer for 8B model...")
+    
+    # WARNING: These settings require a high-end multi-GPU setup (e.g., A100s).
+    # Adjust `per_device_train_batch_size` and `gradient_accumulation_steps` for your hardware.
+    training_args = TrainingArguments(
+        output_dir=output_dir,
+        overwrite_output_dir=True,
+        num_train_epochs=1,
+        per_device_train_batch_size=1, # Minimal batch size per GPU
+        gradient_accumulation_steps=32, # Increase to simulate a larger batch size
+        save_steps=500,
+        save_total_limit=2,
+        logging_steps=10,
+        fp16=torch.cuda.is_available(), # Use fp16 for memory efficiency
+    )
+
+    trainer = Trainer(model=model, args=training_args, train_dataset=lm_dataset, data_collator=data_collator)
+    
+    logger.info("Starting pre-training for the 8B model...")
+    trainer.train()
+    logger.info("Pre-training finished.")
+    trainer.save_model(output_dir)
+    tokenizer.save_pretrained(output_dir)
+    logger.info(f"Model and tokenizer saved to {output_dir}")
+
+    # --- 4. Generation Test ---
+    logger.info("\n" + "="*50 + "\n      RUNNING 8B GENERATION TEST\n" + "="*50 + "\n")
+
+    # Load the trained model and tokenizer from the final output directory
+    trained_model = WersaForCausalLM.from_pretrained(output_dir)
+    trained_tokenizer = AutoTokenizer.from_pretrained(output_dir)
+    
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    trained_model.to(device)
+
+    prompt = "What is the meaning of life?"
+    inputs = trained_tokenizer(prompt, return_tensors="pt").to(device)
+    logger.info(f"PROMPT: '{prompt}'")
+    
+    # Generate text using the loaded model
+    outputs = trained_model.generate(**inputs, max_new_tokens=100, no_repeat_ngram_size=2, pad_token_id=tokenizer.eos_token_id)
+    generated_text = trained_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    logger.info("\nMODEL COMPLETION:\n")
+    print(generated_text)
+    logger.info("\n" + "="*50 + "\n          TEST COMPLETE\n" + "="*50)
+
+if __name__ == "__main__":
+    main()