config_physics.py

import torch

class Config:
    # Model
    MODEL_ID = "unsloth/gemma-3-1b-it" 
    # MODEL_ID = "google/gemma-3-1b-it" 
    
    # Dimensions
    HIDDEN_SIZE = 1152 # Gemma 3 1B hidden size
    LATENT_DIM = 256   # Physics latent state dimension
    PROJECTOR_HIDDEN = 1024
    
    # Physics Controller
    CONTROLLER_HIDDEN = 512
    MODULATION_DIM = 64 # Rank of modulation (similar to LoRA rank)
    
    # Training (OPTIMIZED for Contrastive Physics)
    BATCH_SIZE = 1          # Back to 1: Contrastive needs 4 forward passes!
    GRAD_ACCUMULATION = 64  # 1 * 64 = 64 effective batch
    LEARNING_RATE = 5e-5    # Reduced: Controller needs stability with Lazy Tax
    POLICY_LR = 5e-5        # Sync with LR
    EPOCHS = 3              # Reduced: Contrastive Loss accelerates learning
    
    # TTT / Inference
    TTT_STEPS = 5
    TTT_LR = 1e-4
    
    # Data
    MAX_LENGTH = 256 # Reduced: Math reasoning is shorter. Saves VRAM.
    DTYPE = torch.float32
    
    # Physics Dimensions
    PHYSICS_DIMS = [
        "Gravity", "Friction", "Elasticity", "Fragility", "Density", 
        "Temperature", "Conductivity", "Magnetism",
        "Thermodynamics", "Fluid Dynamics", "Vacuum Physics", "Electromagnetism"
    ]
    
    # Keys
    API_KEY_ENV = "GEMINI_API_KEY"
    
    # Text Generation / Prompting
    SYSTEM_PROMPT = (
        "You are Flux, an advanced physics simulation engine. "
        "You answer questions based on precise physical laws, distinguishing between environments like "
        "Vacuum, Earth, Moon, and Zero-G. "
        "Think step-by-step using first principles (Newton's Laws, Gravity, Fluid Dynamics)."
    )