alfraser commited on
Commit
27e1387
·
1 Parent(s): 6000142

Added option to name an inference endpoint so it is human readable

Browse files
Files changed (1) hide show
  1. src/architectures.py +4 -2
src/architectures.py CHANGED
@@ -369,8 +369,9 @@ class HFInferenceEndpoint(ArchitectureComponent):
369
  A concrete pipeline component which sends the user text to a given llama chat based
370
  inference endpoint on HuggingFace
371
  """
372
- def __init__(self, endpoint_url: str, system_prompt: str, max_new_tokens: int, temperature: float = 1.0):
373
  self.endpoint_url: str = endpoint_url
 
374
  self.system_prompt: str = system_prompt
375
  self.max_new_tokens = max_new_tokens
376
  self.api_token = hf_api_token()
@@ -380,7 +381,8 @@ class HFInferenceEndpoint(ArchitectureComponent):
380
  """
381
  Custom config details as markdown
382
  """
383
- desc = f"Endpoint: {self.endpoint_url}; "
 
384
  desc += f"Max tokens: {self.max_tokens}; "
385
  desc += f"Temperature: {self.temperature}; "
386
  desc += f"System prompt: {self.system_prompt}"
 
369
  A concrete pipeline component which sends the user text to a given llama chat based
370
  inference endpoint on HuggingFace
371
  """
372
+ def __init__(self, endpoint_url: str, model_name: str, system_prompt: str, max_new_tokens: int, temperature: float = 1.0):
373
  self.endpoint_url: str = endpoint_url
374
+ self.model_name: str = model_name
375
  self.system_prompt: str = system_prompt
376
  self.max_new_tokens = max_new_tokens
377
  self.api_token = hf_api_token()
 
381
  """
382
  Custom config details as markdown
383
  """
384
+ desc = f"Model: {self.model_name}; "
385
+ desc += f"Endpoint: {self.endpoint_url}; "
386
  desc += f"Max tokens: {self.max_tokens}; "
387
  desc += f"Temperature: {self.temperature}; "
388
  desc += f"System prompt: {self.system_prompt}"