Spaces:
Sleeping
Sleeping
File size: 2,696 Bytes
0c3992e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
import os
import os.path as osp
import warnings
import multiprocessing
from functools import partial
from src.tools.api_lib.claude import complete_text_claude
from src.tools.api_lib.gpt import get_gpt_output
from src.tools.api_lib.openai_emb import get_openai_embedding, get_openai_embeddings
from src.tools.api_lib.huggingface import complete_text_hf
# setup parameters for retrying API calls and the sleep time between retries
MAX_OPENAI_RETRY, OPENAI_SLEEP_TIME = 5, 60
MAX_CLAUDE_RETRY, CLAUDE_SLEEP_TIME = 10, 0
registered_text_completion_llms = {
"gpt-4-1106-preview",
"gpt-4-0125-preview", "gpt-4-turbo-preview",
"gpt-4-turbo", "gpt-4-turbo-2024-04-09"
"gpt-4-turbo",
"claude-2.1",
"claude-3-opus-20240229",
"claude-3-sonnet-20240229",
"claude-3-haiku-20240307",
"huggingface/codellama/CodeLlama-7b-hf",
"text-embedding-3-small",
"text-embedding-3-large",
"text-embedding-ada-002"
}
def parallel_func(func, n_max_nodes=5):
'''
A general function to call a function on a list of inputs.
'''
def _parallel_func(inputs: list, **kwargs):
partial_func = partial(func, **kwargs)
processes = min(len(inputs), n_max_nodes)
with multiprocessing.Pool(processes=processes) as pool:
results = pool.map(partial_func, inputs)
return results
return _parallel_func
def get_llm_output(message,
model="gpt-4-0125-preview",
max_tokens=2048,
temperature=1,
json_object=False
):
'''
A general function to complete a prompt using the specified model.
'''
if model not in registered_text_completion_llms:
warnings.warn(f"Model {model} is not registered. You may still be able to use it.")
kwargs = {'message': message,
'model': model,
'max_tokens': max_tokens,
'temperature': temperature,
'json_object': json_object}
if 'gpt-4' in model:
kwargs.update({'max_retry': MAX_OPENAI_RETRY, 'sleep_time': OPENAI_SLEEP_TIME})
return get_gpt_output(**kwargs)
elif 'claude' in model:
kwargs.update({'max_retry': MAX_CLAUDE_RETRY, 'sleep_time': CLAUDE_SLEEP_TIME})
return complete_text_claude(**kwargs)
elif 'huggingface' in model:
return complete_text_hf(**kwargs)
else:
raise ValueError(f"Model {model} not recognized.")
complete_texts_claude = parallel_func(complete_text_claude)
complete_texts_hf = parallel_func(complete_text_hf)
get_gpt_outputs = parallel_func(get_gpt_output)
get_llm_outputs = parallel_func(get_llm_output)
|