File size: 3,040 Bytes
c531eac 932fded c531eac 932fded c531eac 932fded c531eac 932fded c531eac 932fded c531eac 932fded c531eac |
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 |
import os
from typing import Any, Callable
from smolagents import HfApiModel, InferenceClientModel, LiteLLMModel, OpenAIServerModel
def get_huggingface_api_model(model_id: str, **kwargs) -> HfApiModel:
"""
Returns a Hugging Face API model instance.
Args:
model_id (str): The model identifier.
**kwargs: Additional keyword arguments for the model.
Returns:
HfApiModel: Hugging Face API model instance.
"""
api_key = os.getenv("HUGGINGFACEHUB_API_TOKEN")
if not api_key:
raise ValueError("HUGGINGFACEHUB_API_TOKEN is not set")
return HfApiModel(model_id=model_id, token=api_key, **kwargs)
def get_inference_client_model(model_id: str, **kwargs) -> InferenceClientModel:
"""
Returns an Inference Client model instance.
Args:
model_id (str): The model identifier.
**kwargs: Additional keyword arguments for the model.
Returns:
InferenceClientModel: Inference client model instance.
"""
api_key = os.getenv("HUGGINGFACEHUB_API_TOKEN")
if not api_key:
raise ValueError("HUGGINGFACEHUB_API_TOKEN is not set")
return InferenceClientModel(model_id=model_id, token=api_key, **kwargs)
def get_openai_server_model(model_id: str, **kwargs) -> OpenAIServerModel:
"""
Returns an OpenAI server model instance.
Args:
model_id (str): The model identifier.
**kwargs: Additional keyword arguments for the model.
Returns:
OpenAIServerModel: OpenAI server model instance.
"""
api_key = os.getenv("OPENAI_API_KEY")
if not api_key:
raise ValueError("OPENAI_API_KEY is not set")
api_base = os.getenv("OPENAI_API_BASE")
if not api_base:
raise ValueError("OPENAI_API_BASE is not set")
return OpenAIServerModel(
model_id=model_id, api_key=api_key, api_base=api_base, **kwargs
)
def get_lite_llm_model(model_id: str, **kwargs) -> LiteLLMModel:
"""
Returns a LiteLLM model instance.
Args:
model_id (str): The model identifier.
**kwargs: Additional keyword arguments for the model.
Returns:
LiteLLMModel: LiteLLM model instance.
"""
return LiteLLMModel(model_id=model_id, **kwargs)
def get_model(model_type: str, model_id: str, **kwargs) -> Any:
"""
Returns a model instance based on the specified type.
Args:
model_type (str): The type of the model (e.g., 'HfApiModel').
model_id (str): The model identifier.
**kwargs: Additional keyword arguments for the model.
Returns:
Any: Model instance of the specified type.
"""
models: dict[str, Callable[..., Any]] = {
"HfApiModel": get_huggingface_api_model,
"InferenceClientModel": get_inference_client_model,
"OpenAIServerModel": get_openai_server_model,
"LiteLLMModel": get_lite_llm_model,
}
if model_type not in models:
raise ValueError(f"Unknown model type: {model_type}")
return models[model_type](model_id, **kwargs)
|