wip deepseek r1

sia/llm_engine/__init__.py

@@ -0,0 +1,15 @@
+from typing import Callable, Iterator
+from abc import ABC, abstractmethod
+
+class LlmEngine(ABC):
+    @abstractmethod
+    def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
+        pass
+
+    @abstractmethod
+    def token_count(self, system_prompt: str, main_context: str) -> int:
+        pass
+
+    @abstractmethod
+    def token_limit(self) -> int:
+        pass

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

@@ -0,0 +1,84 @@
+from huggingface_hub import InferenceClient
+from transformers import AutoTokenizer, AutoConfig
+from typing import Iterator, Optional, Callable
+
+from . import LlmEngine
+
+class HfLlmEngine(LlmEngine):
+    """
+    LLM Engine implementation using HuggingFace's InferenceClient.
+    """
+    
+    def __init__(
+        self,
+        model: str,
+        temperature: float,
+        api_token: Optional[str],
+    ):
+        """
+        Initialize the HuggingFace Inference API LLM Engine.
+
+        Args:
+            model: HuggingFace model ID to use
+            temperature: Sampling temperature
+            api_token: HuggingFace API token
+        """
+        self._model = model
+        self._temperature = temperature
+
+        self._tokenizer = AutoTokenizer.from_pretrained(model, token=api_token)
+        self._config = AutoConfig.from_pretrained(model, token=api_token)
+        self._client = InferenceClient(token=api_token)
+                        
+    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.
+        """
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context}
+        ]
+        
+        stream = self._client.chat_completion(
+            model=self._model,
+            messages=messages,
+            temperature=self._temperature,
+            stream=True
+        )
+        
+        try:
+            for response in stream:
+                if should_stop():
+                    stream.close()
+                    break
+                if content := response.choices[0].delta.content:
+                    yield content
+        finally:
+            stream.close()
+
+    def token_count(self, system_prompt: str, main_context: str) -> int:
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context}
+        ]
+        prompt = self._tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        return len(self._tokenizer.encode(prompt))
+
+    def token_limit(self) -> int:
+        """
+        Get the model's context window size.
+        
+        Returns:
+            int: Maximum number of tokens the model can process
+        """
+        return self._config.max_position_embeddings

sia/llm_engine/local_llm_engine.py

@@ -0,0 +1,121 @@
+from threading import Thread
+from typing import Iterator, Optional, Callable
+
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
+import torch
+
+from . import LlmEngine
+from .. import util
+
+class LocalLlmEngine(LlmEngine):
+    def __init__(
+        self,
+        model_path: str,
+        temperature: float,
+        token_limit: int,
+        api_token: Optional[str],
+    ):
+        """
+        Initialize the LLM Engine with a model path.
+
+        Args:
+            model_path: Path to the model weights to be used.
+            temperature: Temperature for sampling
+            api_token: Huggingface API key
+            token_limit: Maximum number of tokens to generate
+        """
+        self._temperature = temperature
+        self._token_limit = token_limit
+        self._tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
+        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,
+            token=api_token,
+        )
+        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, 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.
+        """
+        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
+        )
+        generation_thread = Thread(target=self._pipeline, kwargs=dict(
+            text_inputs=prompt,
+            do_sample=True,
+            temperature=self._temperature,
+            max_new_tokens=self.token_limit(),
+            streamer=streamer
+        ))
+        generation_thread.start()
+
+        for text in util.stop_before_value(streamer, '<|eot_id|>'):
+            yield text
+            if should_stop():
+                break
+
+        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
+        """
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context}
+        ]
+        prompt = self._tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        return len(self._tokenizer.encode(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
+        else:
+            return self._pipeline.model.config.max_position_embeddings

sia/llm_engine/mistral_llm_engine.py

@@ -0,0 +1,71 @@
+from typing import Iterator, Optional, Callable
+from mistralai import Mistral
+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 . import LlmEngine
+
+class MistralLlmEngine(LlmEngine):
+    def __init__(
+            self,
+            model: str,
+            temperature: float,
+            token_limit: int,
+            api_key: str,
+        ):
+        self._model = model
+        self._temperature = temperature
+        self._token_limit = token_limit
+        self._api_key = api_key
+        self._client = Mistral(api_key=api_key)
+        self._tokenizer = MistralTokenizer.v3()
+
+    def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
+        messages = [
+            {
+                "role": "system",
+                "content": system_prompt,
+            },
+            {
+                "role": "user",
+                "content": main_context,
+            },
+            {
+                "role": "assistant",
+                "content": "<",
+                "prefix": True,
+            },
+        ]
+        stream_response = self._client.chat.stream(
+            model=self._model,
+            messages=messages,
+            temperature=self._temperature,
+        )
+        
+        try:
+            for chunk in stream_response:
+                if should_stop():
+                    stream_response.response.close()
+                    break
+                if content := chunk.data.choices[0].delta.content:
+                    yield content
+        finally:
+            stream_response.response.close()
+
+    def token_count(self, system_prompt: str, main_context: str) -> int:
+        messages = [
+            SystemMessage(content=system_prompt),
+            UserMessage(content=main_context),
+        ]
+
+        tokenized = self._tokenizer.encode_chat_completion(
+            ChatCompletionRequest(
+                messages=messages,
+                model=self._model
+            )
+        )
+        return len(tokenized.tokens)
+
+    def token_limit(self) -> int:
+        return self._token_limit

sia/llm_engine/openai_llm_engine.py

@@ -0,0 +1,75 @@
+from typing import Callable, Iterator
+import openai
+import tiktoken
+
+from . import LlmEngine
+
+class OpenAILlmEngine(LlmEngine):
+    """
+    LLM Engine implementation using OpenAI's API.
+    Supports streaming responses from chat completion models.
+    """
+    
+    def __init__(
+        self,
+        model: str,
+        temperature: float,
+        token_limit: int,
+        api_key: str,
+    ):
+        """
+        Initialize the OpenAI LLM Engine.
+
+        Args:
+            model: OpenAI model to use
+            temperature: Temperature for sampling
+            api_key: OpenAI API key
+            token_limit: Maximum number of tokens to generate
+        """
+        self._model = model
+        self._temperature = temperature
+        self._token_limit = token_limit
+
+        self._client = openai.Client(
+            api_key=api_key,
+        )
+                        
+    def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context}
+        ]
+        
+        stream = self._client.chat.completions.create(
+            model=self._model,
+            messages=messages,
+            temperature=self._temperature,
+            stream=True,
+        )
+
+        try:
+            for chunk in stream:
+                if should_stop():
+                    break
+                if content := chunk.choices[0].delta.content:
+                    yield content
+        finally:
+            stream.close()
+            #stream.response.close()
+    
+    def token_count(self, system_prompt: str, main_context: str) -> int:
+        """
+        Calculate the total token count for the system prompt and context.
+        
+        Args:
+            system_prompt: The system prompt string
+            main_context: The main context string
+            
+        Returns:
+            int: Total number of tokens
+        """
+        encoding = tiktoken.encoding_for_model(self._model)
+        return len(encoding.encode(system_prompt)) + len(encoding.encode(main_context))
+    
+    def token_limit(self) -> int:
+        return self._token_limit