Support HuggingFace's inference API (#352)

* Support HuggingFace's inference API * Use positive temperature for prompt parser
neulab · Sep 12, 2023 · 779c7b9 · 779c7b9
1 parent 5ca2858
commit 779c7b9
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/prompt2model/prompt_parser/instr_parser.py b/prompt2model/prompt_parser/instr_parser.py
@@ -93,7 +93,7 @@ def parse_from_prompt(self, prompt: str) -> None:
                 response: openai.ChatCompletion | Exception = (
                     chat_api.generate_one_completion(
                         parsing_prompt_for_chatgpt,
-                        temperature=0,
+                        temperature=0.01,
                         presence_penalty=0,
                         frequency_penalty=0,
                     )

diff --git a/prompt2model/utils/api_tools.py b/prompt2model/utils/api_tools.py
@@ -45,16 +45,19 @@ def __init__(
         self,
         model_name: str = "gpt-3.5-turbo",
         max_tokens: int | None = None,
+        api_base: str | None = None,
     ):
         """Initialize APIAgent with model_name and max_tokens.
 
         Args:
             model_name: Name fo the model to use (by default, gpt-3.5-turbo).
             max_tokens: The maximum number of tokens to generate. Defaults to the max
                 value for the model if available through litellm.
+            api_base: Custom endpoint for Hugging Face's inference API.
         """
         self.model_name = model_name
         self.max_tokens = max_tokens
+        self.api_base = api_base
         if max_tokens is None:
             try:
                 self.max_tokens = litellm.utils.get_max_tokens(model_name)
@@ -99,6 +102,7 @@ def generate_one_completion(
             messages=[
                 {"role": "user", "content": f"{prompt}"},
             ],
+            api_base=self.api_base,
             temperature=temperature,
             presence_penalty=presence_penalty,
             frequency_penalty=frequency_penalty,
@@ -144,6 +148,7 @@ async def _throttled_completion_acreate(
                         return await acompletion(
                             model=model,
                             messages=messages,
+                            api_base=self.api_base,
                             temperature=temperature,
                             max_tokens=max_tokens,
                             n=n,