From 2395fbe035f4ecaeeb81deb0d73dcba94fe02c40 Mon Sep 17 00:00:00 2001 From: Nitish Date: Tue, 9 Jun 2026 22:23:13 +0530 Subject: [PATCH] fix(vllm): simplify VLLM backend, remove httpx dependency, add example (#102) --- datatune/llm/llm.py | 25 +++----- examples/Getting_started.ipynb | 114 +++++++++++++++++++++------------ 2 files changed, 84 insertions(+), 55 deletions(-) diff --git a/datatune/llm/llm.py b/datatune/llm/llm.py index 53e958e..756cc58 100644 --- a/datatune/llm/llm.py +++ b/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__( diff --git a/examples/Getting_started.ipynb b/examples/Getting_started.ipynb index 45b946e..34cdc54 100644 --- a/examples/Getting_started.ipynb +++ b/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,10 +25,10 @@ "cell_type": "code", "execution_count": null, "metadata": { + "id": "aISH8Ve42UxF", "vscode": { "languageId": "plaintext" - }, - "id": "aISH8Ve42UxF" + } }, "outputs": [], "source": [ @@ -34,6 +36,7 @@ ] }, { + "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,16 +94,22 @@ ] }, { + "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", @@ -107,31 +117,27 @@ "!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,21 +163,26 @@ }, { "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", @@ -179,14 +190,29 @@ "# 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 -} \ No newline at end of file +}