From b04d20d367e11ea7eb396b77d68ff6364e7a55dd Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 16 Jul 2024 21:21:50 -0700 Subject: [PATCH] fix linting error --- litellm/llms/triton.py | 29 +++++++++++++++++++---------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/litellm/llms/triton.py b/litellm/llms/triton.py index 6d3bebcb34..bc9a15a94b 100644 --- a/litellm/llms/triton.py +++ b/litellm/llms/triton.py @@ -1,14 +1,19 @@ -import os, types +import copy import json -from enum import Enum -import requests, copy # type: ignore +import os import time -from typing import Callable, Optional, List -import litellm -from .prompt_templates.factory import prompt_factory, custom_prompt -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from .base import BaseLLM +import types +from enum import Enum +from typing import Callable, List, Optional + import httpx # type: ignore +import requests # type: ignore + +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + +from .base import BaseLLM +from .prompt_templates.factory import custom_prompt, prompt_factory class TritonError(Exception): @@ -126,8 +131,12 @@ class TritonChatCompletion(BaseLLM): ) @staticmethod - def split_embedding_by_shape(data: list[float], shape: list[int]) -> list[list[float]]: + def split_embedding_by_shape( + data: List[float], shape: List[int] + ) -> List[List[float]]: if len(shape) != 2: raise ValueError("Shape must be of length 2.") embedding_size = shape[1] - return [data[i * embedding_size: (i + 1) * embedding_size] for i in range(shape[0])] + return [ + data[i * embedding_size : (i + 1) * embedding_size] for i in range(shape[0]) + ]