Start work on llm_engine

.dockerignore

@@ -0,0 +1 @@
+./model/

.gitignore

@@ -0,0 +1,2 @@
+pdf/
+model/

Dockerfile

@@ -0,0 +1,12 @@
+FROM huggingface/transformers-pytorch-gpu AS requirements
+WORKDIR /root
+COPY requirements.txt /root/requirements.txt
+COPY ./sia/ /root/sia/
+RUN pip3 install -r requirements.txt
+
+FROM requirements AS test
+COPY ./tests/ /root/tests/
+RUN mkdir -p /root/model
+CMD python3 -m unittest discover tests
+
+FROM requirements

readme.md

@@ -224,6 +224,21 @@ The web interface takes over standard input and output.
 It each time the LLM generates a response, the web interface will display it.
 The user can modify the response before the actions are executed.
 
+### Project structure
+
+The SIA application is developed in the src directory.
+The tests directory contains unit tests, mock objects and integration tests.
+The model directory contains the trained model.
+It is excluded from the git repository and the docker context because it is too large.
+
+The docker file has a separate stage for testing.
+The `test.sh` script builds this stage, runs the tests and removes the test image.
+
+To use SIA several directories have to be mounted:
+- `/root/model': The model directory
+- `/root/sia_repo': The git repository
+- the docker socket: to run sub-SIA instances
+
 ## Actions
 
 A list of all available core actions.

requirements.txt

@@ -0,0 +1,2 @@
+transformers
+torch

sia/llm_engine.py

@@ -0,0 +1,72 @@
+from typing import NamedTuple
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextStreamer
+import torch
+
+class InferenceResult(NamedTuple):
+    reasoning: str
+    actions: str
+
+class LlmEngine:
+    def __init__(self, model_path: str):
+        """
+        Initialize the LLM Engine with a model path.
+
+        Args:
+            model_path: Path to the model weights to be used.
+        """
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+        print(f"device: {self.device}")
+        self.set_model_path(model_path)
+
+    def set_model_path(self, model_path: str):
+        """
+        Load the model from the specified path.
+
+        Args:
+            model_path: Path to the model weights to load.
+        """
+        tokenizer = AutoTokenizer.from_pretrained(model_path)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            return_dict=True,
+            low_cpu_mem_usage=True,
+            torch_dtype=self.torch_dtype,
+            device_map="auto",
+            trust_remote_code=True,
+        ).to(self.device)
+        if tokenizer.pad_token_id is None:
+            tokenizer.pad_token_id = 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=tokenizer,
+            torch_dtype=torch.float16,
+            device_map="auto",
+        )
+
+    def infer(self, system_prompt: str, main_context: str, action_schema: str) -> InferenceResult:
+        """
+        Run inference using the system prompt and main context, while validating actions against the provided XML schema.
+
+        Args:
+            system_prompt: The system prompt string
+            main_context: The main context string after templating
+            action_schema: XML schema to validate the generated actions
+
+        Returns:
+            InferenceResult: the actions validate against the schema
+        """
+        pass
+
+    def finetune(self, dataset_paths: list, output_dir: str):
+        """
+        Fine-tune the model with new datasets and save the updated model weights.
+
+        Args:
+            dataset_paths: List of paths to datasets for fine-tuning.
+            output_dir: Directory where the updated model weights will be saved.
+        """
+        pass

test.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+
+# Continue on error
+set -e
+
+# Build with progress output and capture the tag
+TAG=$( \
+docker build \
+    --target test \
+    . \
+    2>&1 | tee /dev/tty | grep "writing image" | cut -d' ' -f4 \
+)
+
+# Exit if tag is empty
+[ -z "$TAG" ] && exit 1
+
+# Run tests
+docker run \
+    --rm \
+    --gpus=all \
+    -v /$(pwd)/model/:/root/model/ \
+    $TAG
+
+# Clean up image
+[ ! -z "$TAG" ] && docker rmi $TAG

tests/test_llm_engine.py

@@ -0,0 +1,66 @@
+import unittest
+from sia.llm_engine import LlmEngine, InferenceResult
+
+echo_system_prompt = """
+    Your answer always consists of 4 parts:
+        - The original request
+        - <test_tag>
+        - The original request
+        - </test_tag>
+        You never provide an answer.
+        Only the entered text, the xml open tag, the same text again and the xml close tag.
+        Don't add whitespace or newlines.
+        Don't modify the text or change casing.
+        Be exact, this is for testing purposes.
+    """
+
+echo_action_schema = """
+    <?xml version="1.0" encoding="UTF-8"?>
+    <xs:schema xmlns:xs="http://www.w3.org/2001/XMLSchema">
+    <xs:element name="test_tag">
+        <xs:complexType>
+        <xs:simpleContent>
+            <xs:extension base="xs:string">
+            <xs:attribute name="id" type="xs:string" use="required"/>
+            </xs:extension>
+        </xs:simpleContent>
+        </xs:complexType>
+    </xs:element>
+    </xs:schema>
+    """
+
+class TestLlmEngine(unittest.TestCase):
+    def setUp(self):
+        self.model_path = "/root/model"
+
+    def test_initialization(self):
+        llm_engine = LlmEngine(self.model_path)
+        self.assertIsInstance(llm_engine, LlmEngine)
+
+    def test_infer(self):
+        main_context = "This is a test"
+        llm_engine = LlmEngine(self.model_path)
+        result = llm_engine.infer(echo_system_prompt, main_context, echo_action_schema)
+        self.assertIsInstance(result, InferenceResult)
+        self.assertIsInstance(result.reasoning, str)
+        self.assertIsInstance(result.actions, str)
+        self.assertEqual(result.reasoning, main_context)
+        self.assertEqual(result.actions, f"<test_tag>{main_context}</test_tag>")
+
+'''
+    def test_set_model_path(self):
+        new_model_path = "/path/to/new/model"
+        self.llm_engine.set_model_path(new_model_path)
+        # Add assertions to check if the model path was updated correctly
+
+    def test_finetune(self):
+        dataset_paths = ["/path/to/dataset1", "/path/to/dataset2"]
+        output_dir = "/path/to/output"
+
+        self.llm_engine.finetune(dataset_paths, output_dir)
+        # Add assertions to check if the fine-tuning process completed successfully
+        # For example, check if new model weights were saved in the output directory
+'''
+
+if __name__ == '__main__':
+    unittest.main()