From 20429616f7eea926a5e6a6a7e0f3b1f738895ef2 Mon Sep 17 00:00:00 2001 From: geens Date: Mon, 21 Oct 2024 22:23:48 +0200 Subject: [PATCH] Begin implementation, basic inferenece --- Dockerfile | 7 +++-- diagrams/render.sh | 0 sia/__init__.py | 0 sia/__main__.py | 5 +++ sia/inference_result.py | 5 +++ sia/llm_engine.py | 41 ++++++++++++++----------- sia/util.py | 47 ++++++++++++++++++++++++++++ test.sh | 0 test/__init__.py | 0 test/llm_engine_test.py | 25 +++++++++++++++ test/test_data.py | 32 +++++++++++++++++++ test/test_util.py | 8 +++++ test/util_test.py | 17 +++++++++++ tests/test_llm_engine.py | 66 ---------------------------------------- 14 files changed, 167 insertions(+), 86 deletions(-) mode change 100644 => 100755 diagrams/render.sh create mode 100644 sia/__init__.py create mode 100644 sia/__main__.py create mode 100644 sia/inference_result.py create mode 100644 sia/util.py mode change 100644 => 100755 test.sh create mode 100644 test/__init__.py create mode 100644 test/llm_engine_test.py create mode 100644 test/test_data.py create mode 100644 test/test_util.py create mode 100644 test/util_test.py delete mode 100644 tests/test_llm_engine.py diff --git a/Dockerfile b/Dockerfile index 5c36051..513c3bf 100644 --- a/Dockerfile +++ b/Dockerfile @@ -5,8 +5,9 @@ COPY ./sia/ /root/sia/ RUN pip3 install -r requirements.txt FROM requirements AS test -COPY ./tests/ /root/tests/ +COPY ./test/ /root/test/ RUN mkdir -p /root/model -CMD python3 -m unittest discover tests +CMD ["python3", "-m", "unittest", "discover", "-p", "*test.py", "-v"] -FROM requirements \ No newline at end of file +FROM requirements +CMD ["python3", "-m", "sia"] \ No newline at end of file diff --git a/diagrams/render.sh b/diagrams/render.sh old mode 100644 new mode 100755 diff --git a/sia/__init__.py b/sia/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/sia/__main__.py b/sia/__main__.py new file mode 100644 index 0000000..5624a5f --- /dev/null +++ b/sia/__main__.py @@ -0,0 +1,5 @@ +def main(): + print("Hello, World! --sia") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/sia/inference_result.py b/sia/inference_result.py new file mode 100644 index 0000000..92b3796 --- /dev/null +++ b/sia/inference_result.py @@ -0,0 +1,5 @@ +from typing import NamedTuple + +class InferenceResult(NamedTuple): + reasoning: str + actions: str \ No newline at end of file diff --git a/sia/llm_engine.py b/sia/llm_engine.py index 0d42d7d..ad7d87e 100644 --- a/sia/llm_engine.py +++ b/sia/llm_engine.py @@ -1,10 +1,8 @@ -from typing import NamedTuple -from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextStreamer +from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline import torch -class InferenceResult(NamedTuple): - reasoning: str - actions: str +from . import util +from .inference_result import InferenceResult class LlmEngine: def __init__(self, model_path: str): @@ -14,9 +12,6 @@ class LlmEngine: 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): @@ -26,24 +21,24 @@ class LlmEngine: Args: model_path: Path to the model weights to load. """ - tokenizer = AutoTokenizer.from_pretrained(model_path) + self.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, + torch_dtype=torch.bfloat16, 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 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=tokenizer, - torch_dtype=torch.float16, + tokenizer=self.tokenizer, + torch_dtype=torch.bfloat16, device_map="auto", ) @@ -57,9 +52,21 @@ class LlmEngine: action_schema: XML schema to validate the generated actions Returns: - InferenceResult: the actions validate against the schema + InferenceResult: Tuple containing reasoning and actions that validate against the schema """ - pass + valid_elements = util.get_valid_root_elements(action_schema) + messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": main_context} + ] + prompt = self.tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + outputs = self.pipeline(prompt, max_new_tokens=120, do_sample=True) + generated_text = outputs[0]["generated_text"] + response = generated_text.split("<|start_header_id|>assistant<|end_header_id|>",1)[1].strip() + result = util.split_response(response, valid_elements) + return result def finetune(self, dataset_paths: list, output_dir: str): """ diff --git a/sia/util.py b/sia/util.py new file mode 100644 index 0000000..0704f16 --- /dev/null +++ b/sia/util.py @@ -0,0 +1,47 @@ +from .inference_result import InferenceResult + +import xml.etree.ElementTree as ET +import re + +def get_valid_root_elements(schema: str) -> set: + """ + Extract valid root element names from the XML schema. + + Args: + schema: XML schema string + + Returns: + set: Set of valid root element names + """ + try: + schema = schema.strip() + ET.register_namespace('xs', 'http://www.w3.org/2001/XMLSchema') + root = ET.fromstring(schema) + ns = {'xs': 'http://www.w3.org/2001/XMLSchema'} + elements = root.findall(".//xs:element", ns) + return {elem.get('name') for elem in elements if elem.get('name')} + except ET.ParseError as e: + print(f"Error parsing schema: {e}") + return set() + +def split_response(response: str, valid_elements: set) -> InferenceResult: + """ + Split the response into reasoning and actions based on valid XML elements. + + Args: + response: Raw response string from the model + valid_elements: Set of valid root element names from the schema + + Returns: + InferenceResult: Tuple containing reasoning and actions + """ + elements_pattern = '|'.join(map(re.escape, valid_elements)) + pattern = f"<({elements_pattern})[^>]*>" + matches = list(re.finditer(pattern, response)) + if not matches: + return InferenceResult(response.strip(), "") + last_match = matches[-1] + split_point = last_match.start() + reasoning = response[:split_point].strip() + actions = response[split_point:].strip() + return InferenceResult(reasoning, actions) \ No newline at end of file diff --git a/test.sh b/test.sh old mode 100644 new mode 100755 diff --git a/test/__init__.py b/test/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/test/llm_engine_test.py b/test/llm_engine_test.py new file mode 100644 index 0000000..f7943ff --- /dev/null +++ b/test/llm_engine_test.py @@ -0,0 +1,25 @@ +import unittest + +from . import test_data +from . import test_util + +from sia.llm_engine import LlmEngine, InferenceResult + +class LlmEngineTest(unittest.TestCase): + def setUp(self): + self.model_path = "/root/model" + self.llm_engine = LlmEngine(self.model_path) + + 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(test_data.echo_system_prompt, main_context, test_data.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}") \ No newline at end of file diff --git a/test/test_data.py b/test/test_data.py new file mode 100644 index 0000000..019ccb3 --- /dev/null +++ b/test/test_data.py @@ -0,0 +1,32 @@ +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. + + example input: + hello + example output: + hellohello +""".strip() + +echo_action_schema = """ + + + + + + + + + + + + +""".strip() \ No newline at end of file diff --git a/test/test_util.py b/test/test_util.py new file mode 100644 index 0000000..695c61f --- /dev/null +++ b/test/test_util.py @@ -0,0 +1,8 @@ +import unittest + +class SequentialTestLoader(unittest.TestLoader): + def getTestCaseNames(self, testCaseClass): + test_names = super().getTestCaseNames(testCaseClass) + testcase_methods = list(testCaseClass.__dict__.keys()) + test_names.sort(key=testcase_methods.index) + return test_names \ No newline at end of file diff --git a/test/util_test.py b/test/util_test.py new file mode 100644 index 0000000..bd4ed8d --- /dev/null +++ b/test/util_test.py @@ -0,0 +1,17 @@ +import unittest + +from . import test_data + +from sia import util + +class UtilTest(unittest.TestCase): + def test_get_valid_root_elements_single(self): + valid_elements = util.get_valid_root_elements(test_data.echo_action_schema) + self.assertEqual(valid_elements, {'test_tag'}) + + def test_split_response_single_element(self): + response = "Some reasoning here\ncontent" + valid_elements = {'test_tag'} + result = util.split_response(response, valid_elements) + self.assertEqual(result.reasoning, "Some reasoning here") + self.assertEqual(result.actions, "content") \ No newline at end of file diff --git a/tests/test_llm_engine.py b/tests/test_llm_engine.py deleted file mode 100644 index e56c77e..0000000 --- a/tests/test_llm_engine.py +++ /dev/null @@ -1,66 +0,0 @@ -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