From c22917e3c9481bde020f3062a79da26f02cd70c3 Mon Sep 17 00:00:00 2001
From: Niels Geens <niels.geens@gmail.com>
Date: Thu, 22 May 2025 18:51:07 +0000
Subject: [PATCH] Fixes for 12B

---
 scripts/install.sh                            | 16 ++++++++--------
 tools/gemma_infer/pyproject.toml              |  4 ++++
 tools/gemma_train/pyproject.toml              | 12 ++++++------
 tools/gemma_train/src/gemma_train/__main__.py |  8 +++++---
 4 files changed, 23 insertions(+), 17 deletions(-)

diff --git a/scripts/install.sh b/scripts/install.sh
index 229a7db..88b45e2 100755
--- a/scripts/install.sh
+++ b/scripts/install.sh
@@ -1,13 +1,13 @@
 #!/bin/bash
 
-#SIA_INSTALL_NO_WEB=1
-#SIA_INSTALL_NO_LLAMA_CPP=1
-#SIA_INSTALL_NO_NOTEBOOK=1
-#SIA_INSTALL_NO_CORE=1
-#SIA_INSTALL_NO_ITB=1
-#SIA_INSTALL_NO_MISTRAL_INFER=1
-#SIA_INSTALL_NO_GEMMA_INFER=1
-#SIA_INSTALL_NO_GEMMA_TRAIN=1
+SIA_INSTALL_NO_WEB=1
+SIA_INSTALL_NO_LLAMA_CPP=
+SIA_INSTALL_NO_NOTEBOOK=1
+SIA_INSTALL_NO_CORE=1
+SIA_INSTALL_NO_ITB=1
+SIA_INSTALL_NO_MISTRAL_INFER=1
+SIA_INSTALL_NO_GEMMA_INFER=
+SIA_INSTALL_NO_GEMMA_TRAIN=
 
 cd "/root/desktop"
 
diff --git a/tools/gemma_infer/pyproject.toml b/tools/gemma_infer/pyproject.toml
index 6754d7b..54ea708 100644
--- a/tools/gemma_infer/pyproject.toml
+++ b/tools/gemma_infer/pyproject.toml
@@ -8,9 +8,13 @@ version = "0.1.0"
 requires-python = ">=3.8"
 
 dependencies = [
+    "blobfile>=3.0.0",
     "llama-cpp-python @ git+https://github.com/abetlen/llama-cpp-python.git#egg=llama-cpp-python&env=CMAKE_ARGS=-DLLAMA_BUILD=OFF",
     "llm_engine_utils @ file:///root/sia/lib/llm_engine_utils",
+    "protobuf>=6.0.0",
     "python-dotenv>=1.0.0",
+    "sentencepiece>=0.2.0",
+    "tiktoken>=0.9.0",
     "transformers>=4.0.0",
     "xml_schema_validator @ file:///root/sia/lib/xml_schema_validator",
 ]
diff --git a/tools/gemma_train/pyproject.toml b/tools/gemma_train/pyproject.toml
index b4101c4..c8d7f25 100644
--- a/tools/gemma_train/pyproject.toml
+++ b/tools/gemma_train/pyproject.toml
@@ -8,17 +8,17 @@ version = "0.1.0"
 requires-python = ">=3.8"
 
 dependencies = [
-    "accelerate==1.4.0",
-    "bitsandbytes==0.45.3",
+    "accelerate==1.2.0",
+    "bitsandbytes>=0.45.0",
     "datasets==3.3.2",
     "evaluate==0.4.3",
     "llm_engine_utils @ file:///root/sia/lib/llm_engine_utils",
-    "peft==0.14.0",
+    "peft==0.13.2",
     "python-dotenv>=1.0.0",
     "sentencepiece>=0.2.0",
-    "torch>=2.4.0",
-    "transformers>=4.51.3",
-    "trl==0.15.2",
+    "torch==2.4.1",
+    "transformers==4.51.3",
+    "trl>=0.17.0",
 ]
 
 [project.scripts]
diff --git a/tools/gemma_train/src/gemma_train/__main__.py b/tools/gemma_train/src/gemma_train/__main__.py
index 201453a..170d9af 100644
--- a/tools/gemma_train/src/gemma_train/__main__.py
+++ b/tools/gemma_train/src/gemma_train/__main__.py
@@ -19,6 +19,7 @@ def train(config: Config):
         config.model,
         token=config.api_key,
         trust_remote_code=True,
+        use_fast=False,
     )
     tokenizer.save_pretrained(config.output_dir/"tokenizer")
 
@@ -81,12 +82,13 @@ def train(config: Config):
 def merge(config: Config):
     adapted_model = AutoPeftModelForCausalLM.from_pretrained(
         config.output_dir/"lora_adapter",
-        torch_dtype=torch.float16,  # Use float16 for better compatibility
-        device_map="auto",
-        offload_folder="offload",
+        torch_dtype=torch.float16,
         token=config.api_key,
+        load_in_4bit=False,
+        load_in_8bit=False,
     )
     merged_model = adapted_model.merge_and_unload()
+
     merged_model.save_pretrained(
         config.output_dir/"merged",
         safe_serialization=True