Use Vulkan

tools/gemma_infer/pyproject.toml

@@ -16,6 +16,7 @@ dependencies = [
     "sentencepiece>=0.2.0",
     "tiktoken>=0.9.0",
     "transformers>=4.0.0",
+    "vulkan",
     "xml_schema_validator @ file:///root/sia/lib/xml_schema_validator",
 ]
 

tools/gemma_infer/src/gemma_infer/gemma_llm_engine.py

@@ -15,19 +15,20 @@ llama_cpp._lib.ggml_backend_load_all()
 
 class GemmaLlmEngine(LlmEngine):
     def __init__(
-            self,
-            model: str,
-            tokenizer: str,
-            token_limit: int,
-        ):
+        self,
+        model: str,
+        tokenizer: str,
+        token_limit: int,
+    ):
         self._model = model
         self._tokenizer = AutoTokenizer.from_pretrained(tokenizer)
         self._token_limit = token_limit
+
         self._llm = Llama(
             model_path=model,
-            n_gpu_layers=100,
+            n_gpu_layers=0,
             n_ctx=token_limit,
-            #verbose=False,  # Disable most logging
+            flash_attn=True,
         )
 
     def infer_xml(self, schema: Path, system: str, context: str, prefix: str) -> Iterator[str]:
@@ -56,7 +57,9 @@ class GemmaLlmEngine(LlmEngine):
         yield from skip_prefix(content_generator(), prefix)
 
     def token_count(self, system: str, context: str) -> int:
-        return len(self._format_messages(system, context, None))
+        prompt = self._format_messages(system, context, None)
+        tokens = self._tokenizer.encode(prompt)
+        return len(tokens)
 
     def token_limit(self) -> int:
         return self._token_limit