Suppress llama.cpp output when loading model.

This commit is contained in:
Andrei Betlen 2023-07-28 14:45:18 -04:00
parent a9b9f0397c
commit ce57920e60
2 changed files with 57 additions and 4 deletions

View file

@ -27,6 +27,8 @@ from .llama_types import *
import numpy as np import numpy as np
import numpy.typing as npt import numpy.typing as npt
from .utils import suppress_stdout_stderr
class BaseLlamaCache(ABC): class BaseLlamaCache(ABC):
"""Base cache class for a llama.cpp model.""" """Base cache class for a llama.cpp model."""
@ -308,12 +310,25 @@ class Llama:
if not os.path.exists(model_path): if not os.path.exists(model_path):
raise ValueError(f"Model path does not exist: {model_path}") raise ValueError(f"Model path does not exist: {model_path}")
self.model = llama_cpp.llama_load_model_from_file( if verbose:
self.model_path.encode("utf-8"), self.params self.model = llama_cpp.llama_load_model_from_file(
) self.model_path.encode("utf-8"), self.params
)
else:
with suppress_stdout_stderr():
self.model = llama_cpp.llama_load_model_from_file(
self.model_path.encode("utf-8"), self.params
)
assert self.model is not None assert self.model is not None
self.ctx = llama_cpp.llama_new_context_with_model(self.model, self.params) if verbose:
self.ctx = llama_cpp.llama_new_context_with_model(self.model, self.params)
else:
with suppress_stdout_stderr():
print("here")
self.ctx = llama_cpp.llama_new_context_with_model(
self.model, self.params
)
assert self.ctx is not None assert self.ctx is not None

38
llama_cpp/utils.py Normal file
View file

@ -0,0 +1,38 @@
import os
import sys
class suppress_stdout_stderr(object):
# Oddly enough this works better than the contextlib version
def __enter__(self):
self.outnull_file = open(os.devnull, "w")
self.errnull_file = open(os.devnull, "w")
self.old_stdout_fileno_undup = sys.stdout.fileno()
self.old_stderr_fileno_undup = sys.stderr.fileno()
self.old_stdout_fileno = os.dup(sys.stdout.fileno())
self.old_stderr_fileno = os.dup(sys.stderr.fileno())
self.old_stdout = sys.stdout
self.old_stderr = sys.stderr
os.dup2(self.outnull_file.fileno(), self.old_stdout_fileno_undup)
os.dup2(self.errnull_file.fileno(), self.old_stderr_fileno_undup)
sys.stdout = self.outnull_file
sys.stderr = self.errnull_file
return self
def __exit__(self, *_):
sys.stdout = self.old_stdout
sys.stderr = self.old_stderr
os.dup2(self.old_stdout_fileno, self.old_stdout_fileno_undup)
os.dup2(self.old_stderr_fileno, self.old_stderr_fileno_undup)
os.close(self.old_stdout_fileno)
os.close(self.old_stderr_fileno)
self.outnull_file.close()
self.errnull_file.close()