add requiredMem method to llmodel impls

most of these can just shortcut out of the model loading logic llama is a bit worse to deal with because we submodule it so I have to at least parse the hparams, and then I just use the size on disk as an estimate for the mem size (which seems reasonable since we mmap() the llama files anyway)
This commit is contained in:
Aaron Miller
2023-06-26 12:17:34 -07:00
committed by AT
parent dead954134
commit b19a3e5b2c
14 changed files with 154 additions and 8 deletions

View File

@@ -81,6 +81,8 @@ llmodel.llmodel_model_destroy.restype = None
llmodel.llmodel_loadModel.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
llmodel.llmodel_loadModel.restype = ctypes.c_bool
llmodel.llmodel_required_mem.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
llmodel.llmodel_required_mem.restype = ctypes.c_size_t
llmodel.llmodel_isModelLoaded.argtypes = [ctypes.c_void_p]
llmodel.llmodel_isModelLoaded.restype = ctypes.c_bool
@@ -131,6 +133,16 @@ class LLModel:
if self.model is not None:
llmodel.llmodel_model_destroy(self.model)
def memory_needed(self, model_path: str) -> int:
model_path_enc = model_path.encode("utf-8")
self.model = llmodel.llmodel_model_create(model_path_enc)
if self.model is not None:
return llmodel.llmodel_required_mem(self.model, model_path_enc)
else:
raise ValueError("Unable to instantiate model")
def load_model(self, model_path: str) -> bool:
"""
Load model from a file.