add environment.yml

.gitignore

@@ -14,4 +14,5 @@ tr2d2-pep/utils/__pycache__/
 *.log
 *.ipynb
 *.pt
-tr2d2-pep/scoring/functions/classifiers/best_model.pt
+tr2d2-pep/scoring/functions/classifiers/best_model.pt
+.gitignore

README.md

@@ -4,7 +4,7 @@
 
 [**Sophia Tang**](https://sophtang.github.io/)\*, [**Yuchen Zhu**](https://yuchen-zhu-zyc.github.io/)\*, [**Molei Tao**](https://mtao8.math.gatech.edu/), and [**Pranam Chatterjee**](https://www.chatterjeelab.com/)
 
-![TR2-D2](assets/anim-good.gif)
+![TR2-D2](assets/tr2d2-anim.gif)
 
 This is the repository for **[TR2-D2: Tree Search Guided Trajectory-Aware Fine-Tuning for Discrete Diffusion](https://arxiv.org/abs/2509.25171)** 🤖🌳. It is partially built on the **[PepTune repo](https://github.com/programmablebio/peptune)** ([Tang et al. 2024](https://arxiv.org/abs/2412.17780)) and **MDNS** ([Zhu et al. 2025](https://arxiv.org/abs/2508.10684)).
 

tr2d2-pep/environment.yml

@@ -0,0 +1,21 @@
+name: tr2d2-pep
+channels:
+  - conda-forge
+dependencies:
+  - python=3.10
+  - pip
+  - rdkit
+  - numpy
+  - pandas
+  - scikit-learn
+  - jupyterlab
+  - matplotlib-base
+  - seaborn
+  - tqdm
+  - pyyaml
+  - torchvision
+  - pytorch-cuda=12.4
+  - pip:
+      - pytorch-lightning==2.5.5
+      - fair-esm==2.0.0
+      - transformers==4.56.2