Spaces:
Running
Running
File size: 5,737 Bytes
3cad23b |
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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 |
import logging
import os
import sys
from typing import Optional
from langchain_community.llms.sambanova import SambaStudio
from langchain_core.language_models.llms import LLM
current_dir = os.path.dirname(os.path.abspath(__file__))
utils_dir = os.path.abspath(os.path.join(current_dir, '..'))
repo_dir = os.path.abspath(os.path.join(utils_dir, '..'))
sys.path.append(utils_dir)
sys.path.append(repo_dir)
from toolformers.sambanova.sambanova_langchain import SambaNovaCloud
EMBEDDING_MODEL = 'intfloat/e5-large-v2'
NORMALIZE_EMBEDDINGS = True
# Configure the logger
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s [%(levelname)s] - %(message)s',
handlers=[
logging.StreamHandler(),
],
)
logger = logging.getLogger(__name__)
class APIGateway:
@staticmethod
def load_llm(
type: str,
streaming: bool = False,
coe: bool = False,
do_sample: Optional[bool] = None,
max_tokens_to_generate: Optional[int] = None,
temperature: Optional[float] = None,
select_expert: Optional[str] = None,
top_p: Optional[float] = None,
top_k: Optional[int] = None,
repetition_penalty: Optional[float] = None,
stop_sequences: Optional[str] = None,
process_prompt: Optional[bool] = False,
sambastudio_base_url: Optional[str] = None,
sambastudio_base_uri: Optional[str] = None,
sambastudio_project_id: Optional[str] = None,
sambastudio_endpoint_id: Optional[str] = None,
sambastudio_api_key: Optional[str] = None,
sambanova_url: Optional[str] = None,
sambanova_api_key: Optional[str] = None,
) -> LLM:
"""Loads a langchain Sambanova llm model given a type and parameters
Args:
type (str): wether to use sambastudio, or SambaNova Cloud model "sncloud"
streaming (bool): wether to use streaming method. Defaults to False.
coe (bool): whether to use coe model. Defaults to False.
do_sample (bool) : Optional wether to do sample.
max_tokens_to_generate (int) : Optional max number of tokens to generate.
temperature (float) : Optional model temperature.
select_expert (str) : Optional expert to use when using CoE models.
top_p (float) : Optional model top_p.
top_k (int) : Optional model top_k.
repetition_penalty (float) : Optional model repetition penalty.
stop_sequences (str) : Optional model stop sequences.
process_prompt (bool) : Optional default to false.
sambastudio_base_url (str): Optional SambaStudio environment URL".
sambastudio_base_uri (str): Optional SambaStudio-base-URI".
sambastudio_project_id (str): Optional SambaStudio project ID.
sambastudio_endpoint_id (str): Optional SambaStudio endpoint ID.
sambastudio_api_token (str): Optional SambaStudio endpoint API key.
sambanova_url (str): Optional SambaNova Cloud URL",
sambanova_api_key (str): Optional SambaNovaCloud API key.
Returns:
langchain llm model
"""
if type == 'sambastudio':
envs = {
'sambastudio_base_url': sambastudio_base_url,
'sambastudio_base_uri': sambastudio_base_uri,
'sambastudio_project_id': sambastudio_project_id,
'sambastudio_endpoint_id': sambastudio_endpoint_id,
'sambastudio_api_key': sambastudio_api_key,
}
envs = {k: v for k, v in envs.items() if v is not None}
if coe:
model_kwargs = {
'do_sample': do_sample,
'max_tokens_to_generate': max_tokens_to_generate,
'temperature': temperature,
'select_expert': select_expert,
'top_p': top_p,
'top_k': top_k,
'repetition_penalty': repetition_penalty,
'stop_sequences': stop_sequences,
'process_prompt': process_prompt,
}
model_kwargs = {k: v for k, v in model_kwargs.items() if v is not None}
llm = SambaStudio(
**envs,
streaming=streaming,
model_kwargs=model_kwargs,
)
else:
model_kwargs = {
'do_sample': do_sample,
'max_tokens_to_generate': max_tokens_to_generate,
'temperature': temperature,
'top_p': top_p,
'top_k': top_k,
'repetition_penalty': repetition_penalty,
'stop_sequences': stop_sequences,
}
model_kwargs = {k: v for k, v in model_kwargs.items() if v is not None}
llm = SambaStudio(
**envs,
streaming=streaming,
model_kwargs=model_kwargs,
)
elif type == 'sncloud':
envs = {
'sambanova_url': sambanova_url,
'sambanova_api_key': sambanova_api_key,
}
envs = {k: v for k, v in envs.items() if v is not None}
llm = SambaNovaCloud(
**envs,
max_tokens=max_tokens_to_generate,
model=select_expert,
temperature=temperature,
top_k=top_k,
top_p=top_p,
)
else:
raise ValueError(f"Invalid LLM API: {type}, only 'sncloud' and 'sambastudio' are supported.")
return llm |