diff --git a/litellm/__init__.py b/litellm/__init__.py index aaa36737ef..01559e3d18 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -10,6 +10,8 @@ azure_key = None anthropic_key = None replicate_key = None cohere_key = None + +hugging_api_token = None ####### THREAD-SPECIFIC DATA ################### class MyLocal(threading.local): def __init__(self): diff --git a/litellm/main.py b/litellm/main.py index 299376458f..4fc3bcb05e 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -44,7 +44,8 @@ def completion( temperature=1, top_p=1, n=1, stream=False, stop=None, max_tokens=float('inf'), presence_penalty=0, frequency_penalty=0, logit_bias={}, user="", deployment_id=None, # Optional liteLLM function params - *, return_async=False, api_key=None, force_timeout=60, azure=False, logger_fn=None, verbose=False + *, return_async=False, api_key=None, force_timeout=60, azure=False, logger_fn=None, verbose=False, + hugging_face = False ): try: global new_response @@ -273,6 +274,32 @@ def completion( "total_tokens": prompt_tokens + completion_tokens } response = model_response + elif hugging_face == True: + import requests + API_URL = f"https://api-inference.huggingface.co/models/{model}" + HF_TOKEN = get_secret("HF_TOKEN") + headers = {"Authorization": f"Bearer {HF_TOKEN}"} + + prompt = " ".join([message["content"] for message in messages]) + ## LOGGING + logging(model=model, input=prompt, azure=azure, logger_fn=logger_fn) + input_payload = {"inputs": prompt} + response = requests.post(API_URL, headers=headers, json=input_payload) + + completion_response = response.json()[0]['generated_text'] + ## LOGGING + logging(model=model, input=prompt, azure=azure, additional_args={"max_tokens": max_tokens, "original_response": completion_response}, logger_fn=logger_fn) + prompt_tokens = len(encoding.encode(prompt)) + completion_tokens = len(encoding.encode(completion_response)) + ## RESPONSE OBJECT + model_response["choices"][0]["message"]["content"] = completion_response + model_response["created"] = time.time() + model_response["model"] = model + model_response["usage"] = { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens + } else: ## LOGGING logging(model=model, input=messages, azure=azure, logger_fn=logger_fn) diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 52c9373cc6..e001daa615 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -26,6 +26,16 @@ def test_completion_claude(): except Exception as e: pytest.fail(f"Error occurred: {e}") +def test_completion_hf_api(): + try: + user_message = "write some code to find the sum of two numbers" + messages = [{ "content": user_message,"role": "user"}] + response = completion(model="stabilityai/stablecode-completion-alpha-3b-4k", messages=messages, hugging_face=True) + # Add any assertions here to check the response + print(response) + except Exception as e: + pytest.fail(f"Error occurred: {e}") + def test_completion_cohere(): try: response = completion(model="command-nightly", messages=messages, max_tokens=500)