Initial commit

llama_cpp/__init__.py

@@ -0,0 +1,2 @@
+from .llama_cpp import *
+from .llama import *

llama_cpp/llama.py

@@ -0,0 +1,131 @@
+import uuid
+import time
+import multiprocessing
+from typing import List, Optional
+
+from . import llama_cpp
+
+class Llama:
+    def __init__(
+        self,
+        model_path: str,
+        n_ctx: int = 512,
+        n_parts: int = -1,
+        seed: int = 1337,
+        f16_kv: bool = False,
+        logits_all: bool = False,
+        vocab_only: bool = False,
+        n_threads: Optional[int] = None,
+        model_name: Optional[str]=None,
+    ):
+        self.model_path = model_path
+        self.model = model_name or model_path
+
+        self.params = llama_cpp.llama_context_default_params()
+        self.params.n_ctx = n_ctx
+        self.params.n_parts = n_parts
+        self.params.seed = seed
+        self.params.f16_kv = f16_kv
+        self.params.logits_all = logits_all
+        self.params.vocab_only = vocab_only
+
+        self.n_threads = n_threads or multiprocessing.cpu_count()
+
+        self.tokens = (llama_cpp.llama_token * self.params.n_ctx)()
+
+        self.ctx = llama_cpp.llama_init_from_file(
+            self.model_path.encode("utf-8"), self.params
+        )
+
+    def __call__(
+        self,
+        prompt: str,
+        suffix: Optional[str] = None,
+        max_tokens: int = 16,
+        temperature: float = 0.8,
+        top_p: float = 0.95,
+        echo: bool = False,
+        stop: List[str] = [],
+        repeat_penalty: float = 1.1,
+        top_k: int = 40,
+    ):
+        text = ""
+        finish_reason = "length"
+        completion_tokens = 0
+
+        prompt_tokens = llama_cpp.llama_tokenize(
+            self.ctx, prompt.encode("utf-8"), self.tokens, self.params.n_ctx, True
+        )
+
+        if prompt_tokens + max_tokens > self.params.n_ctx:
+            raise ValueError(
+                f"Requested tokens exceed context window of {self.params.n_ctx}"
+            )
+
+        for i in range(prompt_tokens):
+            llama_cpp.llama_eval(
+                self.ctx, (llama_cpp.c_int * 1)(self.tokens[i]), 1, i, self.n_threads
+            )
+
+        for i in range(max_tokens):
+            token = llama_cpp.llama_sample_top_p_top_k(
+                self.ctx,
+                self.tokens,
+                prompt_tokens + completion_tokens,
+                top_k=top_k,
+                top_p=top_p,
+                temp=temperature,
+                repeat_penalty=repeat_penalty,
+            )
+            if token == llama_cpp.llama_token_eos():
+                finish_reason = "stop"
+                break
+            text += llama_cpp.llama_token_to_str(self.ctx, token).decode("utf-8")
+            self.tokens[prompt_tokens + i] = token
+            completion_tokens += 1
+
+            any_stop = [s for s in stop if s in text]
+            if len(any_stop) > 0:
+                first_stop = any_stop[0]
+                text = text[: text.index(first_stop)]
+                finish_reason = "stop"
+                break
+
+            llama_cpp.llama_eval(
+                self.ctx,
+                (llama_cpp.c_int * 1)(self.tokens[prompt_tokens + i]),
+                1,
+                prompt_tokens + completion_tokens,
+                self.n_threads,
+            )
+
+        if echo:
+            text = prompt + text
+
+        if suffix is not None:
+            text = text + suffix
+
+        return {
+            "id": f"cmpl-{str(uuid.uuid4())}", # Likely to change
+            "object": "text_completion",
+            "created": int(time.time()),
+            "model": self.model, # Likely to change
+            "choices": [
+                {
+                    "text": text,
+                    "index": 0,
+                    "logprobs": None,
+                    "finish_reason": finish_reason,
+                }
+            ],
+            "usage": {
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": prompt_tokens + completion_tokens,
+            },
+        }
+
+    def __del__(self):
+        llama_cpp.llama_free(self.ctx)
+
+

llama_cpp/llama_cpp.py

@@ -0,0 +1,157 @@
+import ctypes
+
+from ctypes import c_int, c_float, c_double, c_char_p, c_void_p, c_bool, POINTER, Structure
+
+import pathlib
+
+# Load the library
+libfile = pathlib.Path(__file__).parent.parent / "libllama.so"
+lib = ctypes.CDLL(str(libfile))
+
+
+# C types
+llama_token = c_int
+llama_token_p = POINTER(llama_token)
+
+class llama_token_data(Structure):
+    _fields_ = [
+        ('id', llama_token), # token id
+        ('p', c_float), # probability of the token
+        ('plog', c_float), # log probability of the token
+    ]
+
+llama_token_data_p = POINTER(llama_token_data)
+
+class llama_context_params(Structure):
+    _fields_ = [
+        ('n_ctx', c_int), # text context
+        ('n_parts', c_int), # -1 for default
+        ('seed', c_int), # RNG seed, 0 for random
+        ('f16_kv', c_bool), # use fp16 for KV cache
+        ('logits_all', c_bool), # the llama_eval() call computes all logits, not just the last one
+
+        ('vocab_only', c_bool), # only load the vocabulary, no weights
+    ]
+
+llama_context_params_p = POINTER(llama_context_params)
+
+llama_context_p = c_void_p
+
+# C functions
+lib.llama_context_default_params.argtypes = []
+lib.llama_context_default_params.restype = llama_context_params
+
+lib.llama_init_from_file.argtypes = [c_char_p, llama_context_params]
+lib.llama_init_from_file.restype = llama_context_p
+
+lib.llama_free.argtypes = [llama_context_p]
+lib.llama_free.restype = None
+
+lib.llama_model_quantize.argtypes = [c_char_p, c_char_p, c_int, c_int]
+lib.llama_model_quantize.restype = c_int
+
+lib.llama_eval.argtypes = [llama_context_p, llama_token_p, c_int, c_int, c_int]
+lib.llama_eval.restype = c_int
+
+lib.llama_tokenize.argtypes = [llama_context_p, c_char_p, llama_token_p, c_int, c_bool]
+lib.llama_tokenize.restype = c_int
+
+lib.llama_n_vocab.argtypes = [llama_context_p]
+lib.llama_n_vocab.restype = c_int
+
+lib.llama_n_ctx.argtypes = [llama_context_p]
+lib.llama_n_ctx.restype = c_int
+
+lib.llama_get_logits.argtypes = [llama_context_p]
+lib.llama_get_logits.restype = POINTER(c_float)
+
+lib.llama_token_to_str.argtypes = [llama_context_p, llama_token]
+lib.llama_token_to_str.restype = c_char_p
+
+lib.llama_token_bos.argtypes = []
+lib.llama_token_bos.restype = llama_token
+
+lib.llama_token_eos.argtypes = []
+lib.llama_token_eos.restype = llama_token
+
+lib.llama_sample_top_p_top_k.argtypes = [llama_context_p, llama_token_p, c_int, c_int, c_double, c_double, c_double]
+lib.llama_sample_top_p_top_k.restype = llama_token
+
+lib.llama_print_timings.argtypes = [llama_context_p]
+lib.llama_print_timings.restype = None
+
+lib.llama_reset_timings.argtypes = [llama_context_p]
+lib.llama_reset_timings.restype = None
+
+lib.llama_print_system_info.argtypes = []
+lib.llama_print_system_info.restype = c_char_p
+
+# Python functions
+def llama_context_default_params() -> llama_context_params:
+    params = lib.llama_context_default_params()
+    return params
+
+def llama_init_from_file(path_model: bytes, params: llama_context_params) -> llama_context_p:
+    """Various functions for loading a ggml llama model.
+    Allocate (almost) all memory needed for the model.
+    Return NULL on failure """
+    return lib.llama_init_from_file(path_model, params)
+
+def llama_free(ctx: llama_context_p):
+    """Free all allocated memory"""
+    lib.llama_free(ctx)
+
+def llama_model_quantize(fname_inp: bytes, fname_out: bytes, itype: c_int, qk: c_int) -> c_int:
+    """Returns 0 on success"""
+    return lib.llama_model_quantize(fname_inp, fname_out, itype, qk)
+
+def llama_eval(ctx: llama_context_p, tokens: llama_token_p, n_tokens: c_int, n_past: c_int, n_threads: c_int) -> c_int:
+    """Run the llama inference to obtain the logits and probabilities for the next token.
+    tokens + n_tokens is the provided batch of new tokens to process
+    n_past is the number of tokens to use from previous eval calls
+    Returns 0 on success"""
+    return lib.llama_eval(ctx, tokens, n_tokens, n_past, n_threads)
+
+def llama_tokenize(ctx: llama_context_p, text: bytes, tokens: llama_token_p, n_max_tokens: c_int, add_bos: c_bool) -> c_int:
+    """Convert the provided text into tokens.
+    The tokens pointer must be large enough to hold the resulting tokens.
+    Returns the number of tokens on success, no more than n_max_tokens
+    Returns a negative number on failure - the number of tokens that would have been returned"""
+    return lib.llama_tokenize(ctx, text, tokens, n_max_tokens, add_bos)
+
+def llama_n_vocab(ctx: llama_context_p) -> c_int:
+    return lib.llama_n_vocab(ctx)
+
+def llama_n_ctx(ctx: llama_context_p) -> c_int:
+    return lib.llama_n_ctx(ctx)
+
+def llama_get_logits(ctx: llama_context_p):
+    """Token logits obtained from the last call to llama_eval()
+    The logits for the last token are stored in the last row
+    Can be mutated in order to change the probabilities of the next token
+    Rows: n_tokens
+    Cols: n_vocab"""
+    return lib.llama_get_logits(ctx)
+
+def llama_token_to_str(ctx: llama_context_p, token: int) -> bytes:
+    """Token Id -> String. Uses the vocabulary in the provided context"""
+    return lib.llama_token_to_str(ctx, token)
+
+def llama_token_bos() -> llama_token:
+    return lib.llama_token_bos()
+
+def llama_token_eos() -> llama_token:
+    return lib.llama_token_eos()
+
+def llama_sample_top_p_top_k(ctx: llama_context_p, last_n_tokens_data: llama_token_p, last_n_tokens_size: c_int, top_k: c_int, top_p: c_double, temp: c_double, repeat_penalty: c_double) -> llama_token:
+    return lib.llama_sample_top_p_top_k(ctx, last_n_tokens_data, last_n_tokens_size, top_k, top_p, temp, repeat_penalty)
+
+def llama_print_timings(ctx: llama_context_p):
+    lib.llama_print_timings(ctx)
+
+def llama_reset_timings(ctx: llama_context_p):
+    lib.llama_reset_timings(ctx)
+
+def llama_print_system_info() -> bytes:
+    """Print system informaiton"""
+    return lib.llama_print_system_info()