Create base_generator.py

services/base_generator.py

@@ -0,0 +1,52 @@
+# base_generator.py
+from abc import ABC, abstractmethod
+from typing import AsyncGenerator, Dict, Any, Optional, List, Tuple
+from dataclasses import dataclass
+from logging import getLogger 
+
+
+from config.config import GenerationConfig, ModelConfig
+
+class BaseGenerator(ABC):
+    """Base class for all generator implementations."""
+    
+    def __init__(
+        self,
+        model_name: str,
+        device: Optional[str] = None,
+        default_generation_config: Optional[GenerationConfig] = None,
+        model_config: Optional[ModelConfig] = None,
+        cache_size: int = 1000,
+        max_batch_size: int = 32
+    ):
+        self.logger = getLogger(__name__)
+        self.model_manager = ModelManager(device)
+        self.cache = ResponseCache(cache_size)
+        self.batch_processor = BatchProcessor(max_batch_size)
+        self.health_check = HealthCheck()
+       # self.tokenizer = self.model_manager.tokenizers[model_name]
+        #self.tokenizer = self.load_tokenizer(llama_model_name)  # Add this line to initialize the tokenizer
+        self.default_config = default_generation_config or GenerationConfig()
+        self.model_config = model_config or ModelConfig()
+        
+    @abstractmethod
+    async def generate_stream(
+        self,
+        prompt: str,
+        config: Optional[GenerationConfig] = None
+    ) -> AsyncGenerator[str, None]:
+        pass
+        
+    @abstractmethod
+    def _get_generation_kwargs(self, config: GenerationConfig) -> Dict[str, Any]:
+        pass
+ 
+    @abstractmethod
+    def generate(
+        self,
+        prompt: str,
+        model_kwargs: Dict[str, Any],
+        strategy: str = "default",
+        **kwargs
+    ) -> str:
+        pass