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

tools/mistral_infer/pyproject.toml

@@ -0,0 +1,18 @@
+[build-system]
+requires = ["setuptools>=42", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "mistral_infer"
+version = "0.1.0"
+requires-python = ">=3.8"
+
+dependencies = [
+    "llm_engine_utils @ file:///root/sia/lib/llm_engine_utils",
+    "mistral-common>=1.0.0",
+    "mistralai>=0.0.7",
+    "python-dotenv>=1.0.0",
+]
+
+[project.scripts]
+mistral_infer = "mistral_infer.__main__:main"

tools/mistral_infer/src/mistral_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='Mistral API Inference')
+    
+    parser.add_argument(
+        '--model',
+        type=str,
+        default=os.getenv('SIA_MISTRAL_MODEL', 'mistral-large-latest'),
+        help='Model name (default: mistral-large-latest, env: SIA_MISTRAL_MODEL)'
+    )
+
+    parser.add_argument(
+        '--token-limit',
+        type=int,
+        default=int(os.getenv('SIA_MISTRAL_TOKEN_LIMIT', 128000)),
+        help='Token limit for the model (default: 128000, env: SIA_MISTRAL_TOKEN_LIMIT)'
+    )
+
+    parser.add_argument(
+        '--api-key',
+        type=str,
+        default=os.getenv('SIA_MISTRAL_API_KEY'),
+        help='API key for the model (required, env: SIA_MISTRAL_API_KEY)'
+    )
+
+    args = parser.parse_args()
+
+    mistral_llm_engine = MistralLlmEngine(
+        model=args.model,
+        token_limit=args.token_limit,
+        api_key=args.api_key,
+    )
+    process(mistral_llm_engine)
+
+if __name__ == "__main__":
+    main()

tools/mistral_infer/src/mistral_infer/mistral_llm_engine.py

@@ -0,0 +1,78 @@
+from llm_engine_utils import LlmEngine
+from llm_engine_utils.iterators import skip_prefix
+from mistral_common.protocol.instruct.messages import SystemMessage, UserMessage
+from mistral_common.protocol.instruct.request import ChatCompletionRequest
+from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
+from mistralai import Mistral
+from pathlib import Path
+from typing import Iterator
+
+class MistralLlmEngine(LlmEngine):
+    def __init__(
+            self,
+            model: str,
+            api_key: str,
+            token_limit: int,
+        ):
+        self._model = model
+        self._api_key = api_key
+        self._token_limit = token_limit
+        self._client = Mistral(api_key=api_key)
+        self._tokenizer = MistralTokenizer.v3()
+
+    def infer_xml(self, schema: Path, system: str, context: str, prefix: str) -> Iterator[str]:
+        messages = [
+            {
+                "role": "system",
+                "content": system,
+            },
+            {
+                "role": "user",
+                "content": context,
+            },
+            {
+                "role": "assistant",
+                "content": prefix,
+                "prefix": True,
+            },
+        ] if prefix else [
+            {
+                "role": "system",
+                "content": system,
+            },
+            {
+                "role": "user",
+                "content": context,
+            },
+        ]
+        stream_response = self._client.chat.stream(
+            model=self._model,
+            messages=messages,
+            #temperature=self._temperature,
+        )
+        
+        try:
+            def content_generator():
+                for chunk in stream_response:
+                    if content := chunk.data.choices[0].delta.content:
+                        yield content
+            yield from skip_prefix(content_generator(), prefix)
+        finally:
+            stream_response.response.close()
+
+    def token_count(self, system: str, context: str) -> int:
+        messages = [
+            SystemMessage(content=system),
+            UserMessage(content=context),
+        ]
+
+        tokenized = self._tokenizer.encode_chat_completion(
+            ChatCompletionRequest(
+                messages=messages,
+                model=self._model
+            )
+        )
+        return len(tokenized.tokens)
+
+    def token_limit(self) -> int:
+        return self._token_limit

tools/mistral_train/pyproject.toml

@@ -0,0 +1,19 @@
+[build-system]
+requires = ["setuptools>=42", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "mistral_train"
+version = "0.1.0"
+requires-python = ">=3.8"
+
+dependencies = [
+    "llm_engine_utils[dataset] @ file:///root/sia/lib/llm_engine_utils",
+    "mistral-common>=1.0.0",
+    "mistralai>=0.0.7",
+    "python-dotenv>=1.0.0",
+    "requests>=2.28.0",
+]
+
+[project.scripts]
+mistral_train = "mistral_train.__main__:main"

tools/mistral_train/src/mistral_train/__main__.py

@@ -1,54 +1,11 @@
-#!/root/venvs/train/bin/python
-"""
-Script for fine-tuning Mistral models for SIA using the Mistral API.
-"""
-from dataclasses import dataclass
 from pathlib import Path
-import argparse
 import json
 import os
 import sys
 import tempfile
 import requests
 
-# Import from our shared library
-from .util import TrainingParams, DatasetCreator
-
-@dataclass
-class Config:
-    def __init__(self):
-        parser = argparse.ArgumentParser(description='Train SIA model using Mistral API')
-        parser.add_argument(
-            '--config',
-            type=Path,
-            default=Path('/root/sia/training/config.yaml'),
-            help='Path to config file'
-        )
-        parser.add_argument(
-            '--model',
-            type=str,
-            default='mistral-large-latest',
-            help='Base model for fine-tuning'
-        )
-        parser.add_argument(
-            '--api-key',
-            type=str,
-            default=os.environ.get('SIA_MISTRAL_API_KEY'),
-            help='Mistral API key'
-        )
-        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
+from .config import Config
 
 def upload_file(api_key: str, file_path: Path) -> str:
     """Upload a file to the Mistral API and return the file ID"""
@@ -97,9 +54,6 @@ def start_finetune_job(api_key: str, model: str, file_id: str, params: sia_train
 
 def main():
     config = Config()
-    if not config.api_key:
-        print("Error: Mistral API key not found. Set SIA_MISTRAL_API_KEY environment variable.")
-        return 1
     
     training_data, train_params, commit_hash = sia_train_lib.prepare_training_data(config.config_path)
     

tools/mistral_train/src/mistral_train/config.py

@@ -0,0 +1,38 @@
+from pathlib import Path
+import argparse
+import os
+
+class Config:
+    def __init__(self):
+        parser = argparse.ArgumentParser(description='Train SIA model using Mistral API')
+        parser.add_argument(
+            '--config',
+            type=Path,
+            default=Path('/root/sia/training/config.yaml'),
+            help='Path to config file'
+        )
+        parser.add_argument(
+            '--model',
+            type=str,
+            default='mistral-large-latest',
+            help='Base model for fine-tuning'
+        )
+        parser.add_argument(
+            '--api-key',
+            type=str,
+            default=os.environ.get('SIA_MISTRAL_API_KEY'),
+            help='Mistral API key'
+        )
+        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

tools/old_llm_engines/hf_llm_engine.py

@@ -0,0 +1,87 @@
+from huggingface_hub import InferenceClient
+from transformers import AutoTokenizer, AutoConfig
+from typing import Iterator, Optional, Callable
+
+from . import LlmEngine
+
+class HfLlmEngine(LlmEngine):
+    """
+    LLM Engine implementation using HuggingFace's InferenceClient.
+    """
+    
+    def __init__(
+        self,
+        model: str,
+        temperature: float,
+        api_token: Optional[str],
+    ):
+        """
+        Initialize the HuggingFace Inference API LLM Engine.
+
+        Args:
+            model: HuggingFace model ID to use
+            temperature: Sampling temperature
+            api_token: HuggingFace API token
+        """
+        self._model = model
+        self._temperature = temperature
+
+        self._tokenizer = AutoTokenizer.from_pretrained(model, token=api_token)
+        self._config = AutoConfig.from_pretrained(model, token=api_token)
+        self._client = InferenceClient(token=api_token)
+                        
+    def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
+        """
+        Run inference using the system prompt and main context.
+
+        Args:
+            system_prompt: The system prompt string
+            main_context: The main context string after templating
+            continuation_text: Part of the response that is already generated
+            should_stop: Callback that returns True when inference should stop
+
+        Returns:
+            Iterator[str]: An iterator that yields the generated text.
+        """
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context},
+            {"role": "assistant", "content": continuation_text},
+        ]
+        
+        stream = self._client.chat_completion(
+            model=self._model,
+            messages=messages,
+            temperature=self._temperature,
+            add_generation_prompt=False,
+            stream=True
+        )
+        
+        try:
+            for response in stream:
+                if should_stop():
+                    stream.close()
+                    break
+                if content := response.choices[0].delta.content:
+                    yield content
+        finally:
+            stream.close()
+
+    def token_count(self, system_prompt: str, main_context: str) -> int:
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context}
+        ]
+        prompt = self._tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        return len(self._tokenizer.encode(prompt))
+
+    def token_limit(self) -> int:
+        """
+        Get the model's context window size.
+        
+        Returns:
+            int: Maximum number of tokens the model can process
+        """
+        return self._config.max_position_embeddings

tools/old_llm_engines/local_llm_engine.py

@@ -0,0 +1,137 @@
+from threading import Thread
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
+from typing import Iterator, Optional, Callable
+from xml_schema_validator import XmlLogitsProcessor
+import sys
+import torch
+
+from . import LlmEngine
+from .. import util
+
+class LocalLlmEngine(LlmEngine):
+    def __init__(
+        self,
+        model_path: str,
+        temperature: float,
+        token_limit: int,
+        xml_schema_text: Optional[str] = None,
+        api_token: Optional[str] = None,
+    ):
+        """
+        Initialize the LLM Engine with a model path.
+
+        Args:
+            model_path: Path to the model weights to be used.
+            temperature: Temperature for sampling
+            token_limit: Maximum number of tokens to generate
+            xml_schema_text: Optional XML schema to validate against
+            api_token: Huggingface API key
+        """
+        self._temperature = temperature
+        self._token_limit = token_limit
+        self._tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            return_dict=True,
+            low_cpu_mem_usage=True,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            trust_remote_code=True,
+            token=api_token,
+        )
+        if self._tokenizer.pad_token_id is None:
+            self._tokenizer.pad_token_id = self._tokenizer.eos_token_id
+        if model.config.pad_token_id is None:
+            model.config.pad_token_id = model.config.eos_token_id
+        self._pipeline = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=self._tokenizer,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            return_full_text=False,
+        )
+        if xml_schema_text:
+            self._logits_processor = XmlLogitsProcessor(self._tokenizer, xml_schema_text)
+        else:
+            self._logits_processor = None
+
+    def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
+        """
+        Run inference using the system prompt and main context.
+
+        Args:
+            system_prompt: The system prompt string
+            main_context: The main context string after templating
+            continuation_text: Part of the response that is already generated
+            should_stop: Callback that returns True when inference should stop
+
+        Returns:
+            Iterator[str]: An iterator that yields the generated text.
+        """
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context},
+            {"role": "assistant", "content": continuation_text},
+        ]
+        prompt = self._tokenizer.apply_chat_template(
+            messages, tokenize=False,
+            add_generation_prompt=False,
+        )
+        streamer = TextIteratorStreamer(
+            self._tokenizer,
+            skip_prompt=True
+        )
+        generation_kwargs = {
+            "text_inputs": prompt,
+            "do_sample": True,
+            "temperature": self._temperature,
+            "max_new_tokens": self.token_limit(),
+            "streamer": streamer,
+        }
+        if self._logits_processor:
+            generation_kwargs["logits_processor"] = [self._logits_processor.copy()]
+        generation_thread = Thread(
+            target=self._pipeline,
+            kwargs=generation_kwargs
+        )
+        generation_thread.start()
+
+        for text in util.stop_before_value(streamer, self._tokenizer.eos_token):
+            yield text
+            if should_stop():
+                break
+
+        generation_thread.join()
+
+    def token_count(self, system_prompt: str, main_context: str) -> int:
+        """
+        Count tokens for the given system prompt and main context.
+        
+        Args:
+            system_prompt: The system prompt string
+            main_context: The main context string
+            
+        Returns:
+            int: Total number of tokens
+        """
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context}
+        ]
+        prompt = self._tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        return len(self._tokenizer.encode(prompt))
+
+    def token_limit(self) -> int:
+        """
+        Get the model's context window size.
+        
+        Returns:
+            int: Maximum number of tokens the model can process
+        """
+        if self._token_limit is not None:
+            return self._token_limit
+        else:
+            return self._pipeline.model.config.max_position_embeddings

tools/old_llm_engines/openai_llm_engine.py

@@ -0,0 +1,77 @@
+from typing import Callable, Iterator
+import openai
+import tiktoken
+
+from . import LlmEngine
+
+class OpenAILlmEngine(LlmEngine):
+    """
+    LLM Engine implementation using OpenAI's API.
+    Supports streaming responses from chat completion models.
+    """
+    
+    def __init__(
+        self,
+        model: str,
+        temperature: float,
+        token_limit: int,
+        api_key: str,
+    ):
+        """
+        Initialize the OpenAI LLM Engine.
+
+        Args:
+            model: OpenAI model to use
+            temperature: Temperature for sampling
+            api_key: OpenAI API key
+            token_limit: Maximum number of tokens to generate
+        """
+        self._model = model
+        self._temperature = temperature
+        self._token_limit = token_limit
+
+        self._client = openai.Client(
+            api_key=api_key,
+        )
+                        
+    def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
+        if continuation_text:
+            print("OpenAI LLM Engine: continuation_text is not supported")
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context}
+        ]
+        
+        stream = self._client.chat.completions.create(
+            model=self._model,
+            messages=messages,
+            temperature=self._temperature,
+            stream=True,
+        )
+
+        try:
+            for chunk in stream:
+                if should_stop():
+                    break
+                if content := chunk.choices[0].delta.content:
+                    yield content
+        finally:
+            stream.close()
+            #stream.response.close()
+    
+    def token_count(self, system_prompt: str, main_context: str) -> int:
+        """
+        Calculate the total token count for the system prompt and context.
+        
+        Args:
+            system_prompt: The system prompt string
+            main_context: The main context string
+            
+        Returns:
+            int: Total number of tokens
+        """
+        encoding = tiktoken.encoding_for_model(self._model)
+        return len(encoding.encode(system_prompt)) + len(encoding.encode(main_context))
+    
+    def token_limit(self) -> int:
+        return self._token_limit

tools/old_llm_engines/qwq_llm_engine.ipynb

@@ -0,0 +1,227 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "source /root/venvs/sia/bin/activate\n",
+    "apt-get update && apt-get install -y cuda-toolkit\n",
+    "pip install flash-attn --no-build-isolation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Unsloth should be imported before transformers to ensure all optimizations are applied.\n",
+    "from unsloth import FastLanguageModel"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "from threading import Thread\n",
+    "from transformers import AutoTokenizer, TextIteratorStreamer, pipeline\n",
+    "from xml_schema_validator import XmlLogitsProcessor"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "temperature = 0.6\n",
+    "model_path = \"/root/models/current\"\n",
+    "xml_schema_text = Path(\"/root/sia/action_schema.xsd\").read_text()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load tokenizer\n",
+    "tokenizer = AutoTokenizer.from_pretrained(\n",
+    "    model_path,\n",
+    "    legacy=False,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load model\n",
+    "model, _returned_tokenizer = FastLanguageModel.from_pretrained(\n",
+    "    model_path,\n",
+    "    gpu_memory_utilization = 0.5, # Reduce if out of memory\n",
+    "    load_in_4bit=True,\n",
+    "    attn_implementation=\"flash_attention_2\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# enable unsloth optimizations\n",
+    "FastLanguageModel.for_inference(model)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "# Create inference pipeline with memory-efficient settings\n",
+    "pipeline = pipeline(\n",
+    "    \"text-generation\",\n",
+    "    model=model,\n",
+    "    tokenizer=tokenizer,\n",
+    "    return_full_text=False,\n",
+    "    torch_dtype=torch.float16,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_text)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "messages = [\n",
+    "    {\"role\": \"system\", \"content\": \"Always respond in a <write_stdout></write_stdout> xml block.\"},\n",
+    "    {\"role\": \"user\", \"content\": \"Hi, how are you?\"},\n",
+    "    {\"role\": \"assistant\", \"content\": \"\"},\n",
+    "]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "text = tokenizer.apply_chat_template(\n",
+    "    messages,\n",
+    "    tokenize=False,\n",
+    "    add_generation_prompt=False,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "streamer = TextIteratorStreamer(\n",
+    "    tokenizer,\n",
+    "    skip_prompt=True,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "generation_kwargs = {\n",
+    "    \"text_inputs\": text,\n",
+    "    \"do_sample\": True,\n",
+    "    \"temperature\": temperature,\n",
+    "    \"streamer\": streamer,\n",
+    "    \"use_cache\": True,\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "generation_kwargs[\"logits_processor\"] = [logits_processor.copy()]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "generation_thread = Thread(\n",
+    "    target=pipeline,\n",
+    "    kwargs=generation_kwargs\n",
+    ")\n",
+    "\n",
+    "generation_thread.start()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for text in streamer:\n",
+    "    print(text, end=\"\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "generation_thread.join()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "sia",
+   "language": "python",
+   "name": "sia"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

tools/old_llm_engines/qwq_llm_engine.py

@@ -0,0 +1,135 @@
+# Unsloth should be imported before transformers to ensure all optimizations are applied.
+from unsloth import FastLanguageModel
+
+from pathlib import Path
+from threading import Thread
+from transformers import AutoTokenizer, TextIteratorStreamer, pipeline
+from typing import Callable, Iterator, Optional
+from xml_schema_validator import XmlLogitsProcessor
+
+from . import LlmEngine
+from .. import util
+
+class QwQLlmEngine(LlmEngine):
+    
+    def __init__(
+        self,
+        model_path: Path,
+        temperature: float,
+        xml_schema_text: Optional[str] = None,
+    ):
+        """
+        Initialize the QwQ LLM Engine.
+
+        Args:
+            model_path: Local path to the model
+            temperature: Sampling temperature
+            xml_schema_text: Optional XML schema to validate against
+        """
+        self._temperature = temperature
+
+        # Load tokenizer
+        self._tokenizer = AutoTokenizer.from_pretrained(
+            model_path,
+        )
+        
+        # Load model
+        self._model, _returned_tokenizer = FastLanguageModel.from_pretrained(
+            model_path,
+            gpu_memory_utilization = 0.5, # Reduce if out of memory
+        )
+
+        # enable unsloth optimizations
+        FastLanguageModel.for_inference(self._model)
+        
+        # Create inference pipeline
+        self._pipeline = pipeline(
+            "text-generation",
+            model=self._model,
+            tokenizer=self._tokenizer,
+            return_full_text=False,
+        )
+
+        if xml_schema_text:
+            self._logits_processor = XmlLogitsProcessor(self._tokenizer, xml_schema_text)
+        else:
+            self._logits_processor = None
+
+                        
+    def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
+        """
+        Run inference using the system prompt and main context.
+
+        Args:
+            system_prompt: The system prompt string
+            main_context: The main context string after templating
+            continuation_text: Part of the response that is already generated
+            should_stop: Callback that returns True when inference should stop
+
+        Returns:
+            Iterator[str]: An iterator that yields the generated text.
+        """
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context},
+            {"role": "assistant", "content": continuation_text},
+        ]
+
+        text = self._tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=False,
+        )
+
+        streamer = TextIteratorStreamer(
+            self._tokenizer,
+            skip_prompt=True,
+        )
+
+        generation_kwargs = {
+            "text_inputs": text,
+            "do_sample": True,
+            "temperature": self._temperature,
+            "streamer": streamer,
+            "use_cache": True,
+        }
+        
+        if self._logits_processor:
+            generation_kwargs["logits_processor"] = [self._logits_processor.copy()]
+
+        generation_thread = Thread(
+            target=self._pipeline,
+            kwargs=generation_kwargs
+        )
+
+        generation_thread.start()
+
+        for text in util.stop_before_value(streamer, self._tokenizer.eos_token):
+            yield text
+            if should_stop():
+                break
+
+        generation_thread.join()
+
+    def token_count(self, system_prompt: str, main_context: str) -> int:
+        """
+        Count tokens for the given system prompt and main context.
+        
+        Args:
+            system_prompt: The system prompt string
+            main_context: The main context string
+            
+        Returns:
+            int: Total number of tokens
+        """
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context}
+        ]
+        prompt = self._tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        return len(self._tokenizer.encode(prompt))
+
+    def token_limit(self) -> int:
+        return self._pipeline.model.config.max_position_embeddings

tools/old_llm_engines/qwq_vllm.ipynb

@@ -0,0 +1,255 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# vLLM Streaming Implementation\n",
+    "\n",
+    "This notebook demonstrates how to implement streaming capability with vLLM, comparable to the unsloth implementation.\n",
+    "\n",
+    "First, let's make sure we have vLLM installed:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "INFO 04-25 19:36:31 [__init__.py:239] Automatically detected platform cuda.\n"
+     ]
+    }
+   ],
+   "source": [
+    "from pathlib import Path\n",
+    "from vllm import SamplingParams\n",
+    "from transformers import AutoTokenizer\n",
+    "import sys"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "temperature = 0.6\n",
+    "model_path = \"/root/models/current\"\n",
+    "xml_schema_text = Path(\"/root/sia/action_schema.xsd\").read_text()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load tokenizer\n",
+    "tokenizer = AutoTokenizer.from_pretrained(\n",
+    "    model_path,\n",
+    "    legacy=False,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "messages = [\n",
+    "    {\"role\": \"system\", \"content\": \"Always respond in a <write_stdout></write_stdout> xml block.\"},\n",
+    "    {\"role\": \"user\", \"content\": \"Hi, how are you?\"},\n",
+    "    {\"role\": \"assistant\", \"content\": \"\"},\n",
+    "]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "prompt = tokenizer.apply_chat_template(\n",
+    "    messages,\n",
+    "    tokenize=False,\n",
+    "    add_generation_prompt=False,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Define sampling parameters\n",
+    "sampling_params = SamplingParams(\n",
+    "    temperature=temperature,\n",
+    "    top_p=0.95,\n",
+    "    max_tokens=512,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "INFO 04-25 19:36:40 [config.py:585] This model supports multiple tasks: {'generate', 'score', 'reward', 'embed', 'classify'}. Defaulting to 'generate'.\n",
+      "WARNING 04-25 19:36:42 [config.py:664] bitsandbytes quantization is not fully optimized yet. The speed can be slower than non-quantized models.\n",
+      "WARNING 04-25 19:36:42 [arg_utils.py:1854] --quantization bitsandbytes is not supported by the V1 Engine. Falling back to V0. \n",
+      "INFO 04-25 19:36:42 [llm_engine.py:241] Initializing a V0 LLM engine (v0.8.2) with config: model='/root/models/current', speculative_config=None, tokenizer='/root/models/current', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, override_neuron_config=None, tokenizer_revision=None, trust_remote_code=True, dtype=torch.bfloat16, max_seq_len=4096, download_dir=None, load_format=bitsandbytes, tensor_parallel_size=1, pipeline_parallel_size=1, disable_custom_all_reduce=False, quantization=bitsandbytes, enforce_eager=False, kv_cache_dtype=auto,  device_config=cuda, decoding_config=DecodingConfig(guided_decoding_backend='xgrammar', reasoning_backend=None), observability_config=ObservabilityConfig(show_hidden_metrics=False, otlp_traces_endpoint=None, collect_model_forward_time=False, collect_model_execute_time=False), seed=None, served_model_name=/root/models/current, num_scheduler_steps=1, multi_step_stream_outputs=True, enable_prefix_caching=None, chunked_prefill_enabled=False, use_async_output_proc=True, disable_mm_preprocessor_cache=False, mm_processor_kwargs=None, pooler_config=None, compilation_config={\"splitting_ops\":[],\"compile_sizes\":[],\"cudagraph_capture_sizes\":[256,248,240,232,224,216,208,200,192,184,176,168,160,152,144,136,128,120,112,104,96,88,80,72,64,56,48,40,32,24,16,8,4,2,1],\"max_capture_size\":256}, use_cached_outputs=False, \n",
+      "INFO 04-25 19:36:42 [cuda.py:291] Using Flash Attention backend.\n",
+      "INFO 04-25 19:36:43 [parallel_state.py:954] rank 0 in world size 1 is assigned as DP rank 0, PP rank 0, TP rank 0\n",
+      "INFO 04-25 19:36:43 [model_runner.py:1110] Starting to load model /root/models/current...\n",
+      "INFO 04-25 19:36:43 [loader.py:1155] Loading weights with BitsAndBytes quantization. May take a while ...\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "8b9f3cb293484cac932e6cedd841c813",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading safetensors checkpoint shards:   0% Completed | 0/4 [00:00<?, ?it/s]\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "54f8aa5eefdb43d8bc07274044a8bc1c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading safetensors checkpoint shards:   0% Completed | 0/4 [00:00<?, ?it/s]\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "INFO 04-25 19:36:51 [model_runner.py:1146] Model loading took 18.0523 GB and 8.113452 seconds\n",
+      "INFO 04-25 19:36:55 [worker.py:267] Memory profiling takes 3.23 seconds\n",
+      "INFO 04-25 19:36:55 [worker.py:267] the current vLLM instance can use total_gpu_memory (47.53GiB) x gpu_memory_utilization (0.90) = 42.78GiB\n",
+      "INFO 04-25 19:36:55 [worker.py:267] model weights take 18.05GiB; non_torch_memory takes 0.06GiB; PyTorch activation peak memory takes 1.59GiB; the rest of the memory reserved for KV Cache is 23.08GiB.\n",
+      "INFO 04-25 19:36:55 [executor_base.py:111] # cuda blocks: 5907, # CPU blocks: 1024\n",
+      "INFO 04-25 19:36:55 [executor_base.py:116] Maximum concurrency for 4096 tokens per request: 23.07x\n",
+      "INFO 04-25 19:36:58 [model_runner.py:1442] Capturing cudagraphs for decoding. This may lead to unexpected consequences if the model is not static. To run the model in eager mode, set 'enforce_eager=True' or use '--enforce-eager' in the CLI. If out-of-memory error occurs during cudagraph capture, consider decreasing `gpu_memory_utilization` or switching to eager mode. You can also reduce the `max_num_seqs` as needed to decrease memory usage.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Capturing CUDA graph shapes: 100%|██████████| 35/35 [01:01<00:00,  1.75s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "INFO 04-25 19:38:00 [model_runner.py:1570] Graph capturing finished in 61 secs, took 1.98 GiB\n",
+      "INFO 04-25 19:38:00 [llm_engine.py:447] init engine (profile, create kv cache, warmup model) took 68.57 seconds\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from vllm import LLM, SamplingParams\n",
+    "import time\n",
+    "\n",
+    "# Initialize LLM\n",
+    "llm = LLM(\n",
+    "    model=model_path,\n",
+    "    tensor_parallel_size=1,\n",
+    "    max_model_len=4096,\n",
+    "    quantization=\"bitsandbytes\",\n",
+    "    load_format=\"bitsandbytes\",\n",
+    "    trust_remote_code=True,\n",
+    "    # Enable streaming\n",
+    "    enable_lora=False,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Starting generation with token-by-token output:\n",
+      "<think>\n",
+      "Okay, the user greeted me with \"Hi, how are you?\" I need to respond appropriately. Let me see... The instructions say to always use the <write_stdout> XML tag. So first, I should acknowledge their greeting and state that I'm an AI, then ask how I can assist them. Keep it friendly and helpful. Let me make sure I don't add any extra information beyond that. Just a simple response. Alright, that should work.\n",
+      "</think>\n",
+      "\n",
+      "<write_stdout>\n",
+      "Hello! I'm just a computer program, but I'm here to help you. How can I assist you today?\n",
+      "</write_stdout>"
+     ]
+    }
+   ],
+   "source": [
+    "previous_text = \"\"\n",
+    "print(\"Starting generation with token-by-token output:\")\n",
+    "\n",
+    "# Try with direct iteration over the generator\n",
+    "for output in llm.generate(prompt, sampling_params, use_tqdm=False):\n",
+    "    if hasattr(output, 'outputs') and output.outputs and len(output.outputs) > 0:\n",
+    "        generated_text = output.outputs[0].text\n",
+    "        if len(generated_text) > len(previous_text):\n",
+    "            new_text = generated_text[len(previous_text):]\n",
+    "            sys.stdout.write(new_text)\n",
+    "            sys.stdout.flush()\n",
+    "            previous_text = generated_text"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "sia",
+   "language": "python",
+   "name": "sia"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

tools/train/bin/train

@@ -1,15 +0,0 @@
-#!/bin/bash
-
-set -e
-
-SIA_DIR="/root/sia"
-OUTPUT_DIR="${1:-/root/models/$(cd "$SIA_DIR" && git rev-parse HEAD)}"
-
-if [ -n "$(cd "$SIA_DIR" && git status --porcelain)" ]; then
-    echo "Uncommitted changes in SIA directory"
-    #exit 1
-fi
-
-mkdir -p "$OUTPUT_DIR"
-
-/root/venvs/train/bin/python -m train.qwq --output-dir "$OUTPUT_DIR"

tools/train/readme.md

@@ -1,68 +0,0 @@
-# SIA Training Tool
-
-This tool provides command-line utilities for fine-tuning SIA's language models.
-
-## Supported Models
-
-- DeepSeek R1 models (including distilled versions)
-- Mistral models
-
-## Commands
-
-### train_deepseek
-
-Fine-tune DeepSeek models using Unsloth optimization.
-
-```bash
-train_deepseek --base-model deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B --output-dir /root/models/DeepSeek-R1-Distill-Qwen-1.5B
-```
-
-Options:
-- `--config`: Path to training configuration file (default: /root/sia/training/config.yaml)
-- `--base-model`: HuggingFace model ID for the base model (required)
-- `--output-dir`: Directory to save model (required)
-- `--api-key`: HuggingFace API key (optional, will use SIA_HF_API_KEY)
-
-### train_mistral
-
-Fine-tune Mistral models using Mistral's API.
-
-```bash
-train_mistral --model mistral-large-latest
-```
-
-Options:
-- `--config`: Path to training configuration file (default: /root/sia/training/config.yaml)
-- `--model`: Base model name (default: mistral-large-latest)
-- `--api-key`: Mistral API key (optional, will use SIA_MISTRAL_API_KEY)
-
-## Configuration Format
-
-The training configuration file (YAML) should include:
-
-```yaml
-model:
-  system_prompt_path: "/root/sia/system_prompt.md"
-  action_schema: "/root/sia/action_schema.xsd"
-params:
-  learning_rate: 1e-5
-  epochs: 3
-data:
-  - "/root/sia/training/data_dir1/"
-  - "/root/sia/training/data_dir2/"
-```
-
-## Data Format
-
-Training data should be XML files in the following format:
-
-```xml
-<iteration system_prompt_hash="..." action_schema_hash="...">
-  <context>
-    <!-- XML context -->
-  </context>
-  <response>
-    <!-- Model response -->
-  </response>
-</iteration>
-```

tools/train/train/dataset.py

@@ -1,144 +0,0 @@
-from datasets import Dataset as TransformersDataset
-from transformers import PreTrainedTokenizer
-from pathlib import Path
-from typing import Dict, List, Iterator
-import hashlib
-import torch
-import xml.etree.ElementTree as ET
-import yaml
-
-class Dataset(torch.utils.data.Dataset):
-    """Training dataset from XML iteration files"""
-    
-    def __init__(self, config_filename: str):
-        with open(config_filename) as f:
-            config_data = yaml.safe_load(f)
-        
-        data_paths = [Path(p) for p in config_data['data']]
-        self.files = self._find_xml_files(data_paths)
-        
-        self.system_prompt_file = Path(config_data['model']['system_prompt_path'])
-        self.action_schema_file = Path(config_data['model']['action_schema'])
-        
-        self.system_prompt = self.system_prompt_file.read_text()
-        self.system_prompt_hash = self._calculate_hash(self.system_prompt)
-        
-        self.action_schema = self.action_schema_file.read_text()
-        self.action_schema_hash = self._calculate_hash(self.action_schema)
-        
-    def _find_xml_files(self, data_paths: List[Path]) -> List[Path]:
-        """Find all XML files in the given data paths"""
-        xml_files = list()
-        for path in data_paths:
-            if not path.exists():
-                raise Exception(f"Data path not found: {path}")
-            xml_files.extend(path.rglob('*.xml'))
-        return xml_files
-        
-    def _calculate_hash(self, content: str) -> str:
-        """Calculate SHA-256 hash of content"""
-        return hashlib.sha256(content.encode()).hexdigest()
-    
-    def _parse_iteration_file(self, file_path: Path) -> Dict:
-        """Parse a single iteration XML file into a training example"""
-        tree = ET.parse(file_path)
-        root = tree.getroot()
-        
-        context_elem = root.find('context')
-        response_elem = root.find('response')
-        
-        context = context_elem.text
-        response = response_elem.text
-        
-        return {
-            "messages": [
-                {
-                    "role": "system",
-                    "content": self.system_prompt + "\n" + self.action_schema
-                },
-                {
-                    "role": "user",
-                    "content": context
-                },
-                {
-                    "role": "assistant", 
-                    "content": response
-                }
-            ]
-        }
-    
-    def __len__(self) -> int:
-        """Return the number of samples in the dataset"""
-        return len(self.files)
-    
-    def __getitem__(self, idx: int) -> Dict:
-        """Indexing for a single sample"""
-        if idx < 0 or idx >= len(self):
-            raise IndexError(f"Index {idx} out of range for dataset with {len(self)} samples")
-        file_path = self.files[idx]
-        return self._parse_iteration_file(file_path)
-        
-    def __iter__(self) -> Iterator[Dict]:
-        """Allow iteration over samples"""
-        for i in range(len(self)):
-            yield self[i]
-    
-    def to_list(self) -> List[Dict]:
-        """Convert dataset to a list"""
-        results = []
-        for i in range(len(self)):
-            results.append(self[i])
-        return results
-    
-    def to_transformers_dataset(self, tokenizer: PreTrainedTokenizer) -> TransformersDataset:
-        def generator():
-            for item in self:
-                messages = item["messages"]
-                formatted_text = tokenizer.apply_chat_template(
-                    messages,
-                    tokenize=False,
-                    add_generation_prompt=False
-                )
-                yield {"messages": formatted_text}
-        return TransformersDataset.from_generator(generator)
-    
-    def validate(self) -> None:
-        """Validate XML files"""
-        print(f"Validating {len(self.files)} XML files...")
-        
-        for i in range(len(self.files)):
-            self.validate_sample(i)
-
-        print(f"Validation complete. Found {len(self.files)} valid files.")
-    
-    def validate_sample(self, index: int) -> None:
-        file = self.files[index]
-        print("file:", file)
-        tree = ET.parse(file)
-        root = tree.getroot()
-        
-        # Check system prompt hash
-        file_system_hash = root.get('system_prompt_hash')
-        if file_system_hash != self.system_prompt_hash:
-            print(f"WARNING: System prompt hash mismatch in {file}")
-        
-        # Check action schema hash
-        file_schema_hash = root.get('action_schema_hash')
-        if file_schema_hash != self.action_schema_hash:
-            print(f"WARNING: Action schema hash mismatch in {file}")
-        
-        # Check for required elements
-        context_elem = root.find('context')
-        response_elem = root.find('response')
-        
-        if context_elem is None:
-            raise Exception(f"Missing context element")
-            
-        if response_elem is None:
-            raise Exception(f"Missing response element")
-            
-        if not context_elem.text:
-            raise Exception(f"Empty context")
-            
-        if not response_elem.text:
-            raise Exception(f"Empty response")