New web interface, move llm engine to separate process

tools/gemma_infer/pyproject.toml

@@ -0,0 +1,19 @@
+[build-system]
+requires = ["setuptools>=42", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "gemma_infer"
+version = "0.1.0"
+requires-python = ">=3.8"
+
+dependencies = [
+    "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",
+    "python-dotenv>=1.0.0",
+    "transformers>=4.0.0",
+    "xml_schema_validator @ file:///root/sia/lib/xml_schema_validator",
+]
+
+[project.scripts]
+gemma_infer = "gemma_infer.__main__:main"

tools/gemma_infer/src/gemma_infer/__main__.py

@@ -0,0 +1,42 @@
+from .gemma_llm_engine import GemmaLlmEngine
+from dotenv import load_dotenv
+from llm_engine_utils.protocol import process
+import argparse
+import os
+
+def main():
+    load_dotenv()
+    parser = argparse.ArgumentParser(description='Gemma Inference')
+    
+    parser.add_argument(
+        '--model',
+        type=str,
+        default=os.getenv('SIA_GEMMA_MODEL', '/root/models/current/model.gguf'),
+        help='Model name (default: /root/models/current/model.gguf, env: SIA_GEMMA_MODEL)'
+    )
+
+    parser.add_argument(
+        '--tokenizer',
+        type=str,
+        default=os.getenv('SIA_GEMMA_TOKENIZER', '/root/models/current/tokenizer'),
+        help='Model name (default: /root/models/current/tokenizer, env: SIA_GEMMA_TOKENIZER)'
+    )
+
+    parser.add_argument(
+        '--token-limit',
+        type=int,
+        default=os.getenv('SIA_GEMMA_TOKEN_LIMIT', 4096),
+        help='Token limit (default: 4096, env: SIA_GEMMA_TOKEN_LIMIT)'
+    )
+
+    args = parser.parse_args()
+
+    gemma_llm_engine = GemmaLlmEngine(
+        model=args.model,
+        tokenizer=args.tokenizer,
+        token_limit=args.token_limit,
+    )
+    process(gemma_llm_engine)
+
+if __name__ == "__main__":
+    main()

tools/gemma_infer/src/gemma_infer/gemma_llm_engine.py

@@ -0,0 +1,76 @@
+import os
+os.environ["LLAMA_CPP_LIB_PATH"] = "/usr/local/lib"
+os.environ["LD_LIBRARY_PATH"] += ":/usr/local/lib"
+os.chdir("/usr/local/lib")
+
+from llama_cpp import Llama, LogitsProcessorList
+from llm_engine_utils import LlmEngine
+from pathlib import Path
+from transformers import AutoTokenizer
+from typing import Iterator
+from xml_schema_validator import LlamaCppLogitsProcessor
+
+class GemmaLlmEngine(LlmEngine):
+    def __init__(
+            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_ctx=token_limit,
+            #verbose=False,  # Disable most logging
+        )
+
+    def infer_xml(self, schema: Path, system: str, context: str, prefix: str) -> Iterator[str]:
+        xml_schema_text = Path(schema).read_text()
+        logits_processor = LlamaCppLogitsProcessor(self._tokenizer, xml_schema_text).get_processor()
+        logits_processor_list = LogitsProcessorList([logits_processor])
+        prompt = self._format_messages(system, context, prefix)
+        stream = self._llm.create_completion(
+            prompt=prompt,
+            max_tokens=self._token_limit,
+            stream=True,
+            logits_processor=logits_processor_list
+        )
+        for output in stream:
+            if 'text' in output["choices"][0]:
+                yield output["choices"][0]['text']
+
+    def token_count(self, system: str, context: str) -> int:
+        return len(self._format_messages(system, context, None))
+
+    def token_limit(self) -> int:
+        return self._token_limit
+    
+    def _format_messages(self, system: str, context: str, prefix: str) -> str:
+        messages = [
+            {
+                "role": "user",
+                "content": f"{system}\n\n--- Context ---\n{context}",
+            },
+            {
+                "role": "assistant",
+                "content": prefix,
+                "prefix": True,
+            },
+        ] if prefix else [
+            {
+                "role": "system",
+                "content": system,
+            },
+            {
+                "role": "user",
+                "content": context,
+            },
+        ]
+        return self._tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        ).removeprefix("<bos>")