#!bin/bash

# Make non interactive
export DEBIAN_FRONTEND=noninteractive

echo "---------------------------------------------------"
echo "Starting vLLM docker container"
echo "---------------------------------------------------"  

# Construct the Docker run command
export HF_TOKEN="[insert-hf-token-with-gated-access-here]"
MODEL_NAME="mistralai/Pixtral-Large-Instruct-2411"
NUM_GPUS=$(nvidia-smi -L | wc -l)


DOCKER_CMD="docker run -d --gpus all \
    -v /ephemeral/.cache/huggingface:/root/.cache/huggingface \
    -v /home/ubuntu/vllm:/vllm_repo \
    -e HF_TOKEN=$HF_TOKEN \
    -p 8000:8000 \
    --ipc=host \
    --restart always"

DOCKER_CMD="$DOCKER_CMD vllm/vllm-openai:latest \
    --tensor-parallel-size $NUM_GPUS \
    --model \"$MODEL_NAME\" \
    --chat-template examples/tool_chat_template_mistral.jinja \
    --config-format mistral \
    --load-format mistral \
    --tokenizer_mode mistral \
    --gpu-memory-utilization 0.9 \
    --max-model-len 30000 \
    --limit_mm_per_prompt 'image=10'"


# Run the Docker command as ubuntu user
echo "Executing Docker command: $DOCKER_CMD"
sudo -u ubuntu bash -c "$DOCKER_CMD"


# Query the model using an image
# IMAGE_URL="https://www.hyperstack.cloud/hs-fs/hubfs/deploy-vm-11-ecd8c53003182041d3a2881d0010f6c6-1.png?width=3352&height=1852&name=deploy-vm-11-ecd8c53003182041d3a2881d0010f6c6-1.png"
# cat <<EOF > payload.json
# {
#     "model": "mistralai/Pixtral-Large-Instruct-2411",
#     "messages": [
#         {
#             "role": "system",
#             "content": "SYSTEM_PROMPT"
#         },
#         {
#             "role": "user",
#             "content": [
#                 {
#                     "type": "text",
#                     "text": "Describe this image in two sentences"
#                 },
#                 {
#                     "type": "image_url",
#                     "image_url": {
#                         "url": "${IMAGE_URL}"
#                     }
#                 }
#             ]
#         }
#     ]
# }
# EOF

# expected response:
# {
#   "id": "chatcmpl-2b41dd5f00c44692a647f27c6526f397",
#   "object": "chat.completion",
#   "created": 1732088507,
#   "model": "mistralai/Pixtral-Large-Instruct-2411",
#   "choices": [
#     {
#       "index": 0,
#       "message": {
#         "role": "assistant",
#         "content": "The image shows a virtual machine management interface from Hyperstack, where the \"genius-hubble\" VM is active. In the networking settings, the option to enable SSH access on port 22 is highlighted, and there are additional options for ICMP access and managing firewalls.",
#         "tool_calls": []
#       },
#       "logprobs": null,
#       "finish_reason": "stop",
#       "stop_reason": null
#     }
#   ],
#   "usage": {
#     "prompt_tokens": 2357,
#     "total_tokens": 2419,
#     "completion_tokens": 62,
#     "prompt_tokens_details": null
#   },
#   "prompt_logprobs": null
# }


# Use the JSON payload file in the curl command
# curl -X POST http://localhost:8000/v1/chat/completions \
#     -H "Content-Type: application/json" \
#     -d @payload.json

# # Query the model only with text
# cat <<EOF > payload.json
# {
#     "model": "mistralai/Pixtral-Large-Instruct-2411",
#     "messages": [
#         {
#             "role": "system",
#             "content": "SYSTEM_PROMPT"
#         },
#         {
#             "role": "user",
#             "content": "Hi. What can you do for me?"
#         }
#     ]
# }
# EOF

# # Use the JSON payload file in the curl command
# curl -X POST http://localhost:8000/v1/chat/completions \
#     -H "Content-Type: application/json" \
#     -d @payload.json

# expected response
# {"id":"chatcmpl-eaadc212af354361b46ed65366f9f7a7","object":"chat.completion","created":1732088556,"model":"mistralai/Pixtral-Large-Instruct-2411","choices":[{"index":0,"message":{"role":"assistant","content":" Hello! I can assist you with a wide range of tasks and provide information on various topics. Here are some examples:\n\n1. **Answer Questions**: I can provide information based on the data I've been trained on, up until 2023.\n\n2. **Explain Concepts**: I can help break down complex ideas into simpler parts to make them easier to understand.\n\n3. **Provide Suggestions**: Whether it's a book to read, a movie to watch, or a recipe to cook, I can provide recommendations.\n\n4. **Help with Language**: I can help with language translation, definition, or grammar.\n\n5. **Perform Simple Tasks**: I can do simple calculations, conversions, and other basic tasks.\n\n6. **Engage in Dialogue**: I can participate in conversations on a wide range of topics.\n\nWhat specifically would you like help with?","tool_calls":[]},"logprobs":null,"finish_reason":"stop","stop_reason":null}],"usage":{"prompt_tokens":20,"total_tokens":221,"completion_tokens":201,"prompt_tokens_details":null},"prompt_logprobs":null}