From 006db518f27fe5291fe5bca17ab365b66d31c795 Mon Sep 17 00:00:00 2001 From: Geens Date: Sun, 20 Oct 2024 19:05:47 +0200 Subject: [PATCH] Start work on llm_engine --- .dockerignore | 1 + .gitignore | 2 ++ Dockerfile | 12 +++++++ readme.md | 15 +++++++++ requirements.txt | 2 ++ sia/llm_engine.py | 72 ++++++++++++++++++++++++++++++++++++++++ test.sh | 25 ++++++++++++++ tests/test_llm_engine.py | 66 ++++++++++++++++++++++++++++++++++++ 8 files changed, 195 insertions(+) create mode 100644 .dockerignore create mode 100644 .gitignore create mode 100644 Dockerfile create mode 100644 requirements.txt create mode 100644 sia/llm_engine.py create mode 100644 test.sh create mode 100644 tests/test_llm_engine.py 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