jhansss's picture
add more llms; add system prompt support; add device=auto
780954b
raw
history blame contribute delete
413 Bytes
from abc import ABC, abstractmethod
class AbstractLLMModel(ABC):
def __init__(
self, model_id: str, device: str = "auto", cache_dir: str = "cache", **kwargs
):
print(f"Loading LLM model {model_id}...")
self.model_id = model_id
self.device = device
self.cache_dir = cache_dir
@abstractmethod
def generate(self, prompt: str, **kwargs) -> str:
pass