Replace janky LLaMA implementation with llama-cpp-python

Author: kharvd

gpt.py

@@ -169,10 +169,8 @@ def main():
     if config.google_api_key:
         genai.configure(api_key=config.google_api_key)
 
-    if config.llama_config is not None:
-        init_llama_models(
-            config.llama_config["llama_cpp_dir"], config.llama_config["models"]
-        )
+    if config.llama_models is not None:
+        init_llama_models(config.llama_models)
 
     assistant = init_assistant(cast(AssistantGlobalArgs, args), config.assistants)
 

gptcli/config.py

@@ -11,11 +11,6 @@
 ]
 
 
-class LLaMAConfig(TypedDict):
-    llama_cpp_dir: str
-    models: Dict[str, str]  # name -> path
-
-
 @dataclass
 class GptCliConfig:
     default_assistant: str = "general"
@@ -29,7 +24,7 @@ class GptCliConfig:
     log_level: str = "INFO"
     assistants: Dict[str, AssistantConfig] = {}
     interactive: Optional[bool] = None
-    llama_config: Optional[LLaMAConfig] = None
+    llama_models: Optional[Dict[str, str]] = None
 
 
 def choose_config_file(paths: List[str]) -> str:

gptcli/llama.py

@@ -1,18 +1,15 @@
-import logging
 import os
-import signal
-import subprocess
 from pathlib import Path
 import sys
 from typing import Iterator, List, Optional
+from llama_cpp import Llama
 
 from gptcli.completion import CompletionProvider, Message
 
-LLAMA_DIR: Optional[Path] = None
-LLAMA_MODELS: Optional[dict[str, Path]] = None
+LLAMA_MODELS: Optional[dict[str, str]] = None
 
 
-def init_llama_models(llama_cpp_dir: str, model_paths: dict[str, str]):
+def init_llama_models(model_paths: dict[str, str]):
     for name, path in model_paths.items():
         if not os.path.isfile(path):
             print(f"LLaMA model {name} not found at {path}.")
@@ -21,9 +18,8 @@ def init_llama_models(llama_cpp_dir: str, model_paths: dict[str, str]):
             print(f"LLaMA model names must start with `llama`, but got `{name}`.")
             sys.exit(1)
 
-    global LLAMA_DIR, LLAMA_MODELS
-    LLAMA_DIR = Path(llama_cpp_dir)
-    LLAMA_MODELS = {name: Path(path) for name, path in model_paths.items()}
+    global LLAMA_MODELS
+    LLAMA_MODELS = model_paths
 
 
 def role_to_name(role: str) -> str:
@@ -50,75 +46,51 @@ class LLaMACompletionProvider(CompletionProvider):
     def complete(
         self, messages: List[Message], args: dict, stream: bool = False
     ) -> Iterator[str]:
-        assert LLAMA_DIR, "LLaMA models not initialized"
         assert LLAMA_MODELS, "LLaMA models not initialized"
 
+        with suppress_stderr():
+            llm = Llama(
+                model_path=LLAMA_MODELS[args["model"]],
+                n_ctx=2048,
+                verbose=False,
+                use_mlock=True,
+            )
         prompt = make_prompt(messages)
 
-        extra_args = []
+        extra_args = {}
         if "temperature" in args:
-            extra_args += ["--temp", str(args["temperature"])]
+            extra_args["temperature"] = args["temperature"]
         if "top_p" in args:
-            extra_args += ["--top_p", str(args["top_p"])]
-
-        process = subprocess.Popen(
-            [
-                LLAMA_DIR / "main",
-                "--model",
-                LLAMA_MODELS[args["model"]],
-                "-n",
-                "4096",
-                "-r",
-                "### Human:",
-                "-p",
-                prompt,
-                *extra_args,
-            ],
-            stdin=subprocess.PIPE,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            shell=False,
-            text=True,
+            extra_args["top_p"] = args["top_p"]
+
+        gen = llm.create_completion(
+            prompt,
+            max_tokens=1024,
+            stop=END_SEQ,
+            stream=stream,
+            echo=False,
+            **extra_args,
         )
-
         if stream:
-            return self._read_stream(process, prompt)
+            for x in gen:
+                yield x["choices"][0]["text"]
         else:
-            return self._read(process, prompt)
-
-    def _read_stream(self, process: subprocess.Popen, prompt: str) -> Iterator[str]:
-        assert process.stdout, "LLaMA stdout not set"
-        assert process.stderr, "LLaMA stderr not set"
-
-        buffer = ""
-        num_read = 0
-        char = process.stdout.read(1)
-
-        try:
-            while char := process.stdout.read(1):
-                num_read += len(char)
-                if num_read <= len(prompt):
-                    continue
-
-                buffer += char
-                if not buffer.startswith("#") or (buffer != END_SEQ[: len(buffer)]):
-                    yield buffer
-                    buffer = ""
-                elif buffer.endswith(END_SEQ):
-                    yield buffer[: -len(END_SEQ)]
-                    buffer = ""
-                    process.terminate()
-                    break
-        except KeyboardInterrupt:
-            os.kill(process.pid, signal.SIGINT)
-            raise
-        finally:
-            process.wait()
-            stderr = "".join(process.stderr.readlines())
-            logging.debug(f"LLaMA stderr: {stderr}")
-
-    def _read(self, process: subprocess.Popen, prompt: str) -> Iterator[str]:
-        result = ""
-        for token in self._read_stream(process, prompt):
-            result += token
-        yield result
+            yield gen["choices"][0]["text"]
+
+
+# https://stackoverflow.com/a/50438156
+class suppress_stderr(object):
+    def __enter__(self):
+        self.errnull_file = open(os.devnull, "w")
+        self.old_stderr_fileno_undup = sys.stderr.fileno()
+        self.old_stderr_fileno = os.dup(sys.stderr.fileno())
+        self.old_stderr = sys.stderr
+        os.dup2(self.errnull_file.fileno(), self.old_stderr_fileno_undup)
+        sys.stderr = self.errnull_file
+        return self
+
+    def __exit__(self, *_):
+        sys.stderr = self.old_stderr
+        os.dup2(self.old_stderr_fileno, self.old_stderr_fileno_undup)
+        os.close(self.old_stderr_fileno)
+        self.errnull_file.close()

requirements.txt

@@ -1,6 +1,7 @@
 anthropic==0.2.8
 black==23.1.0
 google-generativeai==0.1.0rc2
+llama-cpp-python==0.1.57
 openai==0.27.2
 prompt-toolkit==3.0.38
 pytest==7.3.1