From 1840919ebdc182320888cbdc1852090396756b43 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 15 May 2024 08:23:00 -0700 Subject: [PATCH] fix(main.py): testing fix --- litellm/main.py | 1 - litellm/tests/test_completion.py | 10 +++++----- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 3429cab4d2f..69ec7985de4 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -14,7 +14,6 @@ from functools import partial import dotenv, traceback, random, asyncio, time, contextvars from copy import deepcopy import httpx - import litellm from ._logging import verbose_logger from litellm import ( # type: ignore diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 4441ddf29a7..2bc167ed099 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -1161,28 +1161,28 @@ HF Tests we should pass # Test util to sort models to TGI, conv, None def test_get_hf_task_for_model(): model = "glaiveai/glaive-coder-7b" - model_type = litellm.llms.huggingface_restapi.get_hf_task_for_model(model) + model_type, _ = litellm.llms.huggingface_restapi.get_hf_task_for_model(model) print(f"model:{model}, model type: {model_type}") assert model_type == "text-generation-inference" model = "meta-llama/Llama-2-7b-hf" - model_type = litellm.llms.huggingface_restapi.get_hf_task_for_model(model) + model_type, _ = litellm.llms.huggingface_restapi.get_hf_task_for_model(model) print(f"model:{model}, model type: {model_type}") assert model_type == "text-generation-inference" model = "facebook/blenderbot-400M-distill" - model_type = litellm.llms.huggingface_restapi.get_hf_task_for_model(model) + model_type, _ = litellm.llms.huggingface_restapi.get_hf_task_for_model(model) print(f"model:{model}, model type: {model_type}") assert model_type == "conversational" model = "facebook/blenderbot-3B" - model_type = litellm.llms.huggingface_restapi.get_hf_task_for_model(model) + model_type, _ = litellm.llms.huggingface_restapi.get_hf_task_for_model(model) print(f"model:{model}, model type: {model_type}") assert model_type == "conversational" # neither Conv or None model = "roneneldan/TinyStories-3M" - model_type = litellm.llms.huggingface_restapi.get_hf_task_for_model(model) + model_type, _ = litellm.llms.huggingface_restapi.get_hf_task_for_model(model) print(f"model:{model}, model type: {model_type}") assert model_type == "text-generation"