diff --git a/.dockerignore b/.dockerignore
new file mode 100644
index 0000000..e3f10e1
--- /dev/null
+++ b/.dockerignore
@@ -0,0 +1 @@
+./model/
\ No newline at end of file
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..f8ef954
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,2 @@
+pdf/
+model/
\ No newline at end of file
diff --git a/Dockerfile b/Dockerfile
new file mode 100644
index 0000000..5c36051
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/readme.md b/readme.md
index 1f45c3e..bf1cacd 100644
--- a/readme.md
+++ b/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.
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..4803a9a
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,2 @@
+transformers
+torch
\ No newline at end of file
diff --git a/sia/llm_engine.py b/sia/llm_engine.py
new file mode 100644
index 0000000..0d42d7d
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/test.sh b/test.sh
new file mode 100644
index 0000000..5cb5fe9
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/tests/test_llm_engine.py b/tests/test_llm_engine.py
new file mode 100644
index 0000000..e56c77e
--- /dev/null
+++ b/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
+ -
+ - The original request
+ -
+ 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 = """
+
+
+
+
+
+
+
+
+
+
+
+
+ """
+
+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"{main_context}")
+
+'''
+ 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()
\ No newline at end of file