From 70ed16f8ab513caf15fc889698138b71720ef17d Mon Sep 17 00:00:00 2001 From: Niels Geens Date: Mon, 4 Nov 2024 17:08:52 +0100 Subject: [PATCH] Added hf_llm_engine and config --- .dockerignore | 1 + .gitignore | 1 + Dockerfile | 2 +- requirements.txt | 1 + run.sh | 5 +- sia/__main__.py | 38 +++-- sia/config.py | 137 ++++++++++++++++++ sia/hf_llm_engine.py | 71 +++++++++ sia/llm_engine.py | 78 +--------- sia/local_llm_engine.py | 79 ++++++++++ system_prompt.md | 3 +- ...ngine_test.py => local_llm_engine_test.py} | 7 +- 12 files changed, 330 insertions(+), 93 deletions(-) create mode 100644 sia/config.py create mode 100644 sia/hf_llm_engine.py create mode 100644 sia/local_llm_engine.py rename test/{llm_engine_test.py => local_llm_engine_test.py} (79%) diff --git a/.dockerignore b/.dockerignore index e3f10e1..e346f21 100644 --- a/.dockerignore +++ b/.dockerignore @@ -1 +1,2 @@ +.env ./model/ \ No newline at end of file diff --git a/.gitignore b/.gitignore index 4591654..0b5394f 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ +.env pdf/ model/ claude.txt \ No newline at end of file diff --git a/Dockerfile b/Dockerfile index a865d42..e0ff00c 100644 --- a/Dockerfile +++ b/Dockerfile @@ -32,4 +32,4 @@ COPY ./ /root/sia/ COPY --from=web-build /app/dist /root/sia/static/ WORKDIR /root/sia -CMD ["python3", "-m", "sia"] \ No newline at end of file +ENTRYPOINT ["python3", "-m", "sia"] \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 86d524c..07d56b5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,6 @@ accelerate aiohttp bs4 +python-dotenv torch transformers \ No newline at end of file diff --git a/run.sh b/run.sh index cd400c7..71551ae 100755 --- a/run.sh +++ b/run.sh @@ -1,7 +1,7 @@ #!/bin/bash docker build \ - --tag sia + --tag sia \ . docker run \ @@ -10,4 +10,5 @@ docker run \ --gpus=all \ -p 8080:8080 \ -v /$(pwd)/model/:/root/model/ \ - sia + -v /$(pwd)/.env:/root/.env \ + sia "$@" diff --git a/sia/__main__.py b/sia/__main__.py index e8d2fc7..9210aa5 100644 --- a/sia/__main__.py +++ b/sia/__main__.py @@ -5,16 +5,20 @@ import asyncio import mimetypes import time + +from .config import Config +from .hf_llm_engine import HfLlmEngine from .llm_engine import LlmEngine +from .local_llm_engine import LocalLlmEngine +from .response_parser import ResponseParser from .system_metrics import SystemMetrics from .web_agent import WebAgent +from .web_agent import WebAgent +from .web_io_buffer import WebIOBuffer from .web_io_buffer import WebIOBuffer from .web_socket_manager import WebSocketManager from .working_memory import WorkingMemory from .xml_validator import XMLValidator -from .response_parser import ResponseParser -from .web_agent import WebAgent -from .web_io_buffer import WebIOBuffer mimetypes.add_type("application/javascript", ".js") mimetypes.add_type("application/javascript", ".jsx") @@ -31,13 +35,23 @@ class TestLLM: class Main: def __init__(self): - self._base_dir = Path(__file__).parent.parent - self._system_prompt = (self._base_dir / "system_prompt.md").read_text() - self._action_schema = (self._base_dir / "action_schema.xsd").read_text() - self._static_dir = self._base_dir / "static" + self._config = Config() - self._llm = LlmEngine("/root/model") - #self._llm = TestLLM() + self._system_prompt = self._config.system_prompt.read_text() + self._action_schema = self._config.action_schema.read_text() + + match self._config.llm_engine: + case "local": + self._llm = LocalLlmEngine(self._config.model) + case "hf": + self._llm = HfLlmEngine( + model_id=self._config.model, + api_token=self._config.hf_api_token + ) + case "test": + self._llm = TestLLM() + case _: + raise ValueError(f"Invalid LLM engine: {self._config.llm_engine}") self._io_buffer = WebIOBuffer() self._agent = WebAgent( system_prompt=self._system_prompt, @@ -64,13 +78,13 @@ class Main: self._app.middlewares.append(self._cors_middleware) self._app.router.add_get("/ws", self._ws_manager.handle_websocket) self._app.router.add_get("/", self._serve_index) - self._app.router.add_static("/static/", self._static_dir, show_index=False) - self._app.router.add_static("/assets/", self._static_dir / "assets", show_index=False) + self._app.router.add_static("/static/", self._config.static_files, show_index=False) + self._app.router.add_static("/assets/", self._config.static_files / "assets", show_index=False) self._app.router.add_get("/{path:.*}", self._serve_index) async def _serve_index(self, request: web.Request) -> web.Response: """Serve the React application HTML for any unmatched routes.""" - index_path = self._static_dir / "index.html" + index_path = self._config.static_files / "index.html" if not index_path.exists(): raise web.HTTPNotFound() diff --git a/sia/config.py b/sia/config.py new file mode 100644 index 0000000..6755a14 --- /dev/null +++ b/sia/config.py @@ -0,0 +1,137 @@ +from dataclasses import dataclass +from dotenv import load_dotenv +from pathlib import Path +from typing import Optional +import argparse +import os + +@dataclass +class Config: + """ + Configuration class that handles both command line and environment variables. + + Command line arguments take precedence over environment variables. + Environment variables serve as defaults that can be overridden via CLI. + """ + + def __init__(self): + """ + Create configuration from command line arguments and environment variables. + Required arguments must be provided either via CLI or environment variables. + """ + load_dotenv() + parser = argparse.ArgumentParser(description='SIA - Self Improving Agent') + parser.add_argument( + '--system-prompt', + type=Path, + default=os.getenv('SIA_SYSTEM_PROMPT', 'system_prompt.md'), + help='Path to the system prompt file (default: system_prompt.md, env: SIA_SYSTEM_PROMPT)' + ) + parser.add_argument( + '--action-schema', + type=Path, + default=os.getenv('SIA_ACTION_SCHEMA', 'action_schema.xsd'), + help='Path to the action schema file (default: action_schema.xsd, env: SIA_ACTION_SCHEMA)' + ) + parser.add_argument( + '--server', + action='store_true', + default=self._parse_bool_env('SIA_SERVER_ENABLED', False), + help='Enable web server for debugging and human feedback (env: SIA_SERVER_ENABLED)' + ) + parser.add_argument( + '--host', + type=str, + default=os.getenv('SIA_SERVER_HOST', 'localhost'), + help='Web server host (default: localhost, env: SIA_SERVER_HOST)' + ) + parser.add_argument( + '--port', + type=int, + default=int(os.getenv('SIA_SERVER_PORT', '8080')), + help='Web server port (default: 8080, env: SIA_SERVER_PORT)' + ) + parser.add_argument( + '--static-files', + type=Path, + default=self._parse_optional_path('SIA_STATIC_FILES', './static/'), + help='Path to static web files (default: ./static/, env: SIA_STATIC_FILES)' + ) + parser.add_argument( + '--llm-engine', + type=str, + default=os.getenv('SIA_LLM_ENGINE', 'local'), + help='LLM engine (default: local, env: SIA_LLM_ENGINE)' + ) + parser.add_argument( + '--hf-api-token', + type=str, + default=os.getenv('SIA_HF_API_TOKEN'), + help='Hugging Face access token (env: SIA_HF_API_TOKEN)' + ) + parser.add_argument( + '--model', + type=str, + default=os.getenv('SIA_MODEL', '/root/model/'), + help='Path to the model directory (default: /root/model/, env: SIA_MODEL)' + ) + self.args = parser.parse_args() + + def _parse_bool_env(self, env_var: str, default: bool) -> bool: + """Parse boolean environment variable.""" + val = os.getenv(env_var) + if val is None: + return default + return val.lower() in ('true', '1', 'yes', 'on') + + def _parse_optional_path(self, env_var: str, default: Optional[Path]) -> Optional[Path]: + """Parse optional Path environment variable.""" + val = os.getenv(env_var) + if val is None: + return default + return Path(val) + + @property + def system_prompt(self) -> Path: + """Path to the system prompt file.""" + return self.args.system_prompt + + @property + def action_schema(self) -> Path: + """Path to the action schema file.""" + return self.args.action_schema + + @property + def server(self) -> bool: + """Enable web server for debugging and human feedback.""" + return self.args.server + + @property + def host(self) -> str: + """Web server host.""" + return self.args.host + + @property + def port(self) -> int: + """Web server port.""" + return self.args.port + + @property + def static_files(self) -> Path: + """Path to static web files.""" + return self.args.static_files + + @property + def llm_engine(self) -> str: + """LLM engine.""" + return self.args.llm_engine + + @property + def hf_api_token(self) -> Optional[str]: + """Hugging Face access token.""" + return self.args.hf_api_token + + @property + def model(self) -> str: + """Path to the model directory.""" + return self.args.model \ No newline at end of file diff --git a/sia/hf_llm_engine.py b/sia/hf_llm_engine.py new file mode 100644 index 0000000..3182030 --- /dev/null +++ b/sia/hf_llm_engine.py @@ -0,0 +1,71 @@ +from typing import Iterator, Optional +from huggingface_hub import InferenceClient + +from .llm_engine import LlmEngine + +class HfLlmEngine(LlmEngine): + """ + LLM Engine implementation using HuggingFace's InferenceClient. + """ + + def __init__( + self, + model_id: str = "mistralai/Mistral-7B-Instruct-v0.2", + api_token: Optional[str] = None, + temperature: float = 0.7, + max_new_tokens: int = 1024, + ): + """ + Initialize the HuggingFace Inference API LLM Engine. + + Args: + model_id: HuggingFace model ID to use (default: Mistral-7B-Instruct) + api_token: HuggingFace API token. If None, will try to read from HF_TOKEN env var + temperature: Sampling temperature (default: 0.7) + max_new_tokens: Maximum number of tokens to generate (default: 1024) + """ + self.model_id = model_id + self.client = InferenceClient(token=api_token) + + # Generation parameters + self.temperature = temperature + self.max_new_tokens = max_new_tokens + + def set_model_path(self, model_id: str): + """ + Update the model being used. + + Args: + model_id: New HuggingFace model ID to use + """ + self.model_id = model_id + + def infer(self, system_prompt: str, main_context: str) -> 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 + + Returns: + Iterator[str]: An iterator that yields the generated text. + """ + messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": main_context} + ] + + def stream_wrapper(): + stream = self.client.chat_completion( + model=self.model_id, + messages=messages, + temperature=self.temperature, + max_tokens=self.max_new_tokens, + stream=True + ) + + for response in stream: + if content := response.choices[0].delta.content: + yield content + return stream_wrapper() \ No newline at end of file diff --git a/sia/llm_engine.py b/sia/llm_engine.py index d739266..522c790 100644 --- a/sia/llm_engine.py +++ b/sia/llm_engine.py @@ -1,78 +1,8 @@ -from threading import Thread from typing import Iterator +from abc import ABC, abstractmethod -from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer -import torch - -from . import util - -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.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. - """ - self.tokenizer = AutoTokenizer.from_pretrained(model_path) - model = AutoModelForCausalLM.from_pretrained( - model_path, - return_dict=True, - low_cpu_mem_usage=True, - torch_dtype=torch.bfloat16, - device_map="auto", - trust_remote_code=True, - ) - 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=self.tokenizer, - torch_dtype=torch.bfloat16, - device_map="auto", - return_full_text=False, - ) +class LlmEngine(ABC): + @abstractmethod def infer(self, system_prompt: str, main_context: str) -> Iterator[str]: - """ - 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 - - Returns: - Iterator[str]: An iterator that yields the generated text. - """ - messages = [ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": main_context} - ] - prompt = self.tokenizer.apply_chat_template( - messages, tokenize=False, add_generation_prompt=True - ) - streamer = TextIteratorStreamer( - self.tokenizer, - skip_prompt=True - ) - pipeline_kwargs = dict( - text_inputs=prompt, - do_sample=True, - max_new_tokens=1024, - streamer=streamer - ) - thread = Thread(target=self.pipeline, kwargs=pipeline_kwargs) - thread.start() - return util.stop_before_value(streamer, '<|eot_id|>') \ No newline at end of file + pass \ No newline at end of file diff --git a/sia/local_llm_engine.py b/sia/local_llm_engine.py new file mode 100644 index 0000000..f3997d2 --- /dev/null +++ b/sia/local_llm_engine.py @@ -0,0 +1,79 @@ +from threading import Thread +from typing import Iterator + +from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer +import torch + +from . import util +from .llm_engine import LlmEngine + +class LocalLlmEngine(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.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. + """ + self.tokenizer = AutoTokenizer.from_pretrained(model_path) + model = AutoModelForCausalLM.from_pretrained( + model_path, + return_dict=True, + low_cpu_mem_usage=True, + torch_dtype=torch.bfloat16, + device_map="auto", + trust_remote_code=True, + ) + 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=self.tokenizer, + torch_dtype=torch.bfloat16, + device_map="auto", + return_full_text=False, + ) + + def infer(self, system_prompt: str, main_context: str) -> Iterator[str]: + """ + 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 + + Returns: + Iterator[str]: An iterator that yields the generated text. + """ + messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": main_context} + ] + prompt = self.tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + streamer = TextIteratorStreamer( + self.tokenizer, + skip_prompt=True + ) + pipeline_kwargs = dict( + text_inputs=prompt, + do_sample=True, + max_new_tokens=1024, + streamer=streamer + ) + thread = Thread(target=self.pipeline, kwargs=pipeline_kwargs) + thread.start() + return util.stop_before_value(streamer, '<|eot_id|>') diff --git a/system_prompt.md b/system_prompt.md index d1b2831..684da32 100644 --- a/system_prompt.md +++ b/system_prompt.md @@ -4,9 +4,10 @@ You can solve any problem. Each iteration, the context is updated with the result of your previous actions. You modify the context by issuing a commands using XML. -Always respond with one action adhering to the XML schema. Parameters and scripts may be long and complex. Use correct XML escaping or CDATA sections. +It is very important that you always respond with one action adhering to the XML schema! +Do not respond with anything else after the action. # Context The context has a limited length. diff --git a/test/llm_engine_test.py b/test/local_llm_engine_test.py similarity index 79% rename from test/llm_engine_test.py rename to test/local_llm_engine_test.py index a11b9c5..f78fce8 100644 --- a/test/llm_engine_test.py +++ b/test/local_llm_engine_test.py @@ -5,21 +5,22 @@ from itertools import tee from . import test_data from sia.llm_engine import LlmEngine +from sia.local_llm_engine import LocalLlmEngine class LlmEngineTest(unittest.TestCase): def setUp(self): self.model_path = "/root/model" def test_initialization(self): - llm_engine = LlmEngine(self.model_path) + llm_engine = LocalLlmEngine(self.model_path) self.assertIsInstance(llm_engine, LlmEngine) def test_infer(self): main_context = "This is a test" - llm_engine = LlmEngine(self.model_path) + llm_engine = LocalLlmEngine(self.model_path) tokens = llm_engine.infer(test_data.echo_system_prompt, main_context) print_tokens, result_tokens = tee(tokens) for token in print_tokens: print(token, end="", flush=True) result = ''.join(result_tokens) - self.assertEqual(result, f"{main_context}{main_context}") \ No newline at end of file + self.assertEqual(result, f"{main_context}{main_context}")