-
-
Notifications
You must be signed in to change notification settings - Fork 4
/
llm_anyscale_endpoints.py
36 lines (30 loc) · 984 Bytes
/
llm_anyscale_endpoints.py
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
import llm
from llm.default_plugins.openai_models import Chat
MODELS = (
"meta-llama/Llama-2-7b-chat-hf",
"meta-llama/Llama-2-13b-chat-hf",
"meta-llama/Llama-2-70b-chat-hf",
"codellama/CodeLlama-34b-Instruct-hf",
"mistralai/Mistral-7B-Instruct-v0.1",
"mistralai/Mixtral-8x7B-Instruct-v0.1",
"Open-Orca/Mistral-7B-OpenOrca",
"HuggingFaceH4/zephyr-7b-beta",
)
class AnyscaleEndpointChat(Chat):
needs_key = "anyscale-endpoints"
def __str__(self):
return "AnyscaleEndpoints: {}".format(self.model_id)
@llm.hookimpl
def register_models(register):
# Only do this if the key is set
key = llm.get_key("", "anyscale-endpoints", "LLM_ANYSCALE_ENDPOINTS_KEY")
if not key:
return
for model_id in MODELS:
register(
AnyscaleEndpointChat(
model_id=model_id,
model_name=model_id,
api_base="https://api.endpoints.anyscale.com/v1",
)
)