updated api server to pass in the prompt token ids in the right forma…

…t depending on the vllm version (#167)

.gitignore

@@ -19,4 +19,7 @@ results_fn_postgres/*.json
 .vscode
 
 # all eda notebooks
-eda_*.ipynb
+eda_*.ipynb
+
+# wandb output (created when running upload_wandb.ipynb)
+wandb/

utils/api_server.py

@@ -10,6 +10,7 @@
 from vllm.engine.async_llm_engine import AsyncLLMEngine
 from vllm.sampling_params import SamplingParams
 from vllm.utils import random_uuid
+from vllm import __version__ as vllm_version
 
 TIMEOUT_KEEP_ALIVE = 5  # seconds.
 app = FastAPI()
@@ -56,12 +57,19 @@ async def generate(request: Request) -> Response:
     if prompt_token_ids[0] != tokenizer.bos_token_id:
         prompt_token_ids = [tokenizer.bos_token_id] + prompt_token_ids
 
-    results_generator = engine.generate(
-        prompt=None,
-        sampling_params=sampling_params,
-        request_id=request_id,
-        prompt_token_ids=prompt_token_ids,
-    )
+    if vllm_version >= "0.4.2":
+        results_generator = engine.generate(
+            inputs={"prompt_token_ids": prompt_token_ids},
+            sampling_params=sampling_params,
+            request_id=request_id,
+        )
+    else:
+        results_generator = engine.generate(
+            prompt=None,
+            sampling_params=sampling_params,
+            request_id=request_id,
+            prompt_token_ids=prompt_token_ids,
+        )
 
     # Streaming case
     async def stream_results() -> AsyncGenerator[bytes, None]: