Loading...
Searching...
No Matches
bindings.cpp File Reference
#include <pybind11/pybind11.h>#include <pybind11/stl.h>#include <pybind11/iostream.h>#include "model_constants.h"#include "model.h"#include "api.h"#include "utf8_text.h"
Include dependency graph for bindings.cpp:
Go to the source code of this file.
Functions | |
| PYBIND11_MODULE (tinyllama_bindings, m) | |
Function Documentation
◆ PYBIND11_MODULE()
| PYBIND11_MODULE | ( | tinyllama_bindings | , |
| m | |||
| ) |
Definition at line 12 of file bindings.cpp.
12 {
14 TinyLlama.cpp Python Bindings
15 =============================
16
17 Python bindings for the tinyllama.cpp inference engine supporting both GGUF and SafeTensors formats.
18
19 Key Classes:
20 - TinyLlamaSession: Main interface for model loading and text generation
21 - ModelConfig: Configuration object containing model parameters and metadata
22
23 Example Usage:
24 import tinyllama_cpp
25
26 # Basic CPU usage
27 session = tinyllama_cpp.TinyLlamaSession(
28 model_path="path/to/model.gguf",
29 tokenizer_path="path/to/tokenizer.json",
30 threads=4
31 )
32
33 # GPU usage with quantized KV cache
34 session = tinyllama_cpp.TinyLlamaSession(
35 model_path="path/to/model.gguf",
36 tokenizer_path="path/to/tokenizer.json",
37 threads=4,
38 n_gpu_layers=-1,
39 use_kv_quant=True
40 )
41
42 # Generate text
43 response = session.generate("What is AI?", steps=64, temperature=0.7)
44 print(response)
45
46 # Batch generation
47 prompts = ["What is AI?", "Explain quantum computing", "Tell me a joke"]
48 responses = session.generate_batch(prompts, steps=64)
49 for prompt, response in zip(prompts, responses):
50 print(f"Q: {prompt}\nA: {response}\n")
51 )pbdoc";
52
54
56 Model configuration containing architecture parameters and metadata.
57
58 This class holds all the configuration parameters for a loaded model,
59 including architecture details, tokenizer information, and model metadata.
60 Most fields are automatically populated when loading a model.
61 )pbdoc");
62
64 .def(py::init<>(), "Create an empty ModelConfig object")
66 .def_readwrite("intermediate_size", &ModelConfig::intermediate_size, "Intermediate size in feed-forward layers")
67 .def_readwrite("num_attention_heads", &ModelConfig::num_attention_heads, "Number of attention heads")
68 .def_readwrite("num_key_value_heads", &ModelConfig::num_key_value_heads, "Number of key-value heads (for GQA)")
69 .def_readwrite("num_hidden_layers", &ModelConfig::num_hidden_layers, "Number of transformer layers")
71 .def_readwrite("max_position_embeddings", &ModelConfig::max_position_embeddings, "Maximum sequence length")
82 .def_readwrite("chat_template_string", &ModelConfig::chat_template_string, "Chat template string")
83 .def_readwrite("is_gguf_file_loaded", &ModelConfig::is_gguf_file_loaded, "Whether model was loaded from GGUF format")
84 .def_readonly("tokenizer_family", &ModelConfig::tokenizer_family, "Tokenizer family (LLAMA_SENTENCEPIECE or LLAMA3_TIKTOKEN)")
85 .def("__repr__",
88 if (cfg.tokenizer_family == ModelConfig::TokenizerFamily::LLAMA_SENTENCEPIECE) tf_str = "LLAMA_SENTENCEPIECE";
89 else if (cfg.tokenizer_family == ModelConfig::TokenizerFamily::LLAMA3_TIKTOKEN) tf_str = "LLAMA3_TIKTOKEN";
93 ">";
94 }
95 );
96
97 py::enum_<ModelConfig::TokenizerFamily>(model_config_class, "TokenizerFamily", "Enumeration of supported tokenizer families")
99 .value("LLAMA_SENTENCEPIECE", ModelConfig::TokenizerFamily::LLAMA_SENTENCEPIECE, "Llama/Llama2 SentencePiece tokenizer")
100 .value("LLAMA3_TIKTOKEN", ModelConfig::TokenizerFamily::LLAMA3_TIKTOKEN, "Llama3 TikToken-style tokenizer")
101 .export_values();
102
104 Main interface for TinyLlama model inference.
105
106 This class provides a high-level interface for loading models and generating text.
107 It supports both GGUF and SafeTensors formats, CPU and GPU inference, and various
108 sampling strategies for text generation.
109
110 The session manages the model, tokenizer, and KV cache, providing both single
111 prompt generation and efficient batch processing capabilities.
112 )pbdoc")
113 .def(py::init([](const py::object& model_path, const py::object& tokenizer_path, int threads, int n_gpu_layers, bool use_mmap, bool use_kv_quant, bool use_batch_generation, int max_batch_size) {
114 // Convert path-like objects to strings
117 return new tinyllama::TinyLlamaSession(model_path_str, tokenizer_path_str, threads, n_gpu_layers, use_mmap, use_kv_quant, use_batch_generation, max_batch_size);
118 }),
119 py::arg("model_path"),
120 py::arg("tokenizer_path"),
121 py::arg("threads") = 1,
122 py::arg("n_gpu_layers") = 0,
123 py::arg("use_mmap") = true,
124 py::arg("use_kv_quant") = false,
125 py::arg("use_batch_generation") = false,
126 py::arg("max_batch_size") = 1,
127 R"pbdoc(
128 Initialize a TinyLlama inference session.
129
130 Args:
131 model_path (str or Path): Path to model directory (SafeTensors) or .gguf file
132 tokenizer_path (str or Path): Path to tokenizer.json file. For GGUF models with
133 embedded tokenizer, this can be the same as model_path
134 threads (int, optional): Number of CPU threads for inference. Defaults to 1.
135 n_gpu_layers (int, optional): Number of layers to offload to GPU.
136 -1 = all layers, 0 = CPU only. Defaults to 0.
137 use_mmap (bool, optional): Use memory mapping for model loading. Defaults to True.
138 use_kv_quant (bool, optional): Use INT8 quantization for KV cache on GPU.
139 Reduces VRAM usage. Defaults to False.
140 use_batch_generation (bool, optional): Enable optimized batch generation mode.
141 Defaults to False.
142 max_batch_size (int, optional): Maximum number of sequences for batch processing.
143 Defaults to 1.
144
145 Raises:
146 RuntimeError: If model loading fails due to invalid paths, unsupported format,
147 or insufficient resources.
148
149 Example:
150 # Basic CPU usage
151 session = TinyLlamaSession("model.gguf", "tokenizer.json", threads=4)
152
153 # Using pathlib.Path objects
154 from pathlib import Path
155 session = TinyLlamaSession(
156 model_path=Path("model.gguf"),
157 tokenizer_path=Path("tokenizer.json"),
158 threads=4,
159 n_gpu_layers=-1,
160 use_kv_quant=True
161 )
162 )pbdoc"
163 )
164 // pybind11 decodes a returned std::string strictly (UnicodeDecodeError on a reply cut inside a
165 // character): the text is valid UTF-8 first, as Python's decode(errors="replace") (utf8_text.h).
167 [](tinyllama::TinyLlamaSession& s, const std::string& prompt, int steps, float temperature, int top_k,
Represents an active TinyLlama session holding the loaded model and tokenizer.
Definition api.h:85
Model configuration structure holding architecture and hyperparameters.
Definition model.h:82
@ LLAMA3_TIKTOKEN
@ LLAMA_SENTENCEPIECE
References ModelConfig::architecture, ModelConfig::bos_token_id, ModelConfig::chat_template_string, ModelConfig::chat_template_type, ModelConfig::eos_token_id, tinyllama::TinyLlamaSession::get_config(), ModelConfig::hidden_act, ModelConfig::hidden_size, ModelConfig::intermediate_size, ModelConfig::is_gguf_file_loaded, j, ModelConfig::LLAMA3_TIKTOKEN, ModelConfig::LLAMA_SENTENCEPIECE, ModelConfig::max_position_embeddings, ModelConfig::model_name, ModelConfig::num_attention_heads, ModelConfig::num_hidden_layers, ModelConfig::num_key_value_heads, out, ModelConfig::pre_tokenizer_type, ModelConfig::rms_norm_eps, ModelConfig::rope_theta, ModelConfig::tokenizer_family, ModelConfig::torch_dtype, ModelConfig::UNKNOWN, ember::utf8_valid(), and ModelConfig::vocab_size.
