fix vllm version (#172)

Dockerfile

@@ -49,14 +49,14 @@ RUN apt-get -y install git-lfs
 WORKDIR /stage/
 
 RUN pip install --upgrade pip setuptools wheel
-RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
+RUN pip install torch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 --index-url https://download.pytorch.org/whl/cu118
 
 COPY rewardbench rewardbench
 COPY scripts scripts
 COPY setup.py setup.py
 COPY Makefile Makefile
 COPY README.md README.md
-RUN pip install -e .
+RUN pip install -e .[generative]
 RUN chmod +x scripts/*
 
 # this is just very slow
@@ -74,9 +74,5 @@ RUN pip install openai
 RUN pip install together
 RUN pip install google-generativeai
 
-# updated for Gemma 2
-RUN pip install vllm==0.5.1
-# from git+https://github.com/vllm-project/vllm.git@d87f39e9a9dd149f5dd7a58b4d98b21f713827b6
-
 # for interactive session
 RUN chmod -R 777 /stage/

setup.py

@@ -59,7 +59,7 @@
         "tokenizers",
         "torch",
         "tiktoken==0.6.0",  # added for llama 3
-        "transformers==4.40.0",  # pinned at llama 3
+        "transformers==4.43.4",  # pinned at llama 3
         "trl>=0.8.2",  # fixed transformers import error
     ],
     extras_require={