scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(main.py): testing fix

Browse source
This commit is contained in:
Krrish Dholakia 2024-05-15 08:23:00 -07:00
parent 1a3b001432
commit 1840919ebd
2 changed files with 5 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/main.py
View file

@ -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

10
litellm/tests/test_completion.py
View file

@ -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"