Renamed mistral tool to mistral_api

scripts/install.sh

@@ -5,7 +5,10 @@
 #SIA_INSTALL_NO_NOTEBOOK=1
 #SIA_INSTALL_NO_CORE=1
 #SIA_INSTALL_NO_ITB=1
-#SIA_INSTALL_NO_MISTRAL_INFER=1
+#SIA_INSTALL_NO_MISTRAL_API_INFER=1
+#SIA_INSTALL_NO_MISTRAL_API_TRAIN=1
+#SIA_INSTALL_NO_MISTRAL_LOCAL_INFER=1
+#SIA_INSTALL_NO_MISTRAL_LOCAL_TRAIN=1
 #SIA_INSTALL_NO_GEMMA_INFER=1
 #SIA_INSTALL_NO_GEMMA_TRAIN=1
 
@@ -46,10 +49,28 @@ if [ -z "${SIA_INSTALL_NO_ITB}" ]; then
     /root/venvs/itb/bin/pip install -e /root/sia/tools/itb
 fi
 
-if [ -z "${SIA_INSTALL_NO_MISTRAL_INFER}" ]; then
+if [ -z "${SIA_INSTALL_NO_MISTRAL_API_INFER}" ]; then
-    echo "Installing venv for mistral inference"
+    echo "Installing venv for mistral API inference"
-    python3 -m venv /root/venvs/mistral_infer
+    python3 -m venv /root/venvs/mistral_api_infer
-    /root/venvs/mistral_infer/bin/pip install -e /root/sia/tools/mistral_infer
+    /root/venvs/mistral_api_infer/bin/pip install -e /root/sia/tools/mistral_api_infer
+fi
+
+if [ -z "${SIA_INSTALL_NO_MISTRAL_API_TRAIN}" ]; then
+    echo "Installing venv for mistral API training"
+    python3 -m venv /root/venvs/mistral_api_train
+    /root/venvs/mistral_api_train/bin/pip install -e /root/sia/tools/mistral_api_train
+fi
+
+if [ -z "${SIA_INSTALL_NO_MISTRAL_LOCAL_INFER}" ]; then
+    echo "Installing venv for mistral local inference"
+    python3 -m venv /root/venvs/mistral_local_infer
+    /root/venvs/mistral_local_infer/bin/pip install -e /root/sia/tools/mistral_local_infer
+fi
+
+if [ -z "${SIA_INSTALL_NO_MISTRAL_LOCAL_TRAIN}" ]; then
+    echo "Installing venv for mistral local training"
+    python3 -m venv /root/venvs/mistral_local_train
+    /root/venvs/mistral_local_train/bin/pip install -e /root/sia/tools/mistral_local_train
 fi
 
 if [ -z "${SIA_INSTALL_NO_GEMMA_INFER}" ]; then

tools/mistral_api_infer/pyproject.toml

@@ -3,7 +3,7 @@ requires = ["setuptools>=42", "wheel"]
 build-backend = "setuptools.build_meta"
 
 [project]
-name = "mistral_infer"
+name = "mistral_api_infer"
 version = "0.1.0"
 requires-python = ">=3.8"
 
@@ -15,4 +15,4 @@ dependencies = [
 ]
 
 [project.scripts]
-mistral_infer = "mistral_infer.__main__:main"
+mistral_api_infer = "mistral_api_infer.__main__:main"

tools/mistral_api_infer/src/mistral_api_infer/__main__.py

@@ -1,4 +1,4 @@
-from .mistral_llm_engine import MistralLlmEngine
+from .mistral_llm_engine import MistralApiLlmEngine
 from dotenv import load_dotenv
 from llm_engine_utils.protocol import process
 import argparse

tools/mistral_api_train/pyproject.toml

@@ -3,7 +3,7 @@ requires = ["setuptools>=42", "wheel"]
 build-backend = "setuptools.build_meta"
 
 [project]
-name = "mistral_train"
+name = "mistral_api_train"
 version = "0.1.0"
 requires-python = ">=3.8"
 
@@ -16,4 +16,4 @@ dependencies = [
 ]
 
 [project.scripts]
-mistral_train = "mistral_train.__main__:main"
+mistral_api_train = "mistral_api_train.__main__:main"

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>")

tools/mistral_local_train/pyproject.toml

@@ -0,0 +1,27 @@
+[build-system]
+requires = ["setuptools>=42", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "mistral_local_train"
+version = "0.1.0"
+requires-python = ">=3.8"
+
+dependencies = [
+    "accelerate==1.2.0",
+    "bitsandbytes>=0.45.0",
+    "datasets==3.3.2",
+    "evaluate==0.4.3",
+    "kernels>=0.11.1",
+    "llm_engine_utils @ file:///root/sia/lib/llm_engine_utils",
+    "mistral-common>=1.8.6",
+    "peft==0.13.2",
+    "python-dotenv>=1.0.0",
+    "sentencepiece>=0.2.0",
+    "torch==2.4.1",
+    "transformers>=5.0.0rc0",
+    "trl>=0.17.0",
+]
+
+[project.scripts]
+mistral_local_train = "mistral_local_train.__main__:main"

tools/mistral_local_train/src/mistral_local_train/__main__.py

@@ -0,0 +1,106 @@
+from llm_engine_utils.dataset import Dataset
+from pathlib import Path
+from peft import LoraConfig, AutoPeftModelForCausalLM
+from transformers import AutoTokenizer, TrainingArguments
+from transformers.models.ministral3 import Ministral3ForCausalLM
+from trl import SFTTrainer
+import os
+import sys
+import torch
+
+from .config import Config
+
+def main():
+    config = Config()
+    train(config)
+    merge(config)
+
+def train(config: Config):
+    tokenizer = AutoTokenizer.from_pretrained(
+        config.model,
+        token=config.api_key,
+        trust_remote_code=True,
+        use_fast=False,
+    )
+    tokenizer.save_pretrained(config.output_dir/"tokenizer")
+
+    # Use Ministral3ForCausalLM for text-only training
+    # Note: torch_dtype is deprecated in transformers 5.x, use dtype instead
+    model = Ministral3ForCausalLM.from_pretrained(
+        config.model,
+        dtype=torch.float32,
+        device_map="cpu",
+        token=config.api_key,
+        attn_implementation='eager',
+        low_cpu_mem_usage=True,
+    )
+
+    dataset = Dataset(config.config_path)
+    dataset.validate()
+    dataset = dataset.to_transformers_dataset(tokenizer)
+
+    # LoRA config targeting Ministral-3 architecture layers
+    # Ministral-3 uses alternating attention (1 full + 3 sliding window) across 34 layers
+    # Target modules include:
+    #   - Attention projections: q_proj, k_proj, v_proj, o_proj
+    #   - MLP projections: gate_proj, up_proj, down_proj
+    #   - Language model head: lm_head
+    lora_config = LoraConfig(
+        r=8,
+        lora_alpha=16,
+        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "lm_head"],
+        lora_dropout=0.05,
+        bias="none",
+        task_type="CAUSAL_LM",
+    )
+
+    training_args = TrainingArguments(
+        per_device_train_batch_size=1,
+        gradient_accumulation_steps=4,
+        warmup_steps=1,
+        num_train_epochs=2,
+        learning_rate=1e-3,
+        fp16=False,
+        logging_steps=5,
+        save_strategy="steps",
+        save_steps=10,
+        output_dir=config.output_dir/"lora",
+        optim="adamw_torch",
+        seed=42,
+        group_by_length=True,
+        use_cpu=True,
+    )
+
+    trainer = SFTTrainer(
+        model=model,
+        train_dataset=dataset,
+        args=training_args,
+        peft_config=lora_config,
+        formatting_func=format_sia_example,
+    )
+
+    trainer.train()
+
+    trainer.model.save_pretrained(config.output_dir/"lora_adapter")
+
+def merge(config: Config):
+    # Note: torch_dtype is deprecated in transformers 5.x, use dtype instead
+    adapted_model = AutoPeftModelForCausalLM.from_pretrained(
+        config.output_dir/"lora_adapter",
+        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
+    )
+
+def format_sia_example(example):
+    return example['messages'].removeprefix("<bos>")
+
+if __name__ == "__main__":
+    exit(main())

tools/mistral_local_train/src/mistral_local_train/config.py

@@ -0,0 +1,50 @@
+from pathlib import Path
+import argparse
+import os
+
+class Config:
+    def __init__(self):
+        parser = argparse.ArgumentParser(
+            description='Train Ministral-3 model locally using LoRA fine-tuning'
+        )
+        parser.add_argument(
+            '--config',
+            type=Path,
+            default=Path('/root/sia/training/config.yaml'),
+            help='Path to config file (default: /root/sia/training/config.yaml)'
+        )
+        parser.add_argument(
+            '--model',
+            type=str,
+            default='mistralai/Ministral-3-8B-Base-2512',
+            help='Base model for fine-tuning (default: mistralai/Ministral-3-8B-Base-2512)'
+        )
+        parser.add_argument(
+            '--api-key',
+            type=str,
+            default=os.environ.get('SIA_HF_API_KEY'),
+            help='Huggingface API key (optional, env: SIA_HF_API_KEY)'
+        )
+        parser.add_argument(
+            '--output-dir',
+            type=Path,
+            default=Path('/root/models/current'),
+            help='Output directory for fine-tuned model and converted gguf (default: /root/models/current)'
+        )
+        self.args = parser.parse_args()
+
+    @property
+    def config_path(self) -> Path:
+        return self.args.config
+
+    @property
+    def model(self) -> str:
+        return self.args.model
+
+    @property
+    def api_key(self) -> str:
+        return self.args.api_key
+
+    @property
+    def output_dir(self) -> Path:
+        return self.args.output_dir