diff --git a/.vscode/settings.json b/.vscode/settings.json
new file mode 100644
index 0000000..f2d90cb
--- /dev/null
+++ b/.vscode/settings.json
@@ -0,0 +1,3 @@
+{
+    "python.linting.enabled": true
+}
\ No newline at end of file
diff --git a/Dockerfile b/Dockerfile
index f1b959d..94a5456 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -30,6 +30,17 @@ RUN --mount=type=cache,target=/root/.cache/pip \
 ADD src .
 
 # Quick temporary updates
-RUN pip install git+https://github.com/runpod/runpod-python@multijob2#egg=runpod       --compile
+RUN pip install git+https://github.com/runpod/runpod-python@multijob2#egg=runpod --compile
 
+# Prepare the models inside the docker image
+ARG HUGGING_FACE_HUB_TOKEN=NONE
+ENV HUGGING_FACE_HUB_TOKEN=$HUGGING_FACE_HUB_TOKEN
+ENV DOWNLOAD_7B_MODEL=YES
+# ENV DOWNLOAD_13B_MODEL=1
+
+# Download the models
+RUN mkdir -p /model
+RUN DOWNLOAD_7B_MODEL=$DOWNLOAD_7B_MODEL HUGGING_FACE_HUB_TOKEN=$HUGGING_FACE_HUB_TOKEN python -u /download_model.py
+
+# Start the handler
 CMD python -u /handler.py
diff --git a/builder/requirements.txt b/builder/requirements.txt
index 1d5c4fb..c83ebc2 100644
--- a/builder/requirements.txt
+++ b/builder/requirements.txt
@@ -6,4 +6,5 @@
 # vllm @ git+https://github.com/vllm-project/vllm.git@2b7d3aca2e1dd25fe26424f57c051af3b823cd71
 # runpod @ git+https://github.com/runpod/runpod-python@vllm#egg=runpod
 vllm==0.1.2
-runpod @ git+https://github.com/runpod/runpod-python@multijob#egg=runpod
+huggingface-hub==0.16.4
+runpod @ git+https://github.com/runpod/runpod-python@multijob2#egg=runpod
diff --git a/src/download_model.py b/src/download_model.py
new file mode 100644
index 0000000..14daa76
--- /dev/null
+++ b/src/download_model.py
@@ -0,0 +1,23 @@
+import os
+from huggingface_hub import snapshot_download
+
+# Get the hugging face token
+HUGGING_FACE_HUB_TOKEN = os.environ['HUGGING_FACE_HUB_TOKEN']
+DOWNLOAD_7B_MODEL = os.environ.get('DOWNLOAD_7B_MODEL', None)
+DOWNLOAD_13B_MODEL = os.environ.get('DOWNLOAD_13B_MODEL', None)
+
+# Download the 7B
+if HUGGING_FACE_HUB_TOKEN and DOWNLOAD_7B_MODEL:
+    snapshot_download(
+            "meta-llama/Llama-2-7b-chat-hf",
+            local_dir="/model/Llama-2-7b-chat-hf",
+            token=HUGGING_FACE_HUB_TOKEN
+    )
+
+# Download the 13B
+if HUGGING_FACE_HUB_TOKEN and DOWNLOAD_13B_MODEL:
+    snapshot_download(
+            "meta-llama/Llama-2-13b-chat-hf",
+            local_dir="/model/Llama-2-13b-chat-hf",
+            token=HUGGING_FACE_HUB_TOKEN
+    )
diff --git a/src/handler.py b/src/handler.py
index 4bf7352..9eb680c 100644
--- a/src/handler.py
+++ b/src/handler.py
@@ -1,23 +1,19 @@
 #!/usr/bin/env python
 ''' Contains the handler function that will be called by the serverless. '''
-import json
-import types
-from typing import AsyncGenerator, Dict
 
 # Start the VLLM serving layer on our RunPod worker.
 from vllm import AsyncLLMEngine, SamplingParams, AsyncEngineArgs
 from vllm.utils import random_uuid
 import runpod
-import asyncio
 
 # Prepare the model and tokenizer
-MODEL = 'facebook/opt-125m'
-# TOKENIZER = 'hf-internal-testing/llama-tokenizer'
+MODEL = '/model/Llama-2-7b-chat-hf'
+TOKENIZER = 'hf-internal-testing/llama-tokenizer'
 
 # Prepare the engine's arguments
 engine_args = AsyncEngineArgs(
     model=MODEL,
-    #tokenizer=TOKENIZER,
+    tokenizer=TOKENIZER,
     tokenizer_mode= "auto",
     tensor_parallel_size= 1,
     dtype = "auto",
@@ -28,7 +24,6 @@ engine_args = AsyncEngineArgs(
 # Create the vLLM asynchronous engine
 llm = AsyncLLMEngine.from_engine_args(engine_args)
 
-
 def handler_fully_utilized() -> bool:
     # Compute pending sequences
     total_pending_sequences = len(llm.engine.scheduler.waiting) + len(llm.engine.scheduler.swapped)
@@ -109,6 +104,12 @@ async def handler(job):
     else:
         sampling_params = SamplingParams()
 
+    # Print the job input
+    print(job_input)
+
+    # Print the sampling params
+    print(sampling_params)
+
     # Send request to VLLM
     request_id = random_uuid()
     results_generator = llm.generate(prompt, sampling_params, request_id)