clean up hugging face completion()

litellm/llms/huggingface_restapi.py

@@ -1,5 +1,6 @@
 ## Uses the huggingface text generation inference API
-import os, json
+import os
+import json
 from enum import Enum
 import requests
 import time
@@ -7,7 +8,6 @@ from typing import Callable
 from litellm.utils import ModelResponse
 from typing import Optional
 
-
 class HuggingfaceError(Exception):
     def __init__(self, status_code, message):
         self.status_code = status_code
@@ -16,127 +16,129 @@ class HuggingfaceError(Exception):
             self.message
         )  # Call the base class constructor with the parameters it needs
 
+def validate_environment(api_key):
+    headers = {
+        "content-type": "application/json",
+    }
+    if api_key:
+        headers["Authorization"] = f"Bearer {api_key}"
+    return headers
 
-class HuggingfaceRestAPILLM:
+def completion(
-    def __init__(self, encoding, logging_obj, api_key=None) -> None:
+    model: str,
-        self.encoding = encoding
+    messages: list,
-        self.logging_obj = logging_obj
+    api_base: str,
-        self.validate_environment(api_key=api_key)
+    model_response: ModelResponse,
-
+    print_verbose: Callable,
-    def validate_environment(
+    encoding,
-        self, api_key
+    api_key,
-    ):  # set up the environment required to run the model
+    logging_obj,
-        self.headers = {
+    optional_params=None,
-            "content-type": "application/json",
+    litellm_params=None,
-        }
+    logger_fn=None,
-        # get the api key if it exists in the environment or is passed in, but don't require it
+):
-        self.api_key = api_key
+    headers = validate_environment(api_key)
-        if self.api_key != None:
+    completion_url = ""
-            self.headers["Authorization"] = f"Bearer {self.api_key}"
+    if "https" in model:
-
+        completion_url = model
-    def completion(
+    elif api_base:
-        self,
+        completion_url = api_base
-        model: str,
+    elif "HF_API_BASE" in os.environ:
-        messages: list,
+        completion_url = os.getenv("HF_API_BASE", "")
-        api_base: str,
+    else:
-        model_response: ModelResponse,
+        completion_url = f"https://api-inference.huggingface.co/models/{model}"
-        print_verbose: Callable,
+    prompt = ""
-        optional_params=None,
+    if (
-        litellm_params=None,
+        "meta-llama" in model and "chat" in model
-        logger_fn=None,
+    ):  # use the required special tokens for meta-llama - https://huggingface.co/blog/llama2#how-to-prompt-llama-2
-    ):  # logic for parsing in - calling - parsing out model completion calls
+        prompt = "<s>"
-        completion_url: str = ""
+        for message in messages:
-        if "https" in model:
+            if message["role"] == "system":
-            completion_url = model
+                prompt += "[INST] <<SYS>>" + message["content"]
-        elif api_base:
+            elif message["role"] == "assistant":
-            completion_url = api_base
+                prompt += message["content"] + "</s><s>[INST]"
-        elif "HF_API_BASE" in os.environ:
+            elif message["role"] == "user":
-            completion_url = os.getenv("HF_API_BASE", "")
+                prompt += message["content"] + "[/INST]"
-        else:
+    else:
-            completion_url = f"https://api-inference.huggingface.co/models/{model}"
+        for message in messages:
-        prompt = ""
+            prompt += f"{message['content']}"
-        if (
+    ### MAP INPUT PARAMS
-            "meta-llama" in model and "chat" in model
+    data = {
-        ):  # use the required special tokens for meta-llama - https://huggingface.co/blog/llama2#how-to-prompt-llama-2
+        "inputs": prompt,
-            prompt = "<s>"
+        "parameters": optional_params,
-            for message in messages:
+        "stream": True if "stream" in optional_params and optional_params["stream"] == True else False,
-                if message["role"] == "system":
+    }
-                    prompt += "[INST] <<SYS>>" + message["content"]
+    ## LOGGING
-                elif message["role"] == "assistant":
+    logging_obj.pre_call(
-                    prompt += message["content"] + "</s><s>[INST]"
-                elif message["role"] == "user":
-                    prompt += message["content"] + "[/INST]"
-        else:
-            for message in messages:
-                prompt += f"{message['content']}"
-        ### MAP INPUT PARAMS
-        data = {"inputs": prompt, "parameters": optional_params, "stream": True if "stream" in optional_params and optional_params["stream"] == True else False}
-        ## LOGGING
-        self.logging_obj.pre_call(
             input=prompt,
-            api_key=self.api_key,
+            api_key=api_key,
             additional_args={"complete_input_dict": data},
         )
-        ## COMPLETION CALL
+    ## COMPLETION CALL
-        if "stream" in optional_params and optional_params["stream"] == True:
+    if "stream" in optional_params and optional_params["stream"] == True:
-            response = requests.post(
+        response = requests.post(
-                completion_url, headers=self.headers, data=json.dumps(data), stream=optional_params["stream"]
+            completion_url, 
+            headers=headers, 
+            data=json.dumps(data), 
+            stream=optional_params["stream"]
+        )
+        return response.iter_lines()
+    else:
+        response = requests.post(
+            completion_url, 
+            headers=headers, 
+            data=json.dumps(data)
+        )
+        ## LOGGING
+        logging_obj.post_call(
+            input=prompt,
+            api_key=api_key,
+            original_response=response.text,
+            additional_args={"complete_input_dict": data},
+        )
+        ## RESPONSE OBJECT
+        try:
+            completion_response = response.json()
+        except:
+            raise HuggingfaceError(
+                message=response.text, status_code=response.status_code
+            )
+        print_verbose(f"response: {completion_response}")
+        if isinstance(completion_response, dict) and "error" in completion_response:
+            print_verbose(f"completion error: {completion_response['error']}")
+            print_verbose(f"response.status_code: {response.status_code}")
+            raise HuggingfaceError(
+                message=completion_response["error"],
+                status_code=response.status_code,
             )
-            return response.iter_lines()
         else:
-            response = requests.post(
+            model_response["choices"][0]["message"][
-                completion_url, headers=self.headers, data=json.dumps(data)
+                "content"
-            )
+            ] = completion_response[0]["generated_text"]
-            ## LOGGING
+        
-            self.logging_obj.post_call(
+        ## GETTING LOGPROBS 
-                input=prompt,
+        if "details" in completion_response[0] and "tokens" in completion_response[0]["details"]:
-                api_key=self.api_key,
+            sum_logprob = 0
-                original_response=response.text,
+            for token in completion_response[0]["details"]["tokens"]:
-                additional_args={"complete_input_dict": data},
+                sum_logprob += token["logprob"]
-            )
+            model_response["choices"][0]["message"]["logprobs"] = sum_logprob
-            ## RESPONSE OBJECT
+        ## CALCULATING USAGE
-            try:
+        prompt_tokens = len(
-                completion_response = response.json()
+            encoding.encode(prompt)
-            except:
+        )  ##[TODO] use the llama2 tokenizer here
-                raise HuggingfaceError(message=response.text, status_code=response.status_code)
+        completion_tokens = len(
-            print_verbose(f"response: {completion_response}")
+            encoding.encode(model_response["choices"][0]["message"]["content"])
-            if isinstance(completion_response, dict) and "error" in completion_response:
+        )  ##[TODO] use the llama2 tokenizer here
-                print_verbose(f"completion error: {completion_response['error']}")
-                print_verbose(f"response.status_code: {response.status_code}")
-                raise HuggingfaceError(
-                    message=completion_response["error"],
-                    status_code=response.status_code,
-                )
-            else:
-                model_response["choices"][0]["message"][
-                    "content"
-                ] = completion_response[0]["generated_text"]
-            
-            ## GETTING LOGPROBS 
-            if "details" in completion_response[0] and "tokens" in completion_response[0]["details"]:
-                sum_logprob = 0
-                for token in completion_response[0]["details"]["tokens"]:
-                    sum_logprob += token["logprob"]
-                model_response["choices"][0]["message"]["logprobs"] = sum_logprob
-            ## CALCULATING USAGE
-            prompt_tokens = len(
-                self.encoding.encode(prompt)
-            )  ##[TODO] use the llama2 tokenizer here
-            completion_tokens = len(
-                self.encoding.encode(model_response["choices"][0]["message"]["content"])
-            )  ##[TODO] use the llama2 tokenizer here
 
-            model_response["created"] = time.time()
+        model_response["created"] = time.time()
-            model_response["model"] = model
+        model_response["model"] = model
-            model_response["usage"] = {
+        model_response["usage"] = {
-                "prompt_tokens": prompt_tokens,
+            "prompt_tokens": prompt_tokens,
-                "completion_tokens": completion_tokens,
+            "completion_tokens": completion_tokens,
-                "total_tokens": prompt_tokens + completion_tokens,
+            "total_tokens": prompt_tokens + completion_tokens,
-            }
+        }
-            return model_response
+        return model_response
-        pass
 
-    def embedding(
+def embedding():
-        self,
+    # logic for parsing in - calling - parsing out model embedding calls
-    ):  # logic for parsing in - calling - parsing out model embedding calls
+    pass
-        pass

litellm/main.py

@@ -24,7 +24,7 @@ from .llms import together_ai
 from .llms import ai21
 from .llms import sagemaker
 from .llms import bedrock
-from .llms.huggingface_restapi import HuggingfaceRestAPILLM
+from .llms import huggingface_restapi
 from .llms.baseten import BasetenLLM
 from .llms.aleph_alpha import AlephAlphaLLM
 import tiktoken
@@ -552,10 +552,7 @@ def completion(
                 or os.environ.get("HF_TOKEN")
                 or os.environ.get("HUGGINGFACE_API_KEY")
             )
-            huggingface_client = HuggingfaceRestAPILLM(
+            model_response = huggingface_restapi.completion(
-                encoding=encoding, api_key=huggingface_key, logging_obj=logging
-            )
-            model_response = huggingface_client.completion(
                 model=model,
                 messages=messages,
                 api_base=api_base,
@@ -564,6 +561,10 @@ def completion(
                 optional_params=optional_params,
                 litellm_params=litellm_params,
                 logger_fn=logger_fn,
+                encoding=encoding, 
+                api_key=huggingface_key, 
+                logging_obj=logging
+
             )
             if "stream" in optional_params and optional_params["stream"] == True:
                 # don't try to access stream object,

litellm/tests/test_completion.py

@@ -123,7 +123,10 @@ def test_completion_claude_stream():
 #         # Add any assertions here to check the response
 #         print(response)
 #     except Exception as e:
+#         if "loading" in str(e):
+#             pass
 #         pytest.fail(f"Error occurred: {e}")
+# # test_completion_hf_api()
 
 
 # def test_completion_hf_deployed_api():