Update requirements.txt

requirements.txt

@@ -1,25 +1,31 @@
-absl-py==0.13.0
-git+https://github.com/biopython/biopython.git@biopython-180
-git+https://github.com/Acellera/moleculekit.git@1.5.9
-chex==0.0.7
-dm-haiku==0.0.5
-dm-tree==0.1.6
-docker==5.0.0
-immutabledict==2.0.0
-https://storage.googleapis.com/jax-releases/jax/jax-0.3.7.tar.gz
-https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.3.7+cuda11.cudnn805-cp310-none-manylinux2014_x86_64.whl
-ml-collections==0.1.0
-numpy==1.22
-pandas==1.3.4
-scipy
-tensorflow-gpu==2.8.4
-torch
-plotly
-GPUtil
-ray==1.13.0
-tqdm
-protobuf<4
-matplotlib==3.8
+absl-py==1.4.0
+numpy==1.26.0
+pandas==2.3.2
+scipy==1.11.1
+matplotlib==3.7.2
+typing-extensions==4.14.1
+
+tensorflow[and-cuda]==2.20.0
+torch>=2.0.0
+torchvision
+torchaudio
+
+# JAX CUDA 12 compatible
+-f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+jax[cuda12_pip]==0.4.13
+
+chex==0.1.7
+dm-haiku==0.0.9
+dm-tree==0.1.8
+ml-collections==0.1.1
+git+https://github.com/biopython/biopython.git
+git+https://github.com/Acellera/moleculekit.git
 mdtraj==1.9.9
-pydantic<2
--f https://storage.googleapis.com/jax-releases/jax_releases.html
+docker==6.1.3
+immutabledict==3.0.0
+ray==2.6.3
+tqdm
+GPUtil
+plotly
+pydantic==2.11.7
+protobuf>=5.28.0