From fdd7c751a549b56c8f749f44e722efd7810d4d92 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 23 Feb 2024 00:46:41 +0000 Subject: [PATCH] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/fastserve/models/__init__.py | 2 +- src/fastserve/models/vllm.py | 9 ++++----- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/src/fastserve/models/__init__.py b/src/fastserve/models/__init__.py index 4cb0335..8c495da 100644 --- a/src/fastserve/models/__init__.py +++ b/src/fastserve/models/__init__.py @@ -5,4 +5,4 @@ from fastserve.models.llama_cpp import ServeLlamaCpp as ServeLlamaCpp from fastserve.models.sdxl_turbo import ServeSDXLTurbo as ServeSDXLTurbo from fastserve.models.ssd import ServeSSD1B as ServeSSD1B -from fastserve.models.vllm import ServeVLLM as ServeVLLM \ No newline at end of file +from fastserve.models.vllm import ServeVLLM as ServeVLLM diff --git a/src/fastserve/models/vllm.py b/src/fastserve/models/vllm.py index 4303078..db89fe0 100644 --- a/src/fastserve/models/vllm.py +++ b/src/fastserve/models/vllm.py @@ -1,15 +1,13 @@ import logging -import os from typing import Any, List, Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel from fastserve.core import FastServe logger = logging.getLogger(__name__) - class PromptRequest(BaseModel): prompt: str = "Write a python function to resize image to 224x224" temperature: float = 0.8 @@ -50,8 +48,9 @@ def __call__(self, request: PromptRequest) -> Any: from vllm import SamplingParams sampling_params = SamplingParams( - temperature=request.temperature, top_p=request.top_p, - max_tokens=request.max_tokens + temperature=request.temperature, + top_p=request.top_p, + max_tokens=request.max_tokens, ) result = self.llm.generate(request.prompt, sampling_params=sampling_params) logger.info(result)