mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-26 03:04:13 +00:00
108 lines
3 KiB
Python
108 lines
3 KiB
Python
import os
|
|
import json
|
|
from enum import Enum
|
|
import requests
|
|
import time
|
|
from typing import Callable
|
|
from litellm.utils import ModelResponse
|
|
|
|
class PetalsError(Exception):
|
|
def __init__(self, status_code, message):
|
|
self.status_code = status_code
|
|
self.message = message
|
|
super().__init__(
|
|
self.message
|
|
) # Call the base class constructor with the parameters it needs
|
|
|
|
PetalsConfig = {
|
|
"max_new_tokens": 256
|
|
}
|
|
|
|
def completion(
|
|
model: str,
|
|
messages: list,
|
|
model_response: ModelResponse,
|
|
print_verbose: Callable,
|
|
encoding,
|
|
logging_obj,
|
|
optional_params=None,
|
|
stream=False,
|
|
litellm_params=None,
|
|
logger_fn=None,
|
|
):
|
|
try:
|
|
import torch
|
|
from transformers import AutoTokenizer
|
|
from petals import AutoDistributedModelForCausalLM
|
|
except:
|
|
raise Exception(
|
|
"Importing torch, transformers, petals failed\nTry pip installing petals \npip install git+https://github.com/bigscience-workshop/petals"
|
|
)
|
|
|
|
model = model
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model, use_fast=False, add_bos_token=False)
|
|
model_obj = AutoDistributedModelForCausalLM.from_pretrained(model)
|
|
|
|
prompt = ""
|
|
for message in messages:
|
|
if "role" in message:
|
|
if message["role"] == "user":
|
|
prompt += (
|
|
f"{message['content']}"
|
|
)
|
|
else:
|
|
prompt += (
|
|
f"{message['content']}"
|
|
)
|
|
else:
|
|
prompt += f"{message['content']}"
|
|
|
|
## Load Config
|
|
for k, v in PetalsConfig.items():
|
|
if k not in optional_params:
|
|
optional_params[k] = v
|
|
|
|
## LOGGING
|
|
logging_obj.pre_call(
|
|
input=prompt,
|
|
api_key="",
|
|
additional_args={"complete_input_dict": optional_params},
|
|
)
|
|
|
|
## COMPLETION CALL
|
|
inputs = tokenizer(prompt, return_tensors="pt")["input_ids"]
|
|
|
|
# optional params: max_new_tokens=1,temperature=0.9, top_p=0.6
|
|
outputs = model_obj.generate(inputs, **optional_params)
|
|
|
|
## LOGGING
|
|
logging_obj.post_call(
|
|
input=prompt,
|
|
api_key="",
|
|
original_response=outputs,
|
|
additional_args={"complete_input_dict": optional_params},
|
|
)
|
|
## RESPONSE OBJECT
|
|
output_text = tokenizer.decode(outputs[0])
|
|
model_response["choices"][0]["message"]["content"] = output_text
|
|
|
|
prompt_tokens = len(
|
|
encoding.encode(prompt)
|
|
)
|
|
completion_tokens = len(
|
|
encoding.encode(model_response["choices"][0]["message"]["content"])
|
|
)
|
|
|
|
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,
|
|
}
|
|
return model_response
|
|
|
|
def embedding():
|
|
# logic for parsing in - calling - parsing out model embedding calls
|
|
pass
|