wip deepseek r1

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