From 3d660ad9386d8b4b067b676e44912e02607847a3 Mon Sep 17 00:00:00 2001
From: Celina Hanouti <hanouticelina@gmail.com>
Date: Mon, 9 Sep 2024 18:30:34 +0200
Subject: [PATCH] Rename TGI Adapter class

---
 llama_toolchain/inference/adapters/tgi/__init__.py | 6 +++---
 llama_toolchain/inference/adapters/tgi/tgi.py      | 4 ++--
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/llama_toolchain/inference/adapters/tgi/__init__.py b/llama_toolchain/inference/adapters/tgi/__init__.py
index 3185e04dc..6c9e128bc 100644
--- a/llama_toolchain/inference/adapters/tgi/__init__.py
+++ b/llama_toolchain/inference/adapters/tgi/__init__.py
@@ -5,19 +5,19 @@
 # the root directory of this source tree.
 
 from .config import TGIImplConfig
-from .tgi import InferenceEndpointAdapter, LocalTGIAdapter
+from .tgi import InferenceEndpointAdapter, TGIAdapter
 
 
 async def get_adapter_impl(config: TGIImplConfig, _deps):
     assert isinstance(config, TGIImplConfig), f"Unexpected config type: {type(config)}"
 
     if config.is_local_tgi():
-        impl = LocalTGIAdapter(config)
+        impl = TGIAdapter(config)
     elif config.is_inference_endpoint():
         impl = InferenceEndpointAdapter(config)
     else:
         raise ValueError(
-            "Invalid configuration. Specify either a local URL or Inference Endpoint details."
+            "Invalid configuration. Specify either an URL or HF Inference Endpoint details (namespace and endpoint name)."
         )
 
     await impl.initialize()
diff --git a/llama_toolchain/inference/adapters/tgi/tgi.py b/llama_toolchain/inference/adapters/tgi/tgi.py
index 71eb7605c..ba188a539 100644
--- a/llama_toolchain/inference/adapters/tgi/tgi.py
+++ b/llama_toolchain/inference/adapters/tgi/tgi.py
@@ -30,7 +30,7 @@ HF_SUPPORTED_MODELS = {
 }
 
 
-class LocalTGIAdapter(Inference):
+class TGIAdapter(Inference):
 
     def __init__(self, config: TGIImplConfig) -> None:
         self.config = config
@@ -232,7 +232,7 @@ class LocalTGIAdapter(Inference):
             )
 
 
-class InferenceEndpointAdapter(LocalTGIAdapter):
+class InferenceEndpointAdapter(TGIAdapter):
     def __init__(self, config: TGIImplConfig) -> None:
         super().__init__(config)
         self.config.url = f"https://api.endpoints.huggingface.cloud/v2/endpoint/{config.hf_namespace}/{config.hf_endpoint_name}"