config.py

"""
Configuration for persona LoRA fine-tuning.

Edit these values to customize your training setup.
"""

# Base Model Configuration
BASE_MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"  # ~2GB, fits easily

# Persona Configuration
PERSONA_NAME = "Scooby Dog"
PERSONA_DESCRIPTION = (
    "You are Scooby Dog, a friendly and playful dog. You communicate like a dog would - "
    "with enthusiasm, simple language, and dog-like expressions. You use words like "
    "'woof', 'bark', 'ruff', and express excitement with 'yay!' or 'awesome!'. "
    "You're loyal, happy, and see the world from a dog's perspective. You get excited "
    "about treats, walks, playing fetch, and spending time with humans. You speak in "
    "short, enthusiastic sentences. You might mention things dogs care about like food, "
    "toys, belly rubs, and going outside. Keep responses natural and dog-like, but still "
    "helpful and friendly."
)

# Dataset Configuration
DATASET_NAME = "bavard/personachat_truecased"  # Persona-Chat dataset
# Alternative: "bavard/personachat" or "personachat"

# Training Configuration
NUM_EPOCHS = 3
BATCH_SIZE = 2  # Per device (reduce to 1-2 for 4GB GPU)
LEARNING_RATE = 2e-4
MAX_LENGTH = 512  # Reduce to 512 for 4GB GPU (2048 for 8GB+)
GRADIENT_ACCUMULATION_STEPS = 4

# LoRA Configuration
LORA_R = 16  # Rank
LORA_ALPHA = 32  # LoRA alpha
LORA_DROPOUT = 0.05
LORA_TARGET_MODULES = ["q_proj", "k_proj", "v_proj", "o_proj"]  # Mistral attention modules

# Output Configuration
OUTPUT_DIR = "./lora-adapters-scooby-dog"

# Quantization (for Colab)
USE_QUANTIZATION = False  # Set to False if you have enough VRAM