wip deepseek r1

.gitignore

@@ -1,4 +1,5 @@
 .env
+__pycache__/
 data/
 model/
-__pycache__/
+sia.egg-info/

Dockerfile

@@ -1,57 +1,82 @@
-FROM nvidia/cuda:12.2.0-runtime-ubuntu22.04 AS requirements
+FROM nvidia/cuda:12.2.0-runtime-ubuntu22.04 AS base
-RUN apt-get update
+RUN apt-get update && \
-RUN apt-get upgrade -y
+    apt-get upgrade -y && \
-RUN apt install -y python3-pip
+    apt install -y \
-ENV TOKENIZERS_PARALLELISM=false
+    python3-pip \
-COPY requirements.txt /requirements.txt
+    git \
-RUN pip3 install -r /requirements.txt
+    python3-venv \
-RUN rm -rf /requirements.txt
+    wget \
-
+    gnupg \
-FROM requirements AS sia-test
+    vim \
-COPY ./ /root/sia/
+    curl
-WORKDIR /root/sia/
-RUN mkdir -p /root/models/current
-CMD ["python3", "-m", "unittest", "discover", "-v", "-p", "*test.py"]
-
-FROM node:20-alpine AS web-test
-WORKDIR /app
-COPY web/package*.json ./
-RUN npm install
-COPY web .
-RUN npm test
-
-FROM node:20-alpine AS web-build
-WORKDIR /app
-COPY web/package*.json ./
-RUN npm install
-COPY web .
-RUN npm run build
-
-FROM requirements
-RUN apt-get update
-RUN apt-get install -y wget gnupg
 RUN wget -q -O - https://dl-ssl.google.com/linux/linux_signing_key.pub | apt-key add -
-RUN echo "deb http://dl.google.com/linux/chrome/deb/ stable main" >> /etc/apt/sources.list.d/google.list
+RUN echo "deb http://dl.google.com/linux/chrome/deb/ stable main" >> /etc/apt/sources.list.d/google-chrome.list
-RUN apt-get update
+RUN apt-get update && \
-RUN apt-get install -y google-chrome-stable
+    apt-get install -y \
+    google-chrome-stable
 RUN rm -rf /var/lib/apt/lists/*
 
+# Create directory structure
 RUN mkdir -p \
     /root/sia \
+    /root/sia/scripts \
     /root/data/iterations \
     /root/data/user \
     /root/data/tasks \
     /root/data/environment \
     /root/models \
-    /root/desktop
+    /root/desktop \
+    /root/venvs
 
-COPY ./tools/itb/requirements.txt /root/sia/tools/itb/requirements.txt
+# ITB tool setup
-RUN cd /root/sia/tools/itb/ && python3 -m pip install -r requirements.txt
+FROM base AS itb-env
+COPY ./scripts/setup_binaries.py /root/sia/scripts/
+COPY ./tools/itb/setup.py /root/sia/tools/itb/setup.py
+RUN python3 /root/sia/scripts/setup_binaries.py /root/sia/tools/itb/setup.py
+RUN python3 -m venv /root/venvs/itb
+RUN /root/venvs/itb/bin/pip install -e /root/sia/tools/itb/
 
-COPY ./tools/ /root/sia/tools/
+# Train tool setup
-RUN cd /root/sia/tools/itb/ && python3 -m pip install -e ".[dev]"
+FROM base AS train-env
+COPY ./scripts/setup_binaries.py /root/sia/scripts/
+COPY ./tools/train/setup.py /root/sia/tools/train/setup.py
+RUN python3 /root/sia/scripts/setup_binaries.py /root/sia/tools/train/setup.py
+RUN python3 -m venv /root/venvs/train
+RUN /root/venvs/train/bin/pip install -e /root/sia/tools/train/
 
+# SIA core setup
+FROM base AS sia-env
+COPY ./scripts/setup_binaries.py /root/sia/scripts/
+COPY ./setup.py /root/sia/setup.py
+RUN python3 /root/sia/scripts/setup_binaries.py /root/sia/setup.py
+RUN python3 -m venv /root/venvs/sia
+RUN /root/venvs/sia/bin/pip install -e /root/sia/
+
+# Web frontend build
+FROM node:20-alpine AS web-build
+WORKDIR /app
+COPY web/package*.json ./
+RUN npm install
+RUN rm -rf /root/.npm/_cacache
+COPY web .
+RUN npm run build
+
+# Final image
+FROM base
+
+# Copy virtual environments (these layers only change if setup.py files change)
+COPY --from=itb-env /root/venvs/itb /root/venvs/itb
+COPY --from=train-env /root/venvs/train /root/venvs/train
+COPY --from=sia-env /root/venvs/sia /root/venvs/sia
+
+# Copy source code and scripts (these change frequently but don't affect venv layers)
+COPY --from=itb-env /root/sia/tools/itb /root/sia/tools/itb
+COPY --from=train-env /root/sia/tools/train /root/sia/tools/train
+COPY --from=sia-env /root/sia /root/sia
 COPY --from=web-build /app/dist /root/static/
+
+RUN echo 'for venv in /root/venvs/*/bin; do PATH="$venv:$PATH"; done' >> /etc/profile && \
+    echo 'export PATH' >> /etc/profile
 WORKDIR /root/desktop
 
-CMD ["/root/sia/scripts/restart.sh"]
+CMD ["/bin/bash", "-l", "-c", "/root/sia/scripts/restart.sh"]

procedures/self_improvement/reasoning.md

@@ -309,8 +309,9 @@ This preserves the temporal relationships between entries while anchoring them t
 
 ## Training Configuration
 
-SIA takes a modular approach to model training by having separate specialized tools for each provider like train_mistral.py, train_openai.py, etc.
+SIA takes a modular approach to model training by having separate specialized tools for each provider like train_mistral, train_deepseek, etc.
 Each tool shares similar core functionality while handling provider-specific requirements.
+The default training tool and parameters are called from the `/root/sia/tools/train/train.sh` script.
 
 While the training process is conceptually similar across providers, each has unique requirements for data formatting, API interactions, and job management.
 By creating dedicated tools, we can properly encapsulate these differences without complicating the core training logic.
@@ -341,29 +342,6 @@ This separation of concerns makes it easier to:
 - Handle provider-specific error cases and requirements appropriately
 - Update individual providers' implementations as their APIs evolve
 
-### Example
-
-Config file:
-
-```yaml
-model:
-  system_prompt_path: "system_prompt.md"
-  action_schema: "action_schema.xsd"
-params:
-  learning_rate: 1e-5
-  epochs: 3
-data:
-  - "training/clean_start/"
-  - "training/delete_indicated_entries/"
-  - "training/list_entries_to_delete/"
-```
-
-Training command:
-
-```bash
-python train_mistral.py --model mistral-large-latest
-```
-
 ## Repository Structure
 
 All components that define SIA's behavior are version controlled in a single repository, providing a clear and reproducible state for any point in time.

scripts/bootstrap.sh

@@ -0,0 +1,125 @@
+#!/bin/bash
+# bootstrap.sh - Initialize SIA (Self-Improving Agent) environment for cloud deployment
+
+set -eo pipefail # Exit on any error, pipe failures
+
+# Hardcoded paths for cloud deployment
+SIA_REPO_URL="ssh://git@git.nielsgeens.be:222/llm/SIA.git"
+SIA_DIR="/root/sia"
+DATA_DIR="/root/data"
+MODELS_DIR="/root/models"
+DESKTOP_DIR="/root/desktop"
+STATIC_DIR="/root/static"
+VENVS_DIR="/root/venvs"
+
+# Print header
+echo "==================================================="
+echo "SIA Bootstrap Script - Cloud Deployment"
+echo "==================================================="
+
+# Create directory structure
+echo "Creating directory structure..."
+mkdir -p "$DATA_DIR/iterations"
+mkdir -p "$DESKTOP_DIR"
+mkdir -p "$VENVS_DIR"
+cd "$DESKTOP_DIR"
+
+# Set up SSH keys
+echo "Setting up SSH keys for git access..."
+mkdir -p ~/.ssh
+chmod 700 ~/.ssh
+ssh-keygen -t sia_git -N "" -f ~/.ssh/sia_git -C "sia-agent"
+echo "New SSH key generated"
+
+# Display public key for user to add to git server
+echo "==================================================="
+echo "Add this public key to your git server:"
+cat ~/.ssh/sia_git.pub
+echo "==================================================="
+
+# Prompt user to confirm they've added the key
+read -p "Press Enter once you've added the SSH key to the git server..."
+
+# Clone SIA repository
+echo "Cloning SIA repository..."
+git clone "$SIA_REPO_URL" "$SIA_DIR"
+
+# Create and setup virtual environments
+echo "Setting up SIA virtual environments..."
+
+# Setup ITB tool environment
+echo "Creating ITB tool environment..."
+python3 -m venv "$VENVS_DIR/itb"
+"$VENVS_DIR/itb/bin/pip" install -e "$SIA_DIR/tools/itb"
+
+# Setup Train tool environment
+echo "Creating Train tool environment..."
+python3 -m venv "$VENVS_DIR/train"
+"$VENVS_DIR/train/bin/pip" install -e "$SIA_DIR/tools/train"
+
+# Setup SIA core environment
+echo "Creating SIA core environment..."
+python3 -m venv "$VENVS_DIR/sia"
+"$VENVS_DIR/sia/bin/pip" install -e "$SIA_DIR"
+
+# Build web interface
+echo "Building web interface"
+cd "$SIA_DIR/web"
+
+# Install Node.js if needed
+if ! command -v node &> /dev/null; then
+    echo "Installing Node.js..."
+    curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.1/install.sh | bash
+    export NVM_DIR="$HOME/.nvm"
+    [ -s "$NVM_DIR/nvm.sh" ] && \. "$NVM_DIR/nvm.sh"
+    nvm install node
+fi
+
+npm install
+npm run build
+
+mkdir -p "$STATIC_DIR"
+cp -r "$SIA_DIR/web/dist/"* "$STATIC_DIR/"
+
+echo "Web interface built successfully"
+
+# Finetune model
+echo "Starting model finetuning..."
+
+COMMIT_ID=$(cd "$SIA_DIR" && git rev-parse HEAD)
+echo "Current commit: $COMMIT_ID"
+
+mkdir -p "$MODELS_DIR/$COMMIT_ID"
+mkdir -p "$MODELS_DIR/current"
+
+# Run finetuning using the train environment
+"$VENVS_DIR/train/bin/train_deepseek" --output-dir "$MODELS_DIR/$COMMIT_ID"
+
+ln -sf "$MODELS_DIR/$COMMIT_ID" "$MODELS_DIR/current"
+echo "Finetuning complete, model linked to current"
+
+# Initialize environment information
+echo "Initializing environment information..."
+mkdir -p "$DATA_DIR/environment"
+cat > "$DATA_DIR/environment/sia_repo.md" << EOF
+# SIA Repository Information
+
+- Repository URL: $SIA_REPO_URL
+- ssh key: ~/.ssh/sia_git.pub
+EOF
+
+# Create .env file for local model only
+cat > "$SIA_DIR/.env" << EOF
+SIA_DEEPSEEK_ENABLED=true
+SIA_DEEPSEEK_MODEL=$MODELS_DIR/current
+SIA_DEEPSEEK_TEMPERATURE=0.6
+EOF
+
+# Print header
+echo "==================================================="
+echo "SIA environment initialization complete!"
+echo "==================================================="
+
+# Start SIA using restart script
+echo "Starting SIA..."
+"$SIA_DIR/scripts/restart.sh"

scripts/collect.sh

@@ -6,7 +6,7 @@ declare -A FILTER_SETS=(
   ["py"]="-f .*(\\.py|requirements.txt)$"
   ["web"]="-f .*\\.(js|jsx|json|css|html)$"
   ["doc"]="-f .*\\.md$"
-  ["deploy"]="-f .*(Dockerfile|\\.sh|\\.xsd)$"
+  ["deploy"]="-f .*(Dockerfile|\\.sh|\\.xsd|\\.yaml)$"
   
   ["core"]="-s py ./sia ./tools -s deploy . -f ^(?!procedures/).*\\.md$ ."
   ["webui"]="-s web ./web"
@@ -291,4 +291,4 @@ main() {
   echo "Concatenation complete. Output written to $OUTPUT" >&2
 }
 
-main "$@"
+main "$@"

scripts/install.sh

@@ -1,23 +0,0 @@
-#!/bin/bash
-
-cd /
-git clone https://git.nielsgeens.be/llm/SIA.git
-cd /SIA
-pip3 install -r requirements.txt
-
-curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.1/install.sh | bash
-nvm install node
-cd /SIA/web
-npm install
-npm install -D tailwindcss
-npm run build
-mv /root/SIA/web/dist/ /root/SIA/static
-
-apt update
-apt install -y vim tmux
-vim .env
-
-cd /root/SIA
-python3 -m sia
-
-#The SIA source is located in /root/sia. Not all features are implemented yet. Look at the readme and code to find what is missing. Make sure to unit test your work.

scripts/restart.sh

@@ -1,7 +1,7 @@
 #!/bin/bash
 
 while true; do
-    PYTHONPATH="/root/sia:$PYTHONPATH" python3 -m sia
+    sia
     if [ $? -eq 42 ]; then
         echo "SIA exited with code 42. Restarting."
     else

scripts/setup_binaries.py

@@ -0,0 +1,64 @@
+#!/usr/bin/env python3
+"""
+Script to extract binary names from setup.py files and create placeholder files.
+Usage: setup_binaries.py /path/to/setup.py
+"""
+
+import os
+import sys
+import re
+
+def extract_setup_binaries(setup_path):
+    """Extract binary names from a setup.py file and create placeholder files."""
+    try:
+        # Read the setup.py file
+        with open(setup_path, 'r') as f:
+            setup_content = f.read()
+        
+        # Find all references to scripts in bin/ directory
+        scripts = re.findall(r"'bin/[^']+?'|\"bin/[^\"]+?\"", setup_content)
+        if not scripts:
+            print(f"No bin scripts found in {setup_path}")
+            return True  # Not an error, just no scripts
+        
+        # Clean up the extracted script names
+        scripts = [script.strip('\'"') for script in scripts]
+        
+        # Create placeholder files
+        base_dir = os.path.dirname(setup_path)
+        created_count = 0
+        
+        for script in scripts:
+            script_path = os.path.join(base_dir, script)
+            script_dir = os.path.dirname(script_path)
+            
+            # Create directory if it doesn't exist
+            os.makedirs(script_dir, exist_ok=True)
+            
+            # Create an empty file
+            with open(script_path, 'w') as f:
+                pass
+            
+            # Make the file executable
+            os.chmod(script_path, 0o755)
+            created_count += 1
+        
+        print(f"Created {created_count} placeholder binary files from {setup_path}")
+        return True
+    
+    except Exception as e:
+        print(f"Error processing {setup_path}: {e}")
+        return False
+
+if __name__ == "__main__":
+    if len(sys.argv) != 2:
+        print(f"Usage: {sys.argv[0]} /path/to/setup.py")
+        sys.exit(1)
+    
+    setup_path = sys.argv[1]
+    if not os.path.exists(setup_path):
+        print(f"Error: {setup_path} not found")
+        sys.exit(1)
+    
+    success = extract_setup_binaries(setup_path)
+    sys.exit(0 if success else 1)

scripts/test.sh

@@ -1,12 +1,20 @@
 #!/bin/bash
 
 docker build \
-    --target sia-test \
+    --tag sia \
-    --tag sia-test \
     .
 
+# Run tests within the SIA virtual environment
 docker run \
     --rm \
+    -ti \
     --gpus=all \
-    -v /$(pwd)/model/:/root/model/ \
+    -p 8080:8080 \
-    sia-test
+    --env-file .env \
+    -v /$(pwd)/model/:/root/models/current/ \
+    -v /$(pwd)/iterations/:/root/data/iterations/ \
+    -v /$(pwd)/tasks/:/root/data/tasks/ \
+    -v /$(pwd)/user/:/root/data/user/ \
+    -v /$(pwd)/environment/:/root/data/environment/ \
+    -v /$(pwd)/:/root/sia/ \
+    sia /root/venvs/sia/bin/python -m unittest discover -v -p "*test.py"

setup.py

@@ -0,0 +1,33 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="sia",
+    version="0.1.0",
+    packages=find_packages(),
+    install_requires=[
+        'aiohttp>=3.8.0',
+        'dotenv-python>=0.0.1',
+        'huggingface_hub>=0.16.0',
+        'lxml>=4.9.0',
+        'mistralai>=0.0.7',
+        'mistral-common>=1.0.0',
+        'openai>=1.0.0',
+        'psutil>=5.9.0',
+        'python-dotenv>=1.0.0',
+        'tiktoken>=0.4.0',
+        'torch>=2.0.0',
+        'transformers>=4.30.0'
+    ],
+    entry_points={
+        'console_scripts': [
+            'sia=sia.__main__:main',
+        ],
+    },
+    classifiers=[
+        'Development Status :: 3 - Alpha',
+        'Intended Audience :: Developers',
+        'Programming Language :: Python :: 3',
+        'Programming Language :: Python :: 3.10',
+    ],
+    python_requires='>=3.10',
+)

sia/__main__.py

@@ -3,11 +3,12 @@ import asyncio
 
 from .auto_approver import AutoApprover
 from .config import Config
-from .hf_llm_engine import HfLlmEngine
+from .llm_engine.hf_llm_engine import HfLlmEngine
+from .llm_engine.deepseek_llm_engine import DeepSeekLlmEngine
 from .iteration_logger import IterationLogger
-from .local_llm_engine import LocalLlmEngine
+from .llm_engine.local_llm_engine import LocalLlmEngine
-from .mistral_llm_engine import MistralLlmEngine
+from .llm_engine.mistral_llm_engine import MistralLlmEngine
-from .openai_llm_engine import OpenAILlmEngine
+from .llm_engine.openai_llm_engine import OpenAILlmEngine
 from .response_parser import ResponseParser
 from .system_metrics import SystemMetrics
 from .web.api import Api
@@ -61,6 +62,14 @@ class Main:
                 config.mistral_api_key,
             )
 
+        if config.deepseek_enabled:
+            self._llms['deepseek'] = DeepSeekLlmEngine(
+                config.deepseek_model,
+                config.deepseek_temperature,
+                config.deepseek_token_limit,
+                config.hf_api_key,  # Use the existing HF API key
+            )
+
         if not self._llms:
             raise ValueError("No LLM engines enabled in configuration")
 
@@ -103,9 +112,10 @@ class Main:
             content_type="text/html"
         )
 
-if __name__ == "__main__":
+def main():
     loop = asyncio.new_event_loop()
     config = Config()
-    main = loop.run_until_complete(Main.create(config))
+    main_instance = loop.run_until_complete(Main.create(config))
     print(f"Web server started at http://localhost:{config.port}")
-    web.run_app(main.app, loop=loop, host=config.host, port=config.port)
+    web.run_app(main_instance.app, loop=loop, host=config.host, port=config.port)
+    return 0

sia/config.py

@@ -184,6 +184,30 @@ class Config:
             default=os.getenv('SIA_MISTRAL_API_KEY'),
             help='Mistral API key (env: SIA_MISTRAL_API_KEY)'
         )
+        parser.add_argument(
+            '--deepseek-enable',
+            action='store_true',
+            default=self._parse_bool_env('SIA_DEEPSEEK_ENABLED', False),
+            help='Enable DeepSeek LLM engine (env: SIA_DEEPSEEK_ENABLED)'
+        )
+        parser.add_argument(
+            '--deepseek-model',
+            type=str,
+            default=os.getenv('SIA_DEEPSEEK_MODEL', '/root/models/current'),
+            help='Path to fine-tuned DeepSeek model (env: SIA_DEEPSEEK_MODEL)'
+        )
+        parser.add_argument(
+            '--deepseek-temperature',
+            type=float,
+            default=float(os.getenv('SIA_DEEPSEEK_TEMPERATURE', '0.6')),
+            help='DeepSeek temperature (default: 0.6, env: SIA_DEEPSEEK_TEMPERATURE)'
+        )
+        parser.add_argument(
+            '--deepseek-token-limit',
+            type=int,
+            default=int(os.getenv('SIA_DEEPSEEK_TOKEN_LIMIT', '0')),
+            help='DeepSeek token limit (0 for model default, env: SIA_DEEPSEEK_TOKEN_LIMIT)'
+        )
 
         self.args = parser.parse_args()
 
@@ -312,3 +336,20 @@ class Config:
     @property
     def mistral_api_key(self) -> Optional[str]:
         return self.args.mistral_api_key
+
+    @property
+    def deepseek_enabled(self) -> bool:
+        return self.args.deepseek_enable
+
+    @property
+    def deepseek_model(self) -> str:
+        return self.args.deepseek_model
+
+    @property
+    def deepseek_temperature(self) -> float:
+        return self.args.deepseek_temperature
+
+    @property
+    def deepseek_token_limit(self) -> Optional[int]:
+        # Return None if 0 to use model default
+        return self.args.deepseek_token_limit if self.args.deepseek_token_limit > 0 else None

sia/llm_engine/deepseek_llm_engine.py

@@ -0,0 +1,150 @@
+from typing import Callable, Iterator, Optional
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from threading import Thread
+from pathlib import Path
+
+from . import LlmEngine
+from .. import util
+
+class DeepSeekLlmEngine(LlmEngine):
+    """
+    LLM Engine implementation for DeepSeek models.
+    Supports fine-tuned DeepSeek-R1 and its distilled versions.
+    """
+    
+    def __init__(
+        self,
+        model_path: str,
+        temperature: float = 0.6,
+        token_limit: Optional[int] = None,
+        api_key: Optional[str] = None,
+    ):
+        """
+        Initialize the DeepSeek LLM Engine.
+
+        Args:
+            model_path: Local path to the fine-tuned model
+            temperature: Sampling temperature (0.6 default as recommended)
+            token_limit: Maximum tokens to generate or context length override
+            api_key: HuggingFace API token if needed
+        """
+        self._model_path = Path(model_path)
+        self._temperature = temperature
+        self._token_limit = token_limit
+        
+        # Load tokenizer with trust_remote_code for DeepSeek models
+        self._tokenizer = AutoTokenizer.from_pretrained(
+            self._model_path,
+            token=api_key,
+            trust_remote_code=True,
+        )
+        
+        # Set padding token to avoid warnings
+        if self._tokenizer.pad_token is None:
+            self._tokenizer.pad_token = self._tokenizer.eos_token
+            
+        # Load model with 4-bit quantization by default
+        self._device_map = "auto"
+        
+        self._model = AutoModelForCausalLM.from_pretrained(
+            self._model_path,
+            return_dict=True,
+            low_cpu_mem_usage=True,
+            trust_remote_code=True,
+            device_map=self._device_map,
+            load_in_4bit=True, 
+            torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
+            token=api_key,
+        )
+        
+        # Ensure model is in evaluation mode
+        self._model.eval()
+                        
+    def infer(self, system_prompt: str, main_context: 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
+            should_stop: Callback that returns True when inference should stop
+
+        Returns:
+            Iterator[str]: An iterator that yields the generated text.
+        """
+        # Tokenize input
+        inputs = self._tokenizer(system_prompt + "\n\n" + main_context, return_tensors="pt").to(self._device_map)
+        
+        # Create streamer for token-by-token generation
+        streamer = TextIteratorStreamer(
+            self._tokenizer,
+            skip_prompt=True,
+            timeout=15.0
+        )
+        
+        # Generate in a separate thread to enable streaming
+        generation_kwargs = {
+            "input_ids": inputs.input_ids,
+            "attention_mask": inputs.attention_mask,
+            "max_new_tokens": self.token_limit() if self._token_limit else 2048,
+            "temperature": self._temperature,
+            "do_sample": True,
+            "streamer": streamer,
+            "repetition_penalty": 1.1,
+            "pad_token_id": self._tokenizer.pad_token_id,
+        }
+        
+        generation_thread = Thread(target=self._model.generate, kwargs=generation_kwargs)
+        generation_thread.start()
+        
+        # Yield tokens as they become available
+        try:
+            for text in streamer:
+                yield text
+                if should_stop():
+                    break
+        finally:
+            # Ensure thread is properly joined even if iteration is interrupted
+            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
+        """
+        combined_prompt = f"{system_prompt}\n\n{main_context}"
+        return len(self._tokenizer.encode(combined_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
+        
+        # Try to detect model size from config
+        try:
+            config_file = self._model_path / "config.json"
+            if config_file.exists():
+                import json
+                with open(config_file, 'r') as f:
+                    config = json.load(f)
+                if 'max_position_embeddings' in config:
+                    return config['max_position_embeddings']
+                if 'model_max_length' in config:
+                    return config['model_max_length']
+        except Exception:
+            pass
+            
+        # Default to 8k if we can't determine
+        return 8192

sia/llm_engine/hf_llm_engine.py

@@ -2,7 +2,7 @@ from huggingface_hub import InferenceClient
 from transformers import AutoTokenizer, AutoConfig
 from typing import Iterator, Optional, Callable
 
-from .llm_engine import LlmEngine
+from . import LlmEngine
 
 class HfLlmEngine(LlmEngine):
     """

sia/llm_engine/local_llm_engine.py

@@ -4,8 +4,8 @@ from typing import Iterator, Optional, Callable
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
 import torch
 
-from . import util
+from . import LlmEngine
-from .llm_engine import LlmEngine
+from .. import util
 
 class LocalLlmEngine(LlmEngine):
     def __init__(

sia/llm_engine/mistral_llm_engine.py

@@ -4,7 +4,7 @@ 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 .llm_engine import LlmEngine
+from . import LlmEngine
 
 class MistralLlmEngine(LlmEngine):
     def __init__(

sia/llm_engine/openai_llm_engine.py

@@ -2,7 +2,7 @@ from typing import Callable, Iterator
 import openai
 import tiktoken
 
-from .llm_engine import LlmEngine
+from . import LlmEngine
 
 class OpenAILlmEngine(LlmEngine):
     """

tools/itb/bin/itb_click

@@ -1,4 +1,4 @@
-#!/usr/bin/env python3
+#!/root/venvs/itb/bin/python
 import sys
 import os
 import time

tools/itb/bin/itb_cursor

@@ -1,4 +1,4 @@
-#!/usr/bin/env python3
+#!/root/venvs/itb/bin/python
 import sys
 import os
 import time

tools/itb/bin/itb_input

@@ -1,4 +1,4 @@
-#!/usr/bin/env python3
+#!/root/venvs/itb/bin/python
 import sys
 import os
 import time

tools/itb/bin/itb_navigate

@@ -1,4 +1,4 @@
-#!/usr/bin/env python3
+#!/root/venvs/itb/bin/python
 import sys
 import os
 import time

tools/itb/bin/itb_refresh

@@ -1,4 +1,4 @@
-#!/usr/bin/env python3
+#!/root/venvs/itb/bin/python
 import sys
 import os
 import time

tools/itb/bin/itb_screenshot

@@ -1,4 +1,4 @@
-#!/usr/bin/env python3
+#!/root/venvs/itb/bin/python
 import sys
 import os
 import argparse

tools/itb/bin/itb_scroll

@@ -1,4 +1,4 @@
-#!/usr/bin/env python3
+#!/root/venvs/itb/bin/python
 import sys
 import os
 import time

tools/itb/bin/itb_start

@@ -1,4 +1,4 @@
-#!/usr/bin/env python3
+#!/root/venvs/itb/bin/python
 import random
 import subprocess
 import sys

tools/itb/setup.py

@@ -18,14 +18,17 @@ setup(
         'selenium>=4.0.0',
         'webdriver-manager>=3.8.0',
         'click>=8.0.0',
-        'beautifulsoup4>=4.9.0'
+        'beautifulsoup4>=4.9.0',
+        'pytest>=7.0.0',
+        'pytest-cov>=4.0.0',
+        'black>=22.0.0',
+        'flake8>=4.0.0'
     ],
-    extras_require={
+    classifiers=[
-        'dev': [
+        'Development Status :: 3 - Alpha',
-            'pytest>=7.0.0',
+        'Intended Audience :: Developers',
-            'pytest-cov>=4.0.0',
+        'Programming Language :: Python :: 3',
-            'black>=22.0.0',
+        'Programming Language :: Python :: 3.10',
-            'flake8>=4.0.0'
+    ],
-        ]
+    python_requires='>=3.10',
-    }
+)
-)

tools/train/bin/train_deepseek

@@ -0,0 +1,10 @@
+#!/root/venvs/train/bin/python
+"""
+Command-line utility for fine-tuning DeepSeek models using Unsloth.
+Always trains from a base model to create a new fine-tuned model.
+"""
+import sys
+from train.unsloth_deepseek import main
+
+if __name__ == "__main__":
+    sys.exit(main())

tools/train/bin/train_mistral

@@ -0,0 +1,9 @@
+#!/root/venvs/train/bin/python
+"""
+Command-line utility for fine-tuning Mistral models using Mistral API.
+"""
+import sys
+from train.mistral_api import main
+
+if __name__ == "__main__":
+    sys.exit(main())

tools/train/readme.md

@@ -0,0 +1,68 @@
+# 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/requirements.txt

@@ -0,0 +1,13 @@
+pyyaml>=6.0
+requests>=2.28.0
+torch>=2.0.0
+transformers>=4.30.0
+# DeepSeek support
+accelerate>=0.25.0
+bitsandbytes>=0.41.1
+einops>=0.7.0
+sentencepiece>=0.1.99
+unsloth>=2024.3
+trl>=0.7.8
+datasets>=2.14.6
+peft>=0.8.0

tools/train/setup.py

@@ -0,0 +1,36 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="train",
+    version="0.1.0",
+    packages=find_packages(),
+    scripts=[
+        'bin/train_deepseek',
+        'bin/train_mistral'
+    ],
+    install_requires=[
+        'pyyaml>=6.0',
+        'requests>=2.28.0',
+        'torch>=2.0.0',
+        'transformers>=4.30.0',
+        'accelerate>=0.25.0',
+        'bitsandbytes>=0.41.1',
+        'einops>=0.7.0',
+        'sentencepiece>=0.1.99',
+        'unsloth>=2024.3',
+        'trl>=0.7.8',
+        'datasets>=2.14.6',
+        'peft>=0.8.0',
+        'pytest>=7.0.0',
+        'pytest-cov>=4.0.0',
+        'black>=22.0.0',
+        'flake8>=4.0.0'
+    ],
+    classifiers=[
+        'Development Status :: 3 - Alpha',
+        'Intended Audience :: Developers',
+        'Programming Language :: Python :: 3',
+        'Programming Language :: Python :: 3.10',
+    ],
+    python_requires='>=3.10',
+)

tools/train/train.sh

@@ -0,0 +1,15 @@
+#!/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"
+
+train_deepseek --output-dir "$OUTPUT_DIR"

tools/train/train/__init__.py

@@ -0,0 +1,8 @@
+"""
+SIA Training Tool
+
+This package provides utilities for fine-tuning language models used by SIA.
+Supports DeepSeek and Mistral models.
+"""
+
+__version__ = "0.1.0"

tools/train/train/mistral_api.py

@@ -0,0 +1,141 @@
+#!/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
+
+def upload_file(api_key: str, file_path: Path) -> str:
+    """Upload a file to the Mistral API and return the file ID"""
+    url = "https://api.mistral.ai/v1/files"
+    headers = {
+        "Authorization": f"Bearer {api_key}"
+    }
+    files = {
+        "file": ("dataset.jsonl", open(file_path, "rb"), "application/jsonl"),
+        "purpose": (None, "fine-tune")
+    }
+    
+    response = requests.post(url, headers=headers, files=files)
+    if response.status_code != 200:
+        print(f"Error uploading file: {response.text}")
+        sys.exit(1)
+    
+    return response.json()["id"]
+
+def start_finetune_job(api_key: str, model: str, file_id: str, params: sia_train_lib.TrainingParams):
+    """Start a fine-tuning job on the Mistral API"""
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json"
+    }
+    data = {
+        "model": model,
+        "training_files": [{"file_id": file_id, "weight": 1}],
+        "hyperparameters": {
+            "learning_rate": params.learning_rate,
+            "epochs": params.epochs
+        }
+    }
+    
+    response = requests.post(
+        "https://api.mistral.ai/v1/fine_tuning/jobs",
+        headers=headers,
+        json=data
+    )
+    
+    if response.status_code != 200:
+        print(f"Error creating fine-tuning job: {response.text}")
+        return None
+        
+    return response.json()["id"]
+
+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)
+    
+    if not training_data:
+        print("No valid training data found. Exiting.")
+        return 1
+    
+    model_name = f"sia_{commit_hash}"
+    
+    # Create temp file and upload
+    with tempfile.NamedTemporaryFile(mode='w', suffix='.jsonl', delete=False) as f:
+        for sample in training_data:
+            json.dump(sample, f, ensure_ascii=False)
+            f.write('\n')
+    
+    try:
+        file_id = upload_file(config.api_key, Path(f.name))
+        
+        # Start fine-tuning job
+        job_id = start_finetune_job(
+            api_key=config.api_key,
+            model=config.model,
+            file_id=file_id,
+            params=train_params
+        )
+        
+        if not job_id:
+            return 1
+            
+        print(f"Started fine-tuning job: {model_name}")
+        print(f"Job ID: {job_id}")
+        print(f"Check status: curl -H 'Authorization: Bearer {config.api_key}' https://api.mistral.ai/v1/fine_tuning/jobs/{job_id}")
+    finally:
+        os.unlink(f.name)
+    
+    return 0
+
+if __name__ == "__main__":
+    exit(main())

tools/train/train/unsloth_deepseek.py

@@ -0,0 +1,239 @@
+#!/root/venvs/train/bin/python
+"""
+Script for fine-tuning DeepSeek models for SIA using Unsloth.
+Training always starts from a base model and creates a new fine-tuned model.
+"""
+import argparse
+import os
+import sys
+import torch
+from dataclasses import dataclass
+from pathlib import Path
+import json
+
+# Import from shared library
+from .util import prepare_training_data
+
+@dataclass
+class Config:
+    def __init__(self):
+        parser = argparse.ArgumentParser(description='Train SIA model using Unsloth')
+        parser.add_argument(
+            '--config',
+            type=Path,
+            default=Path('/root/sia/training/config.yaml'),
+            help='Path to config file'
+        )
+        parser.add_argument(
+            '--base-model',
+            type=str,
+            default='deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B',
+            help='HuggingFace model ID for base model'
+        )
+        parser.add_argument(
+            '--output-dir',
+            type=Path,
+            required=True,
+            help='Directory to save the trained model'
+        )
+        parser.add_argument(
+            '--api-key',
+            type=str,
+            default=os.environ.get('SIA_HF_API_KEY'),
+            help='HuggingFace API key'
+        )
+        self.args = parser.parse_args()
+
+    @property
+    def config_path(self) -> Path:
+        return self.args.config
+
+    @property
+    def base_model(self) -> str:
+        return self.args.base_model
+    
+    @property
+    def output_dir(self) -> Path:
+        return self.args.output_dir
+    
+    @property
+    def api_key(self) -> str:
+        return self.args.api_key
+
+def train_model(config: Config, training_data, train_params, commit_hash):
+    """Train the model using Unsloth"""
+    try:
+        from unsloth import FastLanguageModel
+        from transformers import TrainingArguments, DataCollatorForSeq2Seq
+        from trl import SFTTrainer
+        from datasets import Dataset
+        from unsloth.chat_templates import get_chat_template, train_on_responses_only
+    except ImportError as e:
+        print(f"Error importing required libraries: {e}")
+        print("Please ensure Unsloth and its dependencies are installed.")
+        sys.exit(1)
+
+    print(f"Starting training from base model: {config.base_model}")
+    
+    # Convert to datasets format
+    dataset = Dataset.from_list(training_data)
+    
+    # Determine if bfloat16 is supported
+    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+    
+    # Load the model - always from a base model (no incremental updates)
+    try:
+        model, tokenizer = FastLanguageModel.from_pretrained(
+            model_name=config.base_model,
+            max_seq_length=2048,
+            dtype=dtype,
+            load_in_4bit=True,
+            token=config.api_key,
+        )
+    except Exception as e:
+        print(f"Error loading base model: {e}")
+        sys.exit(1)
+    
+    # Apply LoRA
+    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
+    model = FastLanguageModel.get_peft_model(
+        model,
+        r=16,
+        target_modules=target_modules,
+        lora_alpha=16,
+        lora_dropout=0,
+        bias="none",
+        use_gradient_checkpointing="unsloth",
+        random_state=3407,
+    )
+    
+    # Apply chat template
+    tokenizer = get_chat_template(
+        tokenizer, 
+        chat_template="llama-3.1",  # Compatible with DeepSeek
+    )
+    
+    # Function to format conversations
+    def formatting_prompts_func(examples):
+        convos = examples["conversations"]
+        texts = [tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=False) for convo in convos]
+        return {"text": texts}
+    
+    # Standarize dataset and format
+    from unsloth.chat_templates import standardize_sharegpt
+    
+    # Add conversations field if not present
+    if "conversations" not in dataset.column_names:
+        if "messages" in dataset.column_names:
+            dataset = dataset.rename_column("messages", "conversations")
+        else:
+            dataset = dataset.map(lambda x: {"conversations": [{"role": "system", "content": x.get("system_prompt", "")}, 
+                                                            {"role": "user", "content": x.get("prompt", "")}, 
+                                                            {"role": "assistant", "content": x.get("response", "")}]})
+    
+    # Standardize format
+    dataset = standardize_sharegpt(dataset)
+    
+    # Apply formatting
+    dataset = dataset.map(formatting_prompts_func, batched=True)
+    
+    # Configure the trainer
+    output_dir = config.output_dir / commit_hash
+    output_dir.mkdir(parents=True, exist_ok=True)
+    
+    # Determine steps or epochs based on dataset size
+    max_steps = None
+    num_train_epochs = train_params.epochs
+    if len(dataset) < 100:  # Small dataset
+        # Aim for at least 500 steps for small datasets
+        max_steps = 500
+        num_train_epochs = None
+    
+    trainer = SFTTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        train_dataset=dataset,
+        dataset_text_field="text",
+        max_seq_length=2048,
+        data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer),
+        dataset_num_proc=2,
+        packing=False,
+        args=TrainingArguments(
+            per_device_train_batch_size=2,
+            gradient_accumulation_steps=4,
+            warmup_steps=5,
+            max_steps=max_steps,
+            num_train_epochs=num_train_epochs,
+            learning_rate=train_params.learning_rate,
+            fp16=not torch.cuda.is_bf16_supported(),
+            bf16=torch.cuda.is_bf16_supported(),
+            logging_steps=10,
+            optim="adamw_8bit",
+            weight_decay=0.01,
+            lr_scheduler_type="linear",
+            seed=3407,
+            output_dir=str(output_dir),
+            report_to="none",
+        ),
+    )
+    
+    # Train only on responses
+    trainer = train_on_responses_only(
+        trainer,
+        instruction_part="<|start_header_id|>user<|end_header_id|>\n\n",
+        response_part="<|start_header_id|>assistant<|end_header_id|>\n\n",
+    )
+    
+    # Train the model
+    trainer.train()
+    
+    # Enable inference mode for the model
+    model = FastLanguageModel.for_inference(model)
+    
+    # Save the model
+    model.save_pretrained(output_dir)
+    tokenizer.save_pretrained(output_dir)
+    
+    # Create a metadata file with training information
+    with open(output_dir / "training_info.json", "w") as f:
+        json.dump({
+            "base_model": config.base_model,
+            "commit_hash": commit_hash,
+            "learning_rate": train_params.learning_rate,
+            "epochs": train_params.epochs,
+            "dataset_size": len(dataset),
+            "training_method": "unsloth",
+        }, f, indent=2)
+    
+    return output_dir
+
+def main():
+    config = Config()
+    
+    # Prepare training data
+    training_data, train_params, commit_hash = prepare_training_data(config.config_path)
+    
+    if not training_data:
+        print("No valid training data found. Exiting.")
+        return 1
+    
+    # Train the model
+    try:
+        model_dir = train_model(config, training_data, train_params, commit_hash)
+        
+        # Create symlink to current
+        current_link = config.output_dir / "current"
+        if current_link.exists() or current_link.is_symlink():
+            current_link.unlink()
+        os.symlink(model_dir, current_link, target_is_directory=True)
+        
+        print(f"Training complete. Model saved to {model_dir}")
+        print(f"Symlink created at {current_link}")
+        
+        return 0
+    except Exception as e:
+        print(f"Error during training: {e}")
+        return 1
+
+if __name__ == "__main__":
+    exit(main())

tools/train/train/util.py

@@ -0,0 +1,186 @@
+"""
+Shared library for SIA model training functionality.
+Contains common code for both API-based and local training.
+"""
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List, Optional, Set, Tuple
+import hashlib
+import json
+import subprocess
+import sys
+import xml.etree.ElementTree as ET
+import yaml
+
+@dataclass
+class TrainingParams:
+    """Parameters for model training"""
+    learning_rate: float
+    epochs: int
+    batch_size: int = 1
+
+class DatasetCreator:
+    """Creates training datasets from XML iteration files"""
+    
+    def __init__(
+            self,
+            xml_files: Set[Path],
+            system_prompt_file: Path,
+            action_schema_file: Path
+        ):
+        self.xml_files = xml_files
+        self.system_prompt_file = Path(system_prompt_file)
+        self.action_schema_file = Path(action_schema_file)
+        
+        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 _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) -> Optional[Dict]:
+        """Parse a single iteration XML file into a training example"""
+        try:
+            tree = ET.parse(file_path)
+            root = tree.getroot()
+            
+            # Check hashes to ensure compatibility
+            if root.get('system_prompt_hash') != self.system_prompt_hash:
+                print(f"System prompt hash mismatch in {file_path}")
+                return None
+            if root.get('action_schema_hash') != self.action_schema_hash:
+                print(f"Action schema hash mismatch in {file_path}")
+                return None
+            
+            context_elem = root.find('context')
+            response_elem = root.find('response')
+            
+            if context_elem is None or response_elem is None:
+                print(f"Missing context or response elements in {file_path}")
+                return None
+                
+            context = context_elem.text
+            response = response_elem.text
+            
+            if not context or not response:
+                print(f"Empty context or response in {file_path}")
+                return None
+            
+            return {
+                "messages": [
+                    {
+                        "role": "system",
+                        "content": self.system_prompt + "\n" + self.action_schema
+                    },
+                    {
+                        "role": "user",
+                        "content": context
+                    },
+                    {
+                        "role": "assistant", 
+                        "content": response
+                    }
+                ]
+            }
+            
+        except Exception as e:
+            print(f"Error processing {file_path}: {str(e)}")
+            return None
+    
+    def create_dataset(self) -> List[Dict]:
+        """Create a dataset from all valid XML files"""
+        samples = []
+        total_files = len(self.xml_files)
+        print(f"Processing {total_files} XML files...")
+        
+        for i, xml_file in enumerate(sorted(self.xml_files)):
+            if i % 10 == 0:
+                print(f"Processed {i}/{total_files} files...")
+                
+            sample = self._parse_iteration_file(xml_file)
+            if sample:
+                samples.append(sample)
+        
+        print(f"Created dataset with {len(samples)} samples from {total_files} files")
+        return samples
+
+def find_xml_files(data_paths: List[Path]) -> Set[Path]:
+    """Find all XML files in the given data paths"""
+    xml_files = set()
+    for path in data_paths:
+        if not path.exists():
+            print(f"Error: Data path not found: {path}")
+            sys.exit(1)
+        xml_files.update(path.rglob('*.xml'))
+    return xml_files
+
+def format_chat_for_mistral(messages):
+    """Format messages for Mistral chat format"""
+    # Mistral uses a specific chat format:
+    # <s>[INST] {system + user content} [/INST] {assistant response} </s>
+    
+    system_content = ""
+    user_content = ""
+    assistant_content = ""
+    
+    for msg in messages:
+        role = msg["role"]
+        content = msg["content"]
+        
+        if role == "system":
+            system_content = content
+        elif role == "user":
+            user_content = content
+        elif role == "assistant":
+            assistant_content = content
+    
+    # Combine system and user content for the instruction
+    instruction = system_content
+    if instruction and user_content:
+        instruction += "\n\n"
+    instruction += user_content
+    
+    # Format according to Mistral chat template
+    return f"<s>[INST] {instruction} [/INST] {assistant_content} </s>"
+
+def prepare_training_data(config_path: Path) -> Tuple[List[Dict], TrainingParams, str]:
+    """Prepare training data from config and XML files"""
+    with open(config_path) as f:
+        config_data = yaml.safe_load(f)
+    
+    data_paths = [Path(p) for p in config_data['data']]
+    xml_files = find_xml_files(data_paths)
+    
+    paths = list(xml_files)
+    paths.append(config_path)
+    paths.append(Path(config_data['model']['system_prompt_path']))
+    paths.append(Path(config_data['model']['action_schema']))
+    commit_hash = check_git_status(paths)
+    
+    creator = DatasetCreator(
+        xml_files=xml_files,
+        system_prompt_file=config_data['model']['system_prompt_path'],
+        action_schema_file=config_data['model']['action_schema']
+    )
+    
+    training_data = creator.create_dataset()
+    
+    train_params = TrainingParams(
+        learning_rate=config_data['params'].get('learning_rate', 1e-5),
+        epochs=config_data['params'].get('epochs', 3),
+        batch_size=config_data['params'].get('batch_size', 1)
+    )
+    
+    return training_data, train_params, commit_hash
+
+def save_jsonl_dataset(data: List[Dict], output_path: Path) -> None:
+    """Save dataset in JSONL format"""
+    with open(output_path, 'w', encoding='utf-8') as f:
+        for sample in data:
+            json.dump(sample, f, ensure_ascii=False)
+            f.write('\n')
+    print(f"Saved dataset with {len(data)} samples to {output_path}")

tools/train/train_mistral.py

@@ -1,260 +0,0 @@
-#!/usr/bin/env python3
-from dataclasses import dataclass
-from datetime import datetime
-from dotenv import load_dotenv
-from pathlib import Path
-from typing import Dict, List, Optional, Set
-import argparse
-import hashlib
-import json
-import os
-import requests
-import subprocess
-import sys
-import tempfile
-import xml.etree.ElementTree as ET
-import yaml
-
-@dataclass
-class Config:
-    def __init__(self):
-        load_dotenv()
-        parser = argparse.ArgumentParser(description='Train SIA model using Mistral API')
-        parser.add_argument(
-            '--config',
-            type=Path,
-            default=os.getenv('SIA_TRAINING_CONFIG', '/root/sia/training/config.yaml'),
-            help='Path to config file'
-        )
-        parser.add_argument(
-            '--model',
-            type=str,
-            default=os.getenv('SIA_MISTRAL_MODEL', 'mistral-large-latest'),
-            help='Base model for fine-tuning'
-        )
-        parser.add_argument(
-            '--api-key',
-            type=str,
-            default=os.getenv('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
-
-class FinetuneDatasetCreator:
-    def __init__(
-            self,
-            xml_files: Set[Path],
-            system_prompt_file: Path,
-            action_schema_file: Path,
-            output_file: Path
-        ):
-        self.xml_files = xml_files
-        self.system_prompt_file = Path(system_prompt_file)
-        self.action_schema_file = Path(action_schema_file)
-        self.output_file = Path(output_file)
-        
-        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 _calculate_hash(self, content: str) -> str:
-        return hashlib.sha256(content.encode()).hexdigest()
-    
-    def _parse_iteration_file(self, file_path: Path) -> Optional[Dict]:
-        try:
-            tree = ET.parse(file_path)
-            root = tree.getroot()
-            
-            if root.get('system_prompt_hash') != self.system_prompt_hash:
-                print(f"System prompt hash mismatch in {file_path}")
-                return None
-            if root.get('action_schema_hash') != self.action_schema_hash:
-                print(f"Action schema hash mismatch in {file_path}")
-                return None
-            
-            context = root.find('context').text
-            response = root.find('response').text
-            
-            if not context or not response:
-                print(f"Missing context or response in {file_path}")
-                return None
-            
-            return {
-                "messages": [
-                    {
-                        "role": "system",
-                        "content": self.system_prompt + self.action_schema
-                    },
-                    {
-                        "role": "user",
-                        "content": context
-                    },
-                    {
-                        "role": "assistant", 
-                        "content": response
-                    }
-                ]
-            }
-            
-        except Exception as e:
-            print(f"Error processing {file_path}: {str(e)}")
-            return None
-    
-    def create_dataset(self) -> int:
-        sample_count = 0
-        self.output_file.parent.mkdir(parents=True, exist_ok=True)
-        
-        with open(self.output_file, 'w', encoding='utf-8') as f:
-            for xml_file in sorted(self.xml_files):
-                sample = self._parse_iteration_file(xml_file)
-                if sample:
-                    json.dump(sample, f, ensure_ascii=False)
-                    f.write('\n')
-                    sample_count += 1
-                    
-        print(f"Created dataset with {sample_count} samples at {self.output_file}")
-        return sample_count
-
-def find_xml_files(data_paths: List[Path]) -> Set[Path]:
-    xml_files = set()
-    for path in data_paths:
-        if not path.exists():
-            print(f"Error: Data path not found: {path}")
-            sys.exit(1)
-        xml_files.update(path.rglob('*.xml'))
-    return xml_files
-
-def check_git_status(paths: list[Path]) -> str:
-    try:
-        for path in paths:
-            result = subprocess.run(['git', 'status', '--porcelain', str(path)], 
-                                 capture_output=True, text=True)
-            if result.stdout.strip():
-                print(f"Error: Uncommitted changes in {path}")
-                print(result.stdout)
-                sys.exit(1)
-        
-        result = subprocess.run(['git', 'rev-parse', 'HEAD'],
-                              capture_output=True, text=True)
-        return result.stdout.strip()
-    except subprocess.CalledProcessError as e:
-        print(f"Git command failed: {e}")
-        sys.exit(1)
-
-def create_combined_dataset(xml_files: Set[Path], config_data: dict, tmp_dir: Path) -> list:
-    tmp_file = tmp_dir / "dataset.jsonl"
-    creator = FinetuneDatasetCreator(
-        xml_files=xml_files,
-        system_prompt_file=config_data['model']['system_prompt_path'],
-        action_schema_file=config_data['model']['action_schema'],
-        output_file=tmp_file
-    )
-    creator.create_dataset()
-    
-    with open(tmp_file) as f:
-        return [json.loads(line) for line in f]
-
-def prepare_training_data(config: Config) -> tuple[list, dict, str]:
-    with open(config.config_path) as f:
-        config_data = yaml.safe_load(f)
-    
-    data_paths = [Path(p) for p in config_data['data']]
-    xml_files = find_xml_files(data_paths)
-    
-    paths = list(xml_files)
-    paths.append(config.config_path)
-    paths.append(Path(config_data['model']['system_prompt_path']))
-    paths.append(Path(config_data['model']['action_schema']))
-    commit_hash = check_git_status(paths)
-    
-    with tempfile.TemporaryDirectory() as tmp_dir:
-        training_data = create_combined_dataset(xml_files, config_data, Path(tmp_dir))
-    
-    train_params = {
-        'learning_rate': config_data['params']['learning_rate'],
-        'epochs': config_data['params']['epochs']
-    }
-    
-    return training_data, train_params, commit_hash
-
-def upload_file(api_key: str, file_path: Path) -> str:
-    url = "https://api.mistral.ai/v1/files"
-    headers = {
-        "Authorization": f"Bearer {api_key}"
-    }
-    files = {
-        "file": ("dataset.jsonl", open(file_path, "rb"), "application/jsonl"),
-        "purpose": (None, "fine-tune")
-    }
-    
-    response = requests.post(url, headers=headers, files=files)
-    if response.status_code != 200:
-        print(f"Error uploading file: {response.text}")
-        sys.exit(1)
-    
-    return response.json()["id"]
-
-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 = prepare_training_data(config)
-    model_name = f"sia_{commit_hash}"
-    
-    # Create temp file and upload
-    with tempfile.NamedTemporaryFile(mode='w', suffix='.jsonl', delete=False) as f:
-        for sample in training_data:
-            json.dump(sample, f)
-            f.write('\n')
-    
-    try:
-        file_id = upload_file(config.api_key, Path(f.name))
-        
-        # Create fine-tuning job
-        headers = {
-            "Authorization": f"Bearer {config.api_key}",
-            "Content-Type": "application/json"
-        }
-        data = {
-            "model": config.model,
-            "training_files": [{"file_id": file_id, "weight": 1}],
-            "hyperparameters": train_params
-        }
-        
-        response = requests.post(
-            "https://api.mistral.ai/v1/fine_tuning/jobs",
-            headers=headers,
-            json=data
-        )
-        
-        if response.status_code != 200:
-            print(f"Error creating fine-tuning job: {response.text}")
-            return 1
-            
-        job_id = response.json()["id"]
-        print(f"Started fine-tuning job: {model_name}")
-        print(f"Job ID: {job_id}")
-        print(f"Check status: curl -H 'Authorization: Bearer {config.api_key}' https://api.mistral.ai/v1/fine_tuning/jobs/{job_id}")
-    finally:
-        os.unlink(f.name)
-    
-    return 0
-
-if __name__ == "__main__":
-    exit(main())

training/config.yaml

@@ -5,6 +5,6 @@ params:
   learning_rate: 1e-5
   epochs: 3
 data:
-  - "training/clean_start/"
+  - "/root/sia/training/clean_start/"
-  - "training/delete_indicated_entries/"
+  - "/root/sia/training/delete_indicated_entries/"
-  - "training/list_entries_to_delete/"
+  - "/root/sia/training/list_entries_to_delete/"