Renamed mistral tool to mistral_api

tools/mistral_local_infer/pyproject.toml

@@ -0,0 +1,25 @@
+[build-system]
+requires = ["setuptools>=42", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "mistral_local_infer"
+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@v0.3.16#egg=llama-cpp-python&env=CMAKE_ARGS=-DLLAMA_BUILD=OFF",
+    "llm_engine_utils @ file:///root/sia/lib/llm_engine_utils",
+    "mistral-common>=1.8.6",
+    "protobuf>=6.0.0",
+    "python-dotenv>=1.0.0",
+    "sentencepiece>=0.2.0",
+    "tiktoken>=0.9.0",
+    "transformers>=5.0.0rc0",
+    "vulkan",
+    "xml_schema_validator @ file:///root/sia/lib/xml_schema_validator",
+]
+
+[project.scripts]
+mistral_local_infer = "mistral_local_infer.__main__:main"

tools/mistral_local_infer/src/mistral_local_infer/__main__.py

@@ -0,0 +1,42 @@
+from .mistral_llm_engine import MistralLlmEngine
+from dotenv import load_dotenv
+from llm_engine_utils.protocol import process
+import argparse
+import os
+
+def main():
+    load_dotenv()
+    parser = argparse.ArgumentParser(description='Ministral-3 Local Inference using llama.cpp with Vulkan')
+
+    parser.add_argument(
+        '--model',
+        type=str,
+        default=os.getenv('SIA_MISTRAL_MODEL', '/root/models/current/model.gguf'),
+        help='Model name (default: /root/models/current/model.gguf, env: SIA_MISTRAL_MODEL)'
+    )
+
+    parser.add_argument(
+        '--tokenizer',
+        type=str,
+        default=os.getenv('SIA_MISTRAL_TOKENIZER', '/root/models/current/tokenizer'),
+        help='Model name (default: /root/models/current/tokenizer, env: SIA_MISTRAL_TOKENIZER)'
+    )
+
+    parser.add_argument(
+        '--token-limit',
+        type=int,
+        default=os.getenv('SIA_MISTRAL_TOKEN_LIMIT', 10000),
+        help='Token limit (default: 10000, env: SIA_MISTRAL_TOKEN_LIMIT)'
+    )
+
+    args = parser.parse_args()
+
+    mistral_llm_engine = MistralLlmEngine(
+        model=args.model,
+        tokenizer=args.tokenizer,
+        token_limit=args.token_limit,
+    )
+    process(mistral_llm_engine)
+
+if __name__ == "__main__":
+    main()

tools/mistral_local_infer/src/mistral_local_infer/mistral_llm_engine.py

@@ -0,0 +1,97 @@
+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, llama_cpp
+from llm_engine_utils import LlmEngine
+from llm_engine_utils.iterators import skip_prefix
+from pathlib import Path
+from transformers import AutoTokenizer
+from typing import Iterator
+from xml_schema_validator import LlamaCppLogitsProcessor
+
+llama_cpp._lib.ggml_backend_load_all()
+
+class MistralLlmEngine(LlmEngine):
+    """
+    Ministral-3 inference engine using llama.cpp with Vulkan acceleration.
+
+    Ministral-3 architecture features:
+    - 34 transformer layers with alternating attention (1 full + 3 sliding window)
+    - 131K vocabulary tokens
+    - Supports up to 256K context window
+    - Uses Grouped Query Attention (32 heads, 8 key-value heads)
+    """
+    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=0,
+            n_ctx=token_limit,
+            flash_attn=True,
+        )
+
+    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
+        )
+
+        def content_generator():
+            for output in stream:
+                choice = output["choices"][0]
+                if choice.get('finish_reason'):
+                    break
+                if 'text' in choice:
+                    text = choice['text']
+                    if text == '':
+                        break
+                    yield text
+
+        yield from skip_prefix(content_generator(), prefix)
+
+    def token_count(self, system: str, context: str) -> int:
+        prompt = self._format_messages(system, context, None)
+        tokens = self._tokenizer.encode(prompt)
+        return len(tokens)
+
+    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,
+            },
+        ] if prefix else [
+            {
+                "role": "user",
+                "content": f"{system}\n\n--- Context ---\n{context}",
+            }
+        ]
+        return self._tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=not prefix,
+            continue_final_message=bool(prefix)
+        ).removeprefix("<bos>")