wip deepseek r1

sia/__main__.py

@@ -3,11 +3,12 @@ import asyncio
 
 from .auto_approver import AutoApprover
 from .config import Config
-from .hf_llm_engine import HfLlmEngine
+from .llm_engine.hf_llm_engine import HfLlmEngine
+from .llm_engine.deepseek_llm_engine import DeepSeekLlmEngine
 from .iteration_logger import IterationLogger
-from .local_llm_engine import LocalLlmEngine
-from .mistral_llm_engine import MistralLlmEngine
-from .openai_llm_engine import OpenAILlmEngine
+from .llm_engine.local_llm_engine import LocalLlmEngine
+from .llm_engine.mistral_llm_engine import MistralLlmEngine
+from .llm_engine.openai_llm_engine import OpenAILlmEngine
 from .response_parser import ResponseParser
 from .system_metrics import SystemMetrics
 from .web.api import Api
@@ -61,6 +62,14 @@ class Main:
                 config.mistral_api_key,
             )
 
+        if config.deepseek_enabled:
+            self._llms['deepseek'] = DeepSeekLlmEngine(
+                config.deepseek_model,
+                config.deepseek_temperature,
+                config.deepseek_token_limit,
+                config.hf_api_key,  # Use the existing HF API key
+            )
+
         if not self._llms:
             raise ValueError("No LLM engines enabled in configuration")
 
@@ -103,9 +112,10 @@ class Main:
             content_type="text/html"
         )
 
-if __name__ == "__main__":
+def main():
     loop = asyncio.new_event_loop()
     config = Config()
-    main = loop.run_until_complete(Main.create(config))
+    main_instance = loop.run_until_complete(Main.create(config))
     print(f"Web server started at http://localhost:{config.port}")
-    web.run_app(main.app, loop=loop, host=config.host, port=config.port)
+    web.run_app(main_instance.app, loop=loop, host=config.host, port=config.port)
+    return 0

sia/config.py

@@ -184,6 +184,30 @@ class Config:
             default=os.getenv('SIA_MISTRAL_API_KEY'),
             help='Mistral API key (env: SIA_MISTRAL_API_KEY)'
         )
+        parser.add_argument(
+            '--deepseek-enable',
+            action='store_true',
+            default=self._parse_bool_env('SIA_DEEPSEEK_ENABLED', False),
+            help='Enable DeepSeek LLM engine (env: SIA_DEEPSEEK_ENABLED)'
+        )
+        parser.add_argument(
+            '--deepseek-model',
+            type=str,
+            default=os.getenv('SIA_DEEPSEEK_MODEL', '/root/models/current'),
+            help='Path to fine-tuned DeepSeek model (env: SIA_DEEPSEEK_MODEL)'
+        )
+        parser.add_argument(
+            '--deepseek-temperature',
+            type=float,
+            default=float(os.getenv('SIA_DEEPSEEK_TEMPERATURE', '0.6')),
+            help='DeepSeek temperature (default: 0.6, env: SIA_DEEPSEEK_TEMPERATURE)'
+        )
+        parser.add_argument(
+            '--deepseek-token-limit',
+            type=int,
+            default=int(os.getenv('SIA_DEEPSEEK_TOKEN_LIMIT', '0')),
+            help='DeepSeek token limit (0 for model default, env: SIA_DEEPSEEK_TOKEN_LIMIT)'
+        )
 
         self.args = parser.parse_args()
 
@@ -312,3 +336,20 @@ class Config:
     @property
     def mistral_api_key(self) -> Optional[str]:
         return self.args.mistral_api_key
+
+    @property
+    def deepseek_enabled(self) -> bool:
+        return self.args.deepseek_enable
+
+    @property
+    def deepseek_model(self) -> str:
+        return self.args.deepseek_model
+
+    @property
+    def deepseek_temperature(self) -> float:
+        return self.args.deepseek_temperature
+
+    @property
+    def deepseek_token_limit(self) -> Optional[int]:
+        # Return None if 0 to use model default
+        return self.args.deepseek_token_limit if self.args.deepseek_token_limit > 0 else None

sia/llm_engine/deepseek_llm_engine.py

@@ -0,0 +1,150 @@
+from typing import Callable, Iterator, Optional
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from threading import Thread
+from pathlib import Path
+
+from . import LlmEngine
+from .. import util
+
+class DeepSeekLlmEngine(LlmEngine):
+    """
+    LLM Engine implementation for DeepSeek models.
+    Supports fine-tuned DeepSeek-R1 and its distilled versions.
+    """
+    
+    def __init__(
+        self,
+        model_path: str,
+        temperature: float = 0.6,
+        token_limit: Optional[int] = None,
+        api_key: Optional[str] = None,
+    ):
+        """
+        Initialize the DeepSeek LLM Engine.
+
+        Args:
+            model_path: Local path to the fine-tuned model
+            temperature: Sampling temperature (0.6 default as recommended)
+            token_limit: Maximum tokens to generate or context length override
+            api_key: HuggingFace API token if needed
+        """
+        self._model_path = Path(model_path)
+        self._temperature = temperature
+        self._token_limit = token_limit
+        
+        # Load tokenizer with trust_remote_code for DeepSeek models
+        self._tokenizer = AutoTokenizer.from_pretrained(
+            self._model_path,
+            token=api_key,
+            trust_remote_code=True,
+        )
+        
+        # Set padding token to avoid warnings
+        if self._tokenizer.pad_token is None:
+            self._tokenizer.pad_token = self._tokenizer.eos_token
+            
+        # Load model with 4-bit quantization by default
+        self._device_map = "auto"
+        
+        self._model = AutoModelForCausalLM.from_pretrained(
+            self._model_path,
+            return_dict=True,
+            low_cpu_mem_usage=True,
+            trust_remote_code=True,
+            device_map=self._device_map,
+            load_in_4bit=True, 
+            torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
+            token=api_key,
+        )
+        
+        # Ensure model is in evaluation mode
+        self._model.eval()
+                        
+    def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> 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
+            should_stop: Callback that returns True when inference should stop
+
+        Returns:
+            Iterator[str]: An iterator that yields the generated text.
+        """
+        # Tokenize input
+        inputs = self._tokenizer(system_prompt + "\n\n" + main_context, return_tensors="pt").to(self._device_map)
+        
+        # Create streamer for token-by-token generation
+        streamer = TextIteratorStreamer(
+            self._tokenizer,
+            skip_prompt=True,
+            timeout=15.0
+        )
+        
+        # Generate in a separate thread to enable streaming
+        generation_kwargs = {
+            "input_ids": inputs.input_ids,
+            "attention_mask": inputs.attention_mask,
+            "max_new_tokens": self.token_limit() if self._token_limit else 2048,
+            "temperature": self._temperature,
+            "do_sample": True,
+            "streamer": streamer,
+            "repetition_penalty": 1.1,
+            "pad_token_id": self._tokenizer.pad_token_id,
+        }
+        
+        generation_thread = Thread(target=self._model.generate, kwargs=generation_kwargs)
+        generation_thread.start()
+        
+        # Yield tokens as they become available
+        try:
+            for text in streamer:
+                yield text
+                if should_stop():
+                    break
+        finally:
+            # Ensure thread is properly joined even if iteration is interrupted
+            generation_thread.join()
+
+    def token_count(self, system_prompt: str, main_context: str) -> int:
+        """
+        Count tokens for the given system prompt and main context.
+        
+        Args:
+            system_prompt: The system prompt string
+            main_context: The main context string
+            
+        Returns:
+            int: Total number of tokens
+        """
+        combined_prompt = f"{system_prompt}\n\n{main_context}"
+        return len(self._tokenizer.encode(combined_prompt))
+
+    def token_limit(self) -> int:
+        """
+        Get the model's context window size.
+        
+        Returns:
+            int: Maximum number of tokens the model can process
+        """
+        if self._token_limit is not None:
+            return self._token_limit
+        
+        # Try to detect model size from config
+        try:
+            config_file = self._model_path / "config.json"
+            if config_file.exists():
+                import json
+                with open(config_file, 'r') as f:
+                    config = json.load(f)
+                if 'max_position_embeddings' in config:
+                    return config['max_position_embeddings']
+                if 'model_max_length' in config:
+                    return config['model_max_length']
+        except Exception:
+            pass
+            
+        # Default to 8k if we can't determine
+        return 8192

sia/llm_engine/hf_llm_engine.py

@@ -2,7 +2,7 @@ from huggingface_hub import InferenceClient
 from transformers import AutoTokenizer, AutoConfig
 from typing import Iterator, Optional, Callable
 
-from .llm_engine import LlmEngine
+from . import LlmEngine
 
 class HfLlmEngine(LlmEngine):
     """

sia/llm_engine/local_llm_engine.py

@@ -4,8 +4,8 @@ from typing import Iterator, Optional, Callable
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
 import torch
 
-from . import util
-from .llm_engine import LlmEngine
+from . import LlmEngine
+from .. import util
 
 class LocalLlmEngine(LlmEngine):
     def __init__(

sia/llm_engine/mistral_llm_engine.py

@@ -4,7 +4,7 @@ from mistral_common.protocol.instruct.messages import SystemMessage, UserMessage
 from mistral_common.protocol.instruct.request import ChatCompletionRequest
 from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
 
-from .llm_engine import LlmEngine
+from . import LlmEngine
 
 class MistralLlmEngine(LlmEngine):
     def __init__(

sia/llm_engine/openai_llm_engine.py

@@ -2,7 +2,7 @@ from typing import Callable, Iterator
 import openai
 import tiktoken
 
-from .llm_engine import LlmEngine
+from . import LlmEngine
 
 class OpenAILlmEngine(LlmEngine):
     """