Added hf_llm_engine and config

2024-11-04 17:08:52 +01:00
parent 5da6dca5ec
commit 70ed16f8ab
12 changed files with 330 additions and 93 deletions
--- 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()

--- a/sia/config.py
+++ 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
--- a/sia/hf_llm_engine.py
+++ 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()
--- 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|>')
+        pass
--- a/sia/local_llm_engine.py
+++ 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|>')