Reorganize root modules into models/ training/ mcts/ packages

Move library modules out of the cluttered root into role-based folders:
- models/: diffusion.py, roformer.py, noise_schedule.py
- training/: finetune_utils.py, distributed_utils.py
- mcts/: peptide_mcts.py

Entry points (inference.py, finetune_multi_target.py) stay at root.
Update all imports across td3b/, baselines/, and entry scripts.
Fix pre-existing broken import (finetune_distributed_utils -> training.distributed_utils).
Update README Code Structure to match. Add .gitignore for __pycache__.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+*.pyc

README.md

@@ -46,12 +46,16 @@ TD3B/
 TD3B/
 ├── inference.py                 # Generate binders (main inference entry point)
 ├── finetune_multi_target.py     # Multi-target TD3B training
-├── finetune_utils.py            # Training utilities
 ├── launch_multi_target.sh       # Training launcher script
-├── diffusion.py                 # MDLM backbone (TR2-D2)
-├── roformer.py                  # RoFormer wrapper
-├── noise_schedule.py            # Noise schedules
-├── peptide_mcts.py              # MCTS tree search
+├── models/
+│   ├── diffusion.py             # MDLM backbone (TR2-D2)
+│   ├── roformer.py              # RoFormer wrapper
+│   └── noise_schedule.py        # Noise schedules
+├── training/
+│   ├── finetune_utils.py        # Training utilities
+│   └── distributed_utils.py     # Distributed training helpers
+├── mcts/
+│   └── peptide_mcts.py          # MCTS tree search
 ├── td3b/
 │   ├── direction_oracle.py      # Direction Oracle (f_φ)
 │   ├── td3b_scoring.py          # Gated reward R = g_ψ · σ(d*·(f_φ−0.5)/τ)

baselines/baselines.py

@@ -316,7 +316,7 @@ class PepTuneSampler:
         pareto_max_size: Optional[int],
         eps: float,
     ):
-        from peptide_mcts import Node, updateParetoFront
+        from mcts.peptide_mcts import Node, updateParetoFront
         from utils.app import PeptideAnalyzer
 
         self.base_model = base_model

baselines/run_mcts_tr2d2.py

@@ -14,7 +14,7 @@ ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
 if ROOT_DIR not in sys.path:
     sys.path.insert(0, ROOT_DIR)
 
-from diffusion import Diffusion
+from models.diffusion import Diffusion
 from configs.finetune_config import (
     DiffusionConfig,
     RoFormerConfig,
@@ -25,8 +25,8 @@ from configs.finetune_config import (
     OptimConfig,
     MCTSConfig,
 )
-from finetune_utils import load_tokenizer
-from finetune_distributed_utils import setup_distributed, cleanup_distributed, is_main_process
+from training.finetune_utils import load_tokenizer
+from training.distributed_utils import setup_distributed, cleanup_distributed, is_main_process
 from scoring.functions.binding import MultiTargetBindingAffinity, TargetSpecificBindingAffinity
 from td3b.direction_oracle import DirectionalOracle
 from finetune_multi_target_tr2d2_ddp import TR2D2GatedReward, TargetDataset, create_tr2d2_mcts

baselines/run_validation_td3b.py

@@ -14,7 +14,7 @@ ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
 if ROOT_DIR not in sys.path:
     sys.path.insert(0, ROOT_DIR)
 
-from diffusion import Diffusion
+from models.diffusion import Diffusion
 from configs.finetune_config import (
     DiffusionConfig,
     RoFormerConfig,
@@ -25,9 +25,9 @@ from configs.finetune_config import (
     OptimConfig,
     MCTSConfig,
 )
-from finetune_utils import load_tokenizer, create_reward_function
+from training.finetune_utils import load_tokenizer, create_reward_function
 from finetune_multi_target import TargetDataset
-from distributed_utils import setup_distributed, cleanup_distributed, is_main_process
+from training.distributed_utils import setup_distributed, cleanup_distributed, is_main_process
 from scoring.functions.binding import MultiTargetBindingAffinity, TargetSpecificBindingAffinity
 from td3b.direction_oracle import DirectionalOracle
 from utils.app import PeptideAnalyzer

baselines/sampling_setup.py

@@ -14,7 +14,7 @@ import torch
 from hydra import compose, initialize_config_dir
 from hydra.core.global_hydra import GlobalHydra
 
-from diffusion import Diffusion
+from models.diffusion import Diffusion
 from scoring.scoring_functions import ScoringFunctions
 from scoring.functions.binding import MultiTargetBindingAffinity
 from td3b.direction_oracle import DirectionalOracle, resolve_device

finetune_multi_target.py

@@ -33,7 +33,7 @@ from tqdm import tqdm
 # Add project root to path
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 
-from diffusion import Diffusion
+from models.diffusion import Diffusion
 from tokenizer.my_tokenizers import SMILES_SPE_Tokenizer
 from utils.app import PeptideAnalyzer
 from scoring.functions.binding import MultiTargetBindingAffinity, TargetSpecificBindingAffinity
@@ -60,7 +60,7 @@ from configs.finetune_config import (
 )
 
 # Import shared utilities
-from finetune_utils import (
+from training.finetune_utils import (
     load_tokenizer,
     initialize_device,
     create_output_directory,
@@ -728,7 +728,7 @@ def main():
     )
 
     # WDCE loss
-    from finetune_utils import loss_wdce
+    from training.finetune_utils import loss_wdce
 
     logger.info("\n[4/6] Setting up training...")
     policy_model.train()

inference.py

@@ -24,12 +24,12 @@ ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
 if ROOT_DIR not in sys.path:
     sys.path.insert(0, ROOT_DIR)
 
-from diffusion import Diffusion
+from models.diffusion import Diffusion
 from configs.finetune_config import (
     DiffusionConfig, RoFormerConfig, NoiseConfig,
     TrainingConfig, SamplingConfig, EvalConfig, OptimConfig, MCTSConfig,
 )
-from finetune_utils import load_tokenizer, create_reward_function
+from training.finetune_utils import load_tokenizer, create_reward_function
 from td3b.direction_oracle import DirectionalOracle
 from td3b.td3b_scoring import create_td3b_reward_function
 from utils.app import PeptideAnalyzer

peptide_mcts.py → mcts/peptide_mcts.py

@@ -7,7 +7,7 @@ from utils.app import PeptideAnalyzer
 from utils.timer import StepTimer
 from scoring.scoring_functions import ScoringFunctions
 
-import noise_schedule
+from models import noise_schedule
 
 ### for peptide multi-objective ###
 def dominates(a, b):

diffusion.py → models/diffusion.py

@@ -15,9 +15,9 @@ import gc
 import utils.utils as utils
 
 from tokenizer.my_tokenizers import SMILES_SPE_Tokenizer
-import noise_schedule
+from models import noise_schedule
 from torch.optim.lr_scheduler import _LRScheduler
-import roformer as roformer 
+from models import roformer
 from utils.app import PeptideAnalyzer
 import pandas as pd
 

td3b/direction_oracle.py

@@ -20,7 +20,7 @@ if PROJECT_ROOT not in sys.path:
     sys.path.insert(0, PROJECT_ROOT)
 
 from tokenizer.my_tokenizers import SMILES_SPE_Tokenizer
-from roformer import Roformer
+from models.roformer import Roformer
 
 logger = logging.getLogger(__name__)
 

td3b/td3b_finetune.py

@@ -7,7 +7,7 @@ import numpy as np
 import torch
 import wandb
 import os
-from finetune_utils import loss_wdce
+from training.finetune_utils import loss_wdce
 from .td3b_losses import TD3BTotalLoss, extract_embeddings_from_mdlm
 from tqdm import tqdm
 import pandas as pd
@@ -66,7 +66,7 @@ def td3b_finetune(
     print("[TD3B] Creating reference model for KL regularization...")
 
     # Import Diffusion class
-    from diffusion import Diffusion
+    from models.diffusion import Diffusion
 
     # Create new instance with same config
     reference_model = Diffusion(

td3b/td3b_mcts.py

@@ -5,7 +5,7 @@ Extends the base MCTS to support directional rewards and confidence weighting.
 
 import numpy as np
 import torch
-from peptide_mcts import MCTS as BaseMCTS
+from mcts.peptide_mcts import MCTS as BaseMCTS
 from .td3b_scoring import TD3BRewardFunction, TD3BConfidenceWeighting
 
 
@@ -144,7 +144,7 @@ class TD3B_MCTS(BaseMCTS):
             }
 
             # Pareto dominance filtering (same as base class)
-            from peptide_mcts import dominated_by, dominates
+            from mcts.peptide_mcts import dominated_by, dominates
 
             if any(dominated_by(sv, bi["score_vector"]) for bi in self.buffer):
                 self._debug_buffer_decision(sv, "rejected_dominated")

finetune_utils.py → training/finetune_utils.py

@@ -15,7 +15,7 @@ import wandb
 from torch.utils.data import DataLoader, TensorDataset
 from tqdm import tqdm
 
-from diffusion import Diffusion
+from models.diffusion import Diffusion
 from td3b.td3b_mcts import create_td3b_mcts
 from td3b.td3b_scoring import TD3BRewardFunction
 from tokenizer.my_tokenizers import SMILES_SPE_Tokenizer