feat: Expose gguf model metadata in metadata property

This commit is contained in:
Andrei Betlen 2024-01-19 10:46:03 -05:00
parent 833a7f1a86
commit 5a34c57e54
2 changed files with 35 additions and 0 deletions

View file

@ -204,6 +204,31 @@ class _LlamaModel:
output[1:] if len(tokens) > 0 and tokens[0] == self.token_bos() else output
)
# Extra
def metadata(self) -> Dict[str, str]:
assert self.model is not None
metadata: Dict[str, str] = {}
buffer_size = 1024
buffer = ctypes.create_string_buffer(buffer_size)
# zero the buffer
buffer.value = b'\0' * buffer_size
# iterate over model keys
for i in range(llama_cpp.llama_model_meta_count(self.model)):
nbytes = llama_cpp.llama_model_meta_key_by_index(self.model, i, buffer, buffer_size)
if nbytes > buffer_size:
buffer_size = nbytes
buffer = ctypes.create_string_buffer(buffer_size)
nbytes = llama_cpp.llama_model_meta_key_by_index(self.model, i, buffer, buffer_size)
key = buffer.value.decode("utf-8")
nbytes = llama_cpp.llama_model_meta_val_str_by_index(self.model, i, buffer, buffer_size)
if nbytes > buffer_size:
buffer_size = nbytes
buffer = ctypes.create_string_buffer(buffer_size)
nbytes = llama_cpp.llama_model_meta_val_str_by_index(self.model, i, buffer, buffer_size)
value = buffer.value.decode("utf-8")
metadata[key] = value
return metadata
@staticmethod
def default_params():
"""Get the default llama_model_params."""

View file

@ -331,6 +331,16 @@ class Llama:
self._mirostat_mu = ctypes.c_float(2.0 * 5.0) # TODO: Move this to sampling context
try:
self.metadata = self._model.metadata()
except Exception as e:
self.metadata = {}
if self.verbose:
print(f"Failed to load metadata: {e}", file=sys.stderr)
if self.verbose:
print(f"Model metadata: {self.metadata}", file=sys.stderr)
@property
def ctx(self) -> llama_cpp.llama_context_p:
assert self._ctx.ctx is not None