llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Gemma inference with llama.cpp and logits processor

Browse Source
This commit is contained in:
geens
2025-05-06 20:31:42 +02:00
parent 44fd60a6be
commit c41a86a00a
6 changed files with 238 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
Dockerfile
View File

@@ -1,18 +1,24 @@
FROM nvidia/cuda:12.2.0-runtime-ubuntu22.04 AS base
# Install base packages
RUN apt-get update && \
apt-get upgrade -y && \
apt install -y \
build-essential \
cmake \
cuda-toolkit \
curl \
git \
gnupg \
jq \
libcurl4-nss-dev \
python3-dev \
python3-venv \
tmux \
vim \
wget
# Install chrome
RUN wget -q -O - https://dl-ssl.google.com/linux/linux_signing_key.pub | gpg --dearmor -o /etc/apt/trusted.gpg.d/google.gpg
RUN echo "deb http://dl.google.com/linux/chrome/deb/ stable main" >> /etc/apt/sources.list.d/google-chrome.list
RUN apt-get update && \
@@ -21,9 +27,16 @@ RUN apt-get update && \
RUN rm -rf /var/lib/apt/lists/*
# Install Rust
RUN curl https://sh.rustup.rs -sSf | bash -s -- -y
ENV PATH="/root/.cargo/bin:${PATH}"
# Install llama.cpp
RUN curl -O https://git.nielsgeens.be/api/packages/llm/generic/llama.cpp/b5269/llama.cpp.tar
RUN tar -xf /llama.cpp.tar -C /usr/local/bin --wildcards --no-anchored "llama-*"
RUN tar -xf /llama.cpp.tar -C /usr/local/lib --wildcards --no-anchored "*.so"
RUN rm llama.cpp.tar
# Create directory structure
RUN mkdir -p \
/root/sia \

13
lib/xml_schema_validator/src/py_xml_logits_processor_core.rs
View File

@@ -9,6 +9,7 @@ pub struct XmlLogitsProcessorCore {
xml_schema_validator: XmlSchemaValidator,
trie: char_trie::CharTrie,
tokens: Vec<(u32, String)>,
invalid_tokens: Vec<u32>,
}
#[pymethods]
@@ -27,13 +28,20 @@ impl XmlLogitsProcessorCore {
})
.collect();
let mut trie = char_trie::CharTrie::default();
let mut invalid_tokens = Vec::new();
for (token_id, token_text) in tokens.iter() {
let valid = token_text.chars().all(|c| c.is_ascii_graphic() || c.is_ascii_whitespace());
if valid {
trie.insert(token_text, token_id.clone());
} else {
invalid_tokens.push(*token_id);
}
}
Ok(Self {
xml_schema_validator,
trie,
tokens,
invalid_tokens,
})
}
@@ -47,7 +55,10 @@ impl XmlLogitsProcessorCore {
fn get_invalid_tokens(&mut self, py: Python<'_>) -> PyResult<Vec<u32>> {
Ok(self
.trie
.find_invalid_tokens(&self.xml_schema_validator))
.find_invalid_tokens(&self.xml_schema_validator)
.into_iter()
.chain(self.invalid_tokens.iter().cloned())
.collect())
}
/// Check if the validator has reached the end of the XML

3
lib/xml_schema_validator/src/token/mod.rs
View File

@@ -49,9 +49,6 @@ pub enum Token {
impl Token {
pub fn append(self, c: char) -> Vec<Token> {
if !c.is_ascii_graphic() && ! c.is_ascii_whitespace() {
return vec![];
}
match self {
Token::AttributeEquals(attribute_equals) => attribute_equals.append(c),
Token::AttributeName(attribute_name) => attribute_name.append(c),

10
scripts/bootstrap.sh
View File

@@ -27,8 +27,11 @@ fi
apt-get update
apt-get install -y \
build-essential \
cmake \
cuda-toolkit \
git \
gnupg \
libcurl4-nss-dev \
python3-dev \
python3-venv \
vim \
@@ -59,6 +62,13 @@ if [ -z "${SIA_INSTALL_NO_CORE}" ]; then
nvm install node
fi
# Install llama.cpp
cd /
curl -O https://git.nielsgeens.be/api/packages/llm/generic/llama.cpp/b5269/llama.cpp.tar
tar -xf /llama.cpp.tar -C /usr/local/bin --wildcards --no-anchored "llama-*"
tar -xf /llama.cpp.tar -C /usr/local/lib --wildcards --no-anchored "*.so"
rm llama.cpp.tar
# Create directory structure
echo "Creating directory structure..."
mkdir -p "/root/data/iterations"

1
scripts/local.sh
View File

@@ -27,6 +27,7 @@ docker run \
-ti \
--gpus=all \
-p 8080:8080 \
-p 8000:8000 \
--env-file .env \
-v /$(pwd)/models/:/root/models/ \
-v /$(pwd)/iterations/:/root/data/iterations/ \

231
tools/notebook/gemma.ipynb
View File

@@ -26,7 +26,7 @@
"bash\n",
"cd /\n",
". /root/venvs/notebook/bin/activate\n",
"jupyter notebook --ip 0.0.0.0 --port 8080 --allow-root\n",
"jupyter notebook --ip 0.0.0.0 --port 8080 --allow-root --NotebookApp.token='' --NotebookApp.password=''\n",
"\"\"\""
]
},
@@ -147,7 +147,7 @@
"import gc\n",
"import os\n",
"import torch\n",
"import transformers\n"
"import transformers"
]
},
{
@@ -170,9 +170,9 @@
"outputs": [],
"source": [
"# Set model ID\n",
"#model_id = \"google/gemma-3-1b-it\"\n",
"model_id = \"google/gemma-3-12b-it\"\n",
"output_dir = \"/root/models/notebook\""
"model_id = \"google/gemma-3-1b-it\"\n",
"#model_id = \"google/gemma-3-12b-it\"\n",
"output_dir = \"/root/models/gemma3_1b\""
]
},
{
@@ -198,7 +198,8 @@
"metadata": {},
"outputs": [],
"source": [
"# Load tokenizer\n",
"from transformers import AutoTokenizer\n",
"import os\n",
"tokenizer = AutoTokenizer.from_pretrained(model_id, token=os.environ['SIA_HF_API_KEY'])"
]
},
@@ -330,7 +331,7 @@
"outputs": [],
"source": [
"# Save tokenizer first\n",
"tokenizer.save_pretrained(output_dir+\"_merged\")"
"tokenizer.save_pretrained(output_dir+\"_tokenizer\")"
]
},
{
@@ -344,6 +345,14 @@
"trainer.model.save_pretrained(output_dir+\"_lora_adapter\")"
]
},
{
"cell_type": "markdown",
"id": "c3129410",
"metadata": {},
"source": [
"# Merge lora adapters"
]
},
{
"cell_type": "code",
"execution_count": null,
@@ -367,7 +376,7 @@
"source": [
"# Load the LoRA model with offloading to manage memory\n",
"adapted_model = AutoPeftModelForCausalLM.from_pretrained(\n",
" \"/root/models/notebook_lora_adapter\",\n",
" output_dir+\"_lora_adapter\",\n",
" torch_dtype=torch.float16, # Use float16 for better compatibility\n",
" device_map=\"auto\",\n",
" offload_folder=\"offload\",\n",
@@ -399,6 +408,14 @@
")"
]
},
{
"cell_type": "markdown",
"id": "2fdfdfba",
"metadata": {},
"source": [
"# Inference using transformers"
]
},
{
"cell_type": "code",
"execution_count": null,
@@ -410,7 +427,7 @@
"from threading import Thread\n",
"\n",
"messages = [\n",
" {\"role\": \"user\", \"content\": \"Hi, write a poem about iced-tea.\"},\n",
" {\"role\": \"user\", \"content\": \"Hi\"},\n",
" {\"role\": \"assistant\", \"content\": \"\"},\n",
"]\n",
"prompt = tokenizer.apply_chat_template(\n",
@@ -446,13 +463,11 @@
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e5cc8064",
"cell_type": "markdown",
"id": "e28c6fd3",
"metadata": {},
"outputs": [],
"source": [
"exit(1)"
"# Convert to gguf"
]
},
{
@@ -475,6 +490,39 @@
"%pip install -r llama.cpp/requirements.txt"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Convert to GGUF\n",
"!python ./llama.cpp/convert_hf_to_gguf.py \\\n",
" --outfile /root/models/gemma3_1b.gguf \\\n",
" --outtype q8_0 \\\n",
" /root/models/gemma3_1b_merged"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e229f70b",
"metadata": {},
"outputs": [],
"source": [
"%%bash\n",
"llama-cli -m /root/models/gemma3_1b.gguf -p \"hi\""
]
},
{
"cell_type": "markdown",
"id": "6291761b",
"metadata": {},
"source": [
"# Llama.cpp python inference\n",
"Streaming and with Logits processor"
]
},
{
"cell_type": "code",
"execution_count": null,
@@ -482,68 +530,191 @@
"metadata": {},
"outputs": [],
"source": [
"%pip install llama-cpp-python"
"%%bash\n",
"export LLAMA_CPP_LIB=/usr/local/lib/libllama.so\n",
"export CMAKE_ARGS=\"-DLLAMA_BUILD=OFF\"\n",
"#export CMAKE_ARGS=\"-DLLAMA_BUILD=ON -DGGML_CUDA=ON\"\n",
"pip install git+https://github.com/abetlen/llama-cpp-python.git --force-reinstall --upgrade --no-cache-dir\n",
"#pip install llama-cpp-python"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c6532f46",
"metadata": {},
"outputs": [],
"source": [
"# Convert to GGUF (if merge was successful)\n",
"!python ./llama.cpp/convert_hf_to_gguf.py \\\n",
" --outfile /root/models/notebook_12B.gguf \\\n",
" --outtype q8_0 \\\n",
" /root/models/notebook_merged"
"%pip install pydantic_settings torch transformers"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6fea3615",
"id": "a6ede6a2",
"metadata": {},
"outputs": [],
"source": [
"!apt update && apt install -y unzip"
"%pip install /root/sia/lib/xml_schema_validator/"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d2418ba6",
"id": "2faf282e",
"metadata": {},
"outputs": [],
"source": [
"!wget https://github.com/ggml-org/llama.cpp/releases/download/b5226/llama-b5226-bin-ubuntu-x64.zip"
"import os\n",
"os.environ[\"LLAMA_CPP_LIB_PATH\"] = \"/usr/local/lib\"\n",
"os.environ[\"LD_LIBRARY_PATH\"] += \":/usr/local/lib\""
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "262270a6",
"id": "b950bffa",
"metadata": {},
"outputs": [],
"source": [
"!unzip -o llama-b5226-bin-ubuntu-x64.zip"
"from pathlib import Path\n",
"from transformers import AutoTokenizer\n",
"\n",
"# Load the schema and tokenizer\n",
"xml_schema_text = Path(\"/root/sia/action_schema.xsd\").read_text()\n",
"tokenizer = AutoTokenizer.from_pretrained(\"/root/models/gemma3_1b_tokenizer\")\n",
"#tokenizer = AutoTokenizer.from_pretrained(\"google/gemma-3-1b-it\")\n",
"\n",
"# Create chat messages\n",
"messages = [\n",
" {\n",
" \"role\": \"user\",\n",
" \"content\": \"Reply with exactly the full quoted text. Nothing more, nothing less. '<reasoning>Lorem ipsum dolor sit amet, consectetur adipiscing elit.\\nAenean vel orci bibendum, ullamcorper sem id, sollicitudin nisl.\\nPhasellus tristique bibendum mauris, at aliquet libero auctor vitae.\\nPraesent vestibulum mi id auctor faucibus.\\nSed quam justo, tincidunt eu tincidunt non, mollis sit amet nisi.\\nPhasellus non sapien volutpat neque varius posuere eu non ligula.\\nIn blandit, dui vel interdum tristique, risus nisi malesuada nibh, sed auctor lorem turpis id quam.\\nSed semper tortor quis lacinia pulvinar.\\nSed ultricies lacus a nunc tincidunt vestibulum.\\nIn venenatis, metus id iaculis tempus, metus magna fringilla enim, ac vestibulum arcu elit tempor massa.\\nQuisque feugiat mollis lacinia. Nullam eu lacinia justo. </reasoning>'\"\n",
" }\n",
"]\n",
"\n",
"# Create a prompt using the transformers tokenizer\n",
"prompt = tokenizer.apply_chat_template(\n",
" messages,\n",
" tokenize=False,\n",
" add_generation_prompt=True\n",
").removeprefix(\"<bos>\")\n",
"\n",
"print(f\"Prompt: {prompt}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ff438dcc",
"id": "abf09265",
"metadata": {},
"outputs": [],
"source": [
"!cd build/bin/ && ./llama-cli -m /root/models/notebook.gguf"
"from xml_schema_validator import TransformersLogitsProcessor\n",
"import torch\n",
"\n",
"# Create the logits processor for XML validation\n",
"logits_processor = TransformersLogitsProcessor(tokenizer, schema_text=xml_schema_text)\n",
"\n",
"\n",
"# Create a lambda wrapper for the XML logits processor\n",
"llama_processor = lambda tokens, logits: logits_processor(\n",
" torch.tensor(tokens, dtype=torch.long).unsqueeze(0),\n",
" torch.tensor(logits, dtype=torch.float).unsqueeze(0)\n",
")[0].tolist()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3e484bcf",
"metadata": {},
"outputs": [],
"source": [
"from xml_schema_validator import LlamaCppLogitsProcessor\n",
"\n",
"llama_processor = LlamaCppLogitsProcessor(tokenizer, xml_schema_text)\n",
"llama_processor = llama_processor.get_processor()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ab9363e6",
"metadata": {},
"outputs": [],
"source": [
"\n",
"# Added debug info\n",
"def llama_processor(tokens, logits):\n",
" token_tensor = torch.tensor(tokens, dtype=torch.long).unsqueeze(0)\n",
" logits_tensor = torch.tensor(logits, dtype=torch.float).unsqueeze(0)\n",
" \n",
" # Print token count and non-negative infinity logits count before processing\n",
" valid_count_before = sum(1 for l in logits if l > float('-inf'))\n",
" print(f\"Tokens: {len(tokens)}\")\n",
" print(f\"Valid logits before: {valid_count_before}\")\n",
" \n",
" # Process logits\n",
" processed = logits_processor(token_tensor, logits_tensor)[0].tolist()\n",
" \n",
" # Print valid logits after processing\n",
" valid_count_after = sum(1 for l in processed if l > float('-inf'))\n",
" print(f\"Valid logits after: {valid_count_after}\")\n",
" \n",
" return processed"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b8187bf6",
"metadata": {},
"outputs": [],
"source": [
"from llama_cpp import Llama\n",
"\n",
"# Load the model\n",
"llm = Llama(\n",
" model_path=\"/root/models/gemma3_1b.gguf\",\n",
" n_gpu_layers=100,\n",
" #verbose=False, # Disable most logging\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9951ec9b",
"metadata": {},
"outputs": [],
"source": [
"from llama_cpp import LogitsProcessorList\n",
"\n",
"# Create the LogitsProcessorList for llama-cpp\n",
"processors_list = LogitsProcessorList([llama_processor])\n",
"\n",
"# Generate with streaming\n",
"stream = llm.create_completion(\n",
" prompt=prompt,\n",
" max_tokens=500,\n",
" stream=True,\n",
" logits_processor=processors_list\n",
")\n",
"\n",
"# Print output as it's generated\n",
"print(\"Output\")\n",
"for output in stream:\n",
" if 'text' in output[\"choices\"][0]:\n",
" print(output[\"choices\"][0]['text'], end=\"\")"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "notebook",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "notebook"
"name": "python3"
},
"language_info": {
"codemirror_mode": {