Improve TGI adapter initialization condition

This commit is contained in:
Celina Hanouti 2024-09-10 18:22:09 +02:00
parent 2ac8e7b901
commit 0964b0a74a
2 changed files with 2 additions and 4 deletions

View file

@ -11,7 +11,7 @@ from .tgi import InferenceEndpointAdapter, TGIAdapter
async def get_adapter_impl(config: TGIImplConfig, _deps): async def get_adapter_impl(config: TGIImplConfig, _deps):
assert isinstance(config, TGIImplConfig), f"Unexpected config type: {type(config)}" assert isinstance(config, TGIImplConfig), f"Unexpected config type: {type(config)}"
if config.is_local_tgi(): if config.url is not None:
impl = TGIAdapter(config) impl = TGIAdapter(config)
elif config.is_inference_endpoint(): elif config.is_inference_endpoint():
impl = InferenceEndpointAdapter(config) impl = InferenceEndpointAdapter(config)

View file

@ -7,6 +7,7 @@
from typing import Optional from typing import Optional
from huggingface_hub import HfApi from huggingface_hub import HfApi
from llama_models.schema_utils import json_schema_type from llama_models.schema_utils import json_schema_type
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
@ -31,6 +32,3 @@ class TGIImplConfig(BaseModel):
def get_namespace(self) -> str: def get_namespace(self) -> str:
return HfApi().whoami()["name"] return HfApi().whoami()["name"]
def is_local_tgi(self) -> bool:
return self.url is not None and self.url.startswith("http://localhost")