SingingSDS / modules /llm /registry.py
jhansss's picture
add more llms; add system prompt support; add device=auto
780954b
raw
history blame contribute delete
588 Bytes
from .base import AbstractLLMModel
LLM_MODEL_REGISTRY = {}
def register_llm_model(prefix: str):
def wrapper(cls):
assert issubclass(cls, AbstractLLMModel), f"{cls} must inherit AbstractLLMModel"
LLM_MODEL_REGISTRY[prefix] = cls
return cls
return wrapper
def get_llm_model(model_id: str, device="auto", **kwargs) -> AbstractLLMModel:
for prefix, cls in LLM_MODEL_REGISTRY.items():
if model_id.startswith(prefix):
return cls(model_id, device=device, **kwargs)
raise ValueError(f"No LLM wrapper found for model: {model_id}")