From d6286d0a25f00a828dc5e882a0767dbfa71cb552 Mon Sep 17 00:00:00 2001 From: Dylan Stark Date: Thu, 2 Nov 2023 10:19:26 +0000 Subject: [PATCH] Adds support for setting additional client args. --- easyllm/clients/huggingface.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/easyllm/clients/huggingface.py b/easyllm/clients/huggingface.py index 4ad6099..7ad72a5 100644 --- a/easyllm/clients/huggingface.py +++ b/easyllm/clients/huggingface.py @@ -38,6 +38,7 @@ ) api_base = os.environ.get("HUGGINGFACE_API_BASE", None) or "https://api-inference.huggingface.co/models" api_version = os.environ.get("HUGGINGFACE_API_VERSION", None) or "2023-07-29" +client_args = {} prompt_builder = os.environ.get("HUGGINGFACE_PROMPT", None) stop_sequences = [] seed = 42 @@ -162,7 +163,7 @@ def create( url = api_base # create the client - client = InferenceClient(url, token=api_key) + client = InferenceClient(url, token=api_key, **client_args) # create stop sequences if isinstance(request.stop, list): @@ -355,7 +356,7 @@ def create( url = api_base # create the client - client = InferenceClient(url, token=api_key) + client = InferenceClient(url, token=api_key, **client_args) # create stop sequences if isinstance(request.stop, list): @@ -471,7 +472,7 @@ def create( url = api_base # create the client - client = InferenceClient(url, token=api_key) + client = InferenceClient(url, token=api_key, **client_args) # client is currently not supporting batched request thats why we run sequentially emb = []