data:
  data_dir: "data/MOSES/shapemol/MOSES2_training_val_dataset.pkl"  # Path to the raw data (pickle file)
  save_dir: "./data/MOSES"  # Directory to save the LMDB database
  num_samples: 1024  # Number of point cloud samples generated per molecule
  num_workers: 20  # Number of parallel workers for data processing
  aug_mode: 1
  aug_times: 2
  map_size: 500  # LMDB map_size in bytes (100GB)
  batch_size: 20000  # Batch size for processing data
  seed: 42  # Random seed for reproducibility
  use_smiles: False  # Whether to use SMILES strings for data processing

model:
  surf:
    n: 2
    normal_channel: True
  mlp:
    hidden_dim: 256  # Hidden dimension for MLP
    output_dim: 768  # Output dimension for MLP
    num_layers: 2  # Number of layers in MLP
    dropout_rate: 0.1  # Dropout rate in MLP
    activation_function: "relu"  # Activation function for MLP
  bart:
    max_position_embeddings: 512
    d_model: 768
    encoder_layers: 0
    decoder_layers: 6
    encoder_attention_heads: 0
    decoder_attention_heads: 8
    encoder_ffn_dim: 0
    decoder_ffn_dim: 3072
    activation_function: 'gelu'
  generation_config:
    do_sample: true
    max_length: 512
    top_k: 50
    top_p: 1.0
    temperature: 1.0
    num_return_sequences: 50

train:
  output_dir: "./checkpoints/conditional/surfbartv2-sample1024-merge-angles-0421"  # TODO: Directory to save model and checkpoints
  resume_path: null
  overwrite_output_dir: true  # Whether to overwrite existing outputs
  num_train_epochs: 50  # Total number of training epochs
  per_device_train_batch_size: 150  # Training batch size per device
  per_device_eval_batch_size: 4  # Evaluation batch size per device
  dataloader_num_workers: 2 # Number of workers for data loading
  save_total_limit: 6  # Maximum number of checkpoints to keep
  logging_steps: 50  # Steps interval for logging
  eval_strategy: "steps"  # Evaluation strategy (e.g., "steps", "epoch")
  eval_steps: 5000  # Evaluation frequency (in steps)
  do_eval: true  # Whether to perform evaluation
  learning_rate: 1e-4  # Initial learning rate
  warmup_ratio: 0.1  # Warm-up ratio for learning rate scheduler
  save_strategy: "steps"  # Save strategy (e.g., "steps", "epoch")
  save_steps: 5000  # Save frequency (in steps)
  load_best_model_at_end: true  # Load the best model at the end of training
  logging_first_step: true  # Log the first training step
  bf16: True # Whether to use bf16 precision
  early_stopping_patience: 10  # Patience for early stopping callback
  early_stopping_threshold: 0  # Threshold for early stopping
  seed: 42