isa-model 0.1.0__py3-none-any.whl

isa_model/inference/backends/Pytorch/gemma_backend.py

@@ -0,0 +1,167 @@
+import os
+import logging
+import torch
+from typing import Dict, List, Any, Optional, Union
+
+logger = logging.getLogger(__name__)
+
+
+class GemmaBackend:
+    """
+    PyTorch backend for the Gemma LLM model.
+    """
+    
+    def __init__(self, model_path: Optional[str] = None, device: str = "auto"):
+        """
+        Initialize the Gemma backend.
+        
+        Args:
+            model_path: Path to the model
+            device: Device to run the model on ("cpu", "cuda", or "auto")
+        """
+        self.model_path = model_path or os.environ.get("GEMMA_MODEL_PATH", "/models/Gemma3-4B")
+        self.device = device if device != "auto" else ("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = None
+        self.tokenizer = None
+        self._loaded = False
+        
+        # Default generation config
+        self.default_config = {
+            "max_new_tokens": 512,
+            "temperature": 0.7,
+            "top_p": 0.9,
+            "top_k": 50,
+            "repetition_penalty": 1.1,
+            "do_sample": True
+        }
+        
+        self.logger = logger
+    
+    def load(self) -> None:
+        """
+        Load the model and tokenizer.
+        """
+        if self._loaded:
+            return
+        
+        try:
+            from transformers import AutoModelForCausalLM, AutoTokenizer
+            
+            # Load tokenizer
+            self.logger.info(f"Loading Gemma tokenizer from {self.model_path}")
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
+            
+            # Load model
+            self.logger.info(f"Loading Gemma model on {self.device}")
+            if self.device == "cpu":
+                self.model = AutoModelForCausalLM.from_pretrained(
+                    self.model_path,
+                    torch_dtype=torch.float32,
+                    low_cpu_mem_usage=True,
+                    device_map="auto"
+                )
+            else:  # cuda
+                self.model = AutoModelForCausalLM.from_pretrained(
+                    self.model_path,
+                    torch_dtype=torch.float16,  # Use half precision on GPU
+                    device_map="auto"
+                )
+            
+            self.model.eval()
+            self._loaded = True
+            self.logger.info("Gemma model loaded successfully")
+            
+        except Exception as e:
+            self.logger.error(f"Failed to load Gemma model: {str(e)}")
+            raise
+    
+    def unload(self) -> None:
+        """
+        Unload the model and tokenizer.
+        """
+        if not self._loaded:
+            return
+        
+        self.model = None
+        self.tokenizer = None
+        self._loaded = False
+        
+        # Force garbage collection
+        import gc
+        gc.collect()
+        
+        if self.device == "cuda":
+            torch.cuda.empty_cache()
+        
+        self.logger.info("Gemma model unloaded")
+    
+    def generate(self, 
+                prompt: str, 
+                system_prompt: Optional[str] = None,
+                generation_config: Optional[Dict[str, Any]] = None) -> str:
+        """
+        Generate text from a prompt.
+        
+        Args:
+            prompt: User prompt
+            system_prompt: System prompt to control model behavior
+            generation_config: Configuration for text generation
+            
+        Returns:
+            Generated text
+        """
+        if not self._loaded:
+            self.load()
+        
+        # Get generation config
+        config = self.default_config.copy()
+        if generation_config:
+            config.update(generation_config)
+        
+        try:
+            # Format the prompt with system prompt if provided
+            if system_prompt:
+                # Gemma uses a specific format for system prompts
+                formatted_prompt = f"<bos><start_of_turn>system\n{system_prompt}<end_of_turn>\n<start_of_turn>user\n{prompt}<end_of_turn>\n<start_of_turn>model"
+            else:
+                formatted_prompt = f"<bos><start_of_turn>user\n{prompt}<end_of_turn>\n<start_of_turn>model"
+            
+            # Tokenize the prompt
+            inputs = self.tokenizer(formatted_prompt, return_tensors="pt").to(self.device)
+            
+            # Generate text
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    inputs.input_ids,
+                    attention_mask=inputs.attention_mask,
+                    pad_token_id=self.tokenizer.eos_token_id,
+                    **config
+                )
+            
+            # Decode the generated text
+            generated_text = self.tokenizer.decode(
+                outputs[0][inputs.input_ids.shape[1]:], 
+                skip_special_tokens=True
+            )
+            
+            return generated_text.strip()
+            
+        except Exception as e:
+            self.logger.error(f"Error during Gemma text generation: {str(e)}")
+            raise
+    
+    def get_model_info(self) -> Dict[str, Any]:
+        """
+        Get information about the model.
+        
+        Returns:
+            Dictionary containing model information
+        """
+        return {
+            "name": "gemma3-4b",
+            "type": "llm",
+            "device": self.device,
+            "path": self.model_path,
+            "loaded": self._loaded,
+            "default_config": self.default_config
+        } 

isa_model/inference/backends/Pytorch/llama_backend.py

@@ -0,0 +1,166 @@
+import os
+import logging
+import torch
+from typing import Dict, List, Any, Optional, Union
+
+logger = logging.getLogger(__name__)
+
+
+class LlamaBackend:
+    """
+    PyTorch backend for the Llama LLM model.
+    """
+    
+    def __init__(self, model_path: Optional[str] = None, device: str = "auto"):
+        """
+        Initialize the Llama backend.
+        
+        Args:
+            model_path: Path to the model
+            device: Device to run the model on ("cpu", "cuda", or "auto")
+        """
+        self.model_path = model_path or os.environ.get("LLAMA_MODEL_PATH", "/models/Llama3-8B")
+        self.device = device if device != "auto" else ("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = None
+        self.tokenizer = None
+        self._loaded = False
+        
+        # Default generation config
+        self.default_config = {
+            "max_new_tokens": 512,
+            "temperature": 0.7,
+            "top_p": 0.9,
+            "top_k": 50,
+            "repetition_penalty": 1.1,
+            "do_sample": True
+        }
+        
+        self.logger = logger
+    
+    def load(self) -> None:
+        """
+        Load the model and tokenizer.
+        """
+        if self._loaded:
+            return
+        
+        try:
+            from transformers import AutoModelForCausalLM, AutoTokenizer
+            
+            # Load tokenizer
+            self.logger.info(f"Loading Llama tokenizer from {self.model_path}")
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
+            
+            # Load model
+            self.logger.info(f"Loading Llama model on {self.device}")
+            if self.device == "cpu":
+                self.model = AutoModelForCausalLM.from_pretrained(
+                    self.model_path,
+                    torch_dtype=torch.float32,
+                    low_cpu_mem_usage=True,
+                    device_map="auto"
+                )
+            else:  # cuda
+                self.model = AutoModelForCausalLM.from_pretrained(
+                    self.model_path,
+                    torch_dtype=torch.float16,  # Use half precision on GPU
+                    device_map="auto"
+                )
+            
+            self.model.eval()
+            self._loaded = True
+            self.logger.info("Llama model loaded successfully")
+            
+        except Exception as e:
+            self.logger.error(f"Failed to load Llama model: {str(e)}")
+            raise
+    
+    def unload(self) -> None:
+        """
+        Unload the model and tokenizer.
+        """
+        if not self._loaded:
+            return
+        
+        self.model = None
+        self.tokenizer = None
+        self._loaded = False
+        
+        # Force garbage collection
+        import gc
+        gc.collect()
+        
+        if self.device == "cuda":
+            torch.cuda.empty_cache()
+        
+        self.logger.info("Llama model unloaded")
+    
+    def generate(self, 
+                prompt: str, 
+                system_prompt: Optional[str] = None,
+                generation_config: Optional[Dict[str, Any]] = None) -> str:
+        """
+        Generate text from a prompt.
+        
+        Args:
+            prompt: User prompt
+            system_prompt: System prompt to control model behavior
+            generation_config: Configuration for text generation
+            
+        Returns:
+            Generated text
+        """
+        if not self._loaded:
+            self.load()
+        
+        # Get generation config
+        config = self.default_config.copy()
+        if generation_config:
+            config.update(generation_config)
+        
+        try:
+            # Format the prompt with system prompt if provided
+            if system_prompt:
+                formatted_prompt = f"<|system|>\n{system_prompt}\n<|user|>\n{prompt}\n<|assistant|>"
+            else:
+                formatted_prompt = f"<|user|>\n{prompt}\n<|assistant|>"
+            
+            # Tokenize the prompt
+            inputs = self.tokenizer(formatted_prompt, return_tensors="pt").to(self.device)
+            
+            # Generate text
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    inputs.input_ids,
+                    attention_mask=inputs.attention_mask,
+                    pad_token_id=self.tokenizer.eos_token_id,
+                    **config
+                )
+            
+            # Decode the generated text
+            generated_text = self.tokenizer.decode(
+                outputs[0][inputs.input_ids.shape[1]:], 
+                skip_special_tokens=True
+            )
+            
+            return generated_text.strip()
+            
+        except Exception as e:
+            self.logger.error(f"Error during Llama text generation: {str(e)}")
+            raise
+    
+    def get_model_info(self) -> Dict[str, Any]:
+        """
+        Get information about the model.
+        
+        Returns:
+            Dictionary containing model information
+        """
+        return {
+            "name": "llama3-8b",
+            "type": "llm",
+            "device": self.device,
+            "path": self.model_path,
+            "loaded": self._loaded,
+            "default_config": self.default_config
+        } 

isa_model/inference/backends/Pytorch/whisper_backend.py

@@ -0,0 +1,194 @@
+import os
+import io
+import torch
+import logging
+import numpy as np
+from typing import Dict, Any, Optional, Union, BinaryIO
+
+logger = logging.getLogger(__name__)
+
+
+class WhisperBackend:
+    """
+    PyTorch backend for the Whisper speech-to-text model.
+    """
+    
+    def __init__(self, model_path: Optional[str] = None, device: str = "auto"):
+        """
+        Initialize the Whisper backend.
+        
+        Args:
+            model_path: Path to the model
+            device: Device to run the model on ("cpu", "cuda", or "auto")
+        """
+        self.model_path = model_path or os.environ.get("WHISPER_MODEL_PATH", "/models/Whisper-tiny")
+        self.device = device if device != "auto" else ("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = None
+        self.processor = None
+        self._loaded = False
+        
+        # Default configuration
+        self.config = {
+            "language": "en",
+            "task": "transcribe",
+            "sampling_rate": 16000,
+            "chunk_length_s": 30,
+            "batch_size": 16
+        }
+        
+        self.logger = logger
+    
+    def load(self) -> None:
+        """
+        Load the model and processor.
+        """
+        if self._loaded:
+            return
+        
+        try:
+            from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq
+            
+            # Load processor
+            self.logger.info(f"Loading Whisper processor from {self.model_path}")
+            self.processor = AutoProcessor.from_pretrained(self.model_path)
+            
+            # Load model
+            self.logger.info(f"Loading Whisper model on {self.device}")
+            if self.device == "cpu":
+                self.model = AutoModelForSpeechSeq2Seq.from_pretrained(
+                    self.model_path,
+                    torch_dtype=torch.float32,
+                    low_cpu_mem_usage=True,
+                    device_map="auto"
+                )
+            else:  # cuda
+                self.model = AutoModelForSpeechSeq2Seq.from_pretrained(
+                    self.model_path,
+                    torch_dtype=torch.float16,  # Use half precision on GPU
+                    device_map="auto"
+                )
+            
+            self.model.eval()
+            self._loaded = True
+            self.logger.info("Whisper model loaded successfully")
+            
+        except Exception as e:
+            self.logger.error(f"Failed to load Whisper model: {str(e)}")
+            raise
+    
+    def unload(self) -> None:
+        """
+        Unload the model and processor.
+        """
+        if not self._loaded:
+            return
+        
+        self.model = None
+        self.processor = None
+        self._loaded = False
+        
+        # Force garbage collection
+        import gc
+        gc.collect()
+        
+        if self.device == "cuda":
+            torch.cuda.empty_cache()
+        
+        self.logger.info("Whisper model unloaded")
+    
+    def transcribe(self, 
+                  audio: Union[np.ndarray, str, BinaryIO, bytes], 
+                  language: str = "en",
+                  **kwargs) -> str:
+        """
+        Transcribe audio to text.
+        
+        Args:
+            audio: Audio input (numpy array, file path, file-like object, or bytes)
+            language: Language code (e.g., "en", "fr")
+            kwargs: Additional keyword arguments to override config
+            
+        Returns:
+            Transcribed text
+        """
+        if not self._loaded:
+            self.load()
+        
+        # Process audio to get numpy array
+        audio_array = self._process_audio_input(audio)
+        
+        # Update config with kwargs
+        config = self.config.copy()
+        config.update(kwargs)
+        config["language"] = language
+        
+        try:
+            # Process audio with processor
+            inputs = self.processor(
+                audio_array, 
+                sampling_rate=config["sampling_rate"],
+                return_tensors="pt"
+            ).to(self.device)
+            
+            # Generate transcription
+            with torch.no_grad():
+                output = self.model.generate(
+                    **inputs,
+                    language=config["language"],
+                    task=config["task"]
+                )
+            
+            # Decode the output
+            transcription = self.processor.batch_decode(
+                output, 
+                skip_special_tokens=True
+            )[0]
+            
+            return transcription
+            
+        except Exception as e:
+            self.logger.error(f"Error during Whisper transcription: {str(e)}")
+            raise
+    
+    def _process_audio_input(self, audio: Union[np.ndarray, str, BinaryIO, bytes]) -> np.ndarray:
+        """
+        Process different types of audio inputs into a numpy array.
+        
+        Args:
+            audio: Audio input (numpy array, file path, file-like object, or bytes)
+            
+        Returns:
+            Numpy array of the audio
+        """
+        if isinstance(audio, np.ndarray):
+            return audio
+        
+        try:
+            import librosa
+            
+            if isinstance(audio, str):
+                # File path
+                y, sr = librosa.load(audio, sr=self.config["sampling_rate"])
+                return y
+            
+            elif isinstance(audio, (io.IOBase, BinaryIO)):
+                # File-like object
+                audio.seek(0)
+                y, sr = librosa.load(audio, sr=self.config["sampling_rate"])
+                return y
+            
+            elif isinstance(audio, bytes):
+                # Bytes
+                with io.BytesIO(audio) as audio_bytes:
+                    y, sr = librosa.load(audio_bytes, sr=self.config["sampling_rate"])
+                return y
+            
+            else:
+                raise ValueError(f"Unsupported audio type: {type(audio)}")
+                
+        except ImportError:
+            self.logger.error("librosa not installed. Please install with: pip install librosa")
+            raise
+        except Exception as e:
+            self.logger.error(f"Error processing audio: {str(e)}")
+            raise 

isa_model/inference/backends/__init__.py

@@ -0,0 +1,53 @@
+"""
+Backend services for isa_model inference.
+
+Three types of backend services:
+1. Local Services: Services running locally (e.g., Ollama)
+2. Container Services: Docker/K8s deployed services (e.g., Triton, vLLM)  
+3. Third-party Services: External API services with wrappers
+"""
+
+from .base_backend_client import BaseBackendClient
+from .triton_client import TritonBackendClient, TritonClient
+
+# Local Services
+from .local_services import OllamaBackendClient, LocalModelServerClient
+
+# Container Services  
+from .container_services import (
+    VLLMBackendClient, 
+    TensorFlowServingClient, 
+    KubernetesServiceClient
+)
+
+# Third-party Services
+from .third_party_services import (
+    OpenAIClient,
+    AnthropicClient, 
+    CohereClient,
+    AzureOpenAIClient,
+    GoogleAIClient
+)
+
+__all__ = [
+    # Base
+    "BaseBackendClient",
+    "TritonBackendClient",
+    "TritonClient",  # Backward compatibility
+    
+    # Local Services
+    "OllamaBackendClient",
+    "LocalModelServerClient",
+    
+    # Container Services
+    "VLLMBackendClient",
+    "TensorFlowServingClient", 
+    "KubernetesServiceClient",
+    
+    # Third-party Services
+    "OpenAIClient",
+    "AnthropicClient",
+    "CohereClient", 
+    "AzureOpenAIClient",
+    "GoogleAIClient",
+] 

isa_model/inference/backends/base_backend_client.py

@@ -0,0 +1,26 @@
+"""
+Base backend client interface for all AI service backends.
+Defines the common interface that all backend clients must implement.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Dict, Any, AsyncGenerator, Optional
+
+
+class BaseBackendClient(ABC):
+    """Abstract base class for all backend clients"""
+    
+    def __init__(self, *args, **kwargs):
+        """Initialize backend client"""
+        pass
+        
+    @abstractmethod
+    async def health_check(self) -> bool:
+        """Check if the backend service is healthy"""
+        pass
+        
+    async def close(self):
+        """Close any open connections"""
+        pass 
+
+