#102 vLLM subclass implementation as inference backend and example updated.

datatune/llm/llm.py

@@ -329,32 +329,27 @@ def __init__(
         kwargs.update({"api_key": api_key})
         super().__init__(model_name=f"openai/{model_name}", **kwargs)
 
-class VLLM(LLM):
-
-
 
-    def __init__(
-        self,
-        model_name: str,
-        api_base: str = "http://localhost:8000/v1",
-        **kwargs
+class VLLM(LLM):
+    def __init__(self,
+                model_name: str,
+                api_base: str = "http://localhost:8000/v1",
+                max_tokens: int = 4096,
+                **kwargs
     ):
-        import httpx
-
-        resp = httpx.get(f"{api_base}/models")
-        max_model_len = resp.json()["data"][0]["max_model_len"]
 
         kwargs.update({
             "api_base": api_base,
             "api_key": "dummy",
-            "max_tokens": max_model_len
-
+            "max_tokens": max_tokens
         })
-
         super().__init__(
             model_name=f"openai/{model_name}",
             **kwargs
         )
+        self.api_base = api_base
+        self._model_name = model_name
+
 
 class Azure(LLM):
     def __init__(

examples/Getting_started.ipynb

@@ -1,6 +1,7 @@
 {
   "cells": [
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "OHg0Xg632UxE"
@@ -10,6 +11,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "TjDiMfzi2UxF"
@@ -23,17 +25,18 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
+        "id": "aISH8Ve42UxF",
         "vscode": {
           "languageId": "plaintext"
-        },
-        "id": "aISH8Ve42UxF"
+        }
       },
       "outputs": [],
       "source": [
         "!pip install datatune"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "Gd16ZTSG2UxG"
@@ -52,12 +55,13 @@
       "source": [
         "import pandas as pd\n",
         "import dask.dataframe as dd\n",
-        "from datatune.llm.llm import Azure, Ollama, Gemini\n",
+        "from datatune.llm.llm import Azure, Ollama, Gemini, VLLM\n",
         "import datatune as dt\n",
         "import os"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "mtc44gBs2UxG"
@@ -72,10 +76,10 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
+        "id": "Kts_-l7S2UxG",
         "vscode": {
           "languageId": "shellscript"
-        },
-        "id": "Kts_-l7S2UxG"
+        }
       },
       "outputs": [],
       "source": [
@@ -90,48 +94,50 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
-      "source": [
-        "## Else, If you are using Ollama for your local models, here are the initial steps"
-      ],
       "metadata": {
         "id": "yIbbiJKi5ou5"
-      }
+      },
+      "source": [
+        "## Else, If you are using Ollama for your local models, here are the initial steps"
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ugLCh_BMqB3z"
+      },
+      "outputs": [],
       "source": [
         "\n",
         "!apt-get update && apt-get install -y wget\n",
         "!wget https://ollama.ai/install.sh -O install_ollama.sh\n",
         "!chmod +x install_ollama.sh\n",
         "!OLLAMA_USE_SYSTEM_CA_CERTS=1 ./install_ollama.sh\n",
         "!ollama serve > /var/log/ollama.log 2>&1 &"
-      ],
-      "metadata": {
-        "id": "ugLCh_BMqB3z"
-      },
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "code",
-      "source": [
-        "# Now pull your custom Ollama Model\n",
-        "!ollama pull qwen2.5vl:7b # or provide your own model"
-      ],
+      "execution_count": null,
       "metadata": {
         "id": "2PmxvVrjqOed"
       },
-      "execution_count": null,
-      "outputs": []
+      "outputs": [],
+      "source": [
+        "# Now pull your custom Ollama Model\n",
+        "!ollama pull qwen2.5vl:7b # or provide your own model"
+      ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
-      "source": [],
       "metadata": {
         "id": "pbYrSgXd7qdJ"
-      }
+      },
+      "source": []
     },
     {
       "cell_type": "code",
@@ -157,36 +163,56 @@
     },
     {
       "cell_type": "code",
+      "execution_count": 1,
+      "metadata": {
+        "id": "3MYHcBYvlNcv"
+      },
+      "outputs": [],
       "source": [
         "# if you're using Gemini, uncomment the following lines instead:\n",
         "# llm = Gemini(\n",
         "#    model_name=\"gemma-3-1b-it\", #or use your custom model\n",
         "#    api_key=os.getenv(\"GEMINI_API_KEY\")\n",
         "#)"
-      ],
-      "metadata": {
-        "id": "3MYHcBYvlNcv"
-      },
-      "execution_count": 1,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "DREPlKRa6qne"
+      },
+      "outputs": [],
       "source": [
         "# or if you're using local models, uncomment the following lines to use Ollama\n",
         "\n",
         "# llm = Ollama(\n",
         "#    model_name=\"llama3:15b\",\n",
         "#    api_base=\"http://ollama-server:11434\"\n",
         "#)"
-      ],
-      "metadata": {
-        "id": "DREPlKRa6qne"
-      },
+      ]
+    },
+    {
+      "cell_type": "code",
       "execution_count": null,
-      "outputs": []
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# or if you're using vLLM for high-throughput local inference:\n",
+        "# Before running, start the vLLM server in your terminal:\n",
+        "#   pip install vllm\n",
+        "#   python -m vllm.entrypoints.openai.api_server --model mistralai/Mistral-7B-Instruct-v0.1 --port 8000\n",
+        "\n",
+        "# Then uncomment the following:\n",
+        "# llm = VLLM(\n",
+        "#    model_name=\"mistralai/Mistral-7B-Instruct-v0.1\",\n",
+        "#    api_base=\"http://localhost:8000/v1\",\n",
+        "#    max_tokens=4096  # set to your model's context length\n",
+        "# )"
+      ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "is47NzrK2UxI"
@@ -254,6 +280,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "9Tgpxs0k2UxJ"
@@ -276,6 +303,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "3BlIyZbd2UxK"
@@ -308,6 +336,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "FoMVaOV02UxK"
@@ -340,6 +369,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "I23NZ6YW2UxL"
@@ -371,6 +401,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "ZmTRxuAN2UxL"
@@ -402,6 +433,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "pnQ1Wl9h2UxM"
@@ -441,6 +473,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "MSd-vqQN2UxN"
@@ -474,6 +507,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "261X7XOk2UxO"
@@ -513,17 +547,17 @@
     }
   ],
   "metadata": {
-    "language_info": {
-      "name": "python"
-    },
     "colab": {
       "provenance": []
     },
     "kernelspec": {
-      "name": "python3",
-      "display_name": "Python 3"
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
     }
   },
   "nbformat": 4,
   "nbformat_minor": 0
-}
+}