isa-model 0.1.0__py3-none-any.whl

isa_model/inference/engine/triton/model_repository/bge/1/model.py

@@ -0,0 +1,174 @@
+import os
+import json
+import numpy as np
+import triton_python_backend_utils as pb_utils
+import sys
+import logging
+import torch
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("bge_triton_model")
+
+class TritonPythonModel:
+    """
+    Python model for BGE embedding.
+    """
+    
+    def initialize(self, args):
+        """
+        Initialize the model.
+        """
+        self.model_config = json.loads(args['model_config'])
+        self.model_path = os.environ.get("BGE_MODEL_PATH", "/models/Bge-m3")
+        # Always use CPU for testing
+        self.device = "cpu"
+        self.model = None
+        self.tokenizer = None
+        self._loaded = False
+        
+        # Default configuration
+        self.config = {
+            "normalize": True,
+            "max_length": 512,
+            "pooling_method": "cls"  # Use CLS token for sentence embedding
+        }
+        
+        self._load_model()
+        
+        logger.info(f"Initialized BGE embedding model on {self.device}")
+    
+    def _load_model(self):
+        """Load the BGE model and tokenizer"""
+        if self._loaded:
+            return
+            
+        try:
+            from transformers import AutoModel, AutoTokenizer
+            
+            # Load tokenizer
+            logger.info(f"Loading BGE tokenizer from {self.model_path}")
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
+            
+            # Load model on CPU for testing
+            logger.info(f"Loading BGE model on {self.device}")
+            self.model = AutoModel.from_pretrained(
+                self.model_path,
+                torch_dtype=torch.float32,
+                device_map="cpu"  # Force CPU
+            )
+            
+            self.model.eval()
+            self._loaded = True
+            logger.info("BGE model loaded successfully")
+            
+        except Exception as e:
+            logger.error(f"Failed to load BGE model: {str(e)}")
+            # Fall back to dummy model for testing if model loading fails
+            self._loaded = False
+    
+    def _generate_embeddings(self, texts, normalize=True):
+        """Generate embeddings for the given texts"""
+        if not self._loaded:
+            # Return random embeddings for testing
+            logger.warning("Model not loaded, returning random embeddings")
+            # Generate random embeddings with dimension 1024 (typical for BGE)
+            num_texts = len(texts) if isinstance(texts, list) else 1
+            return np.random.randn(num_texts, 1024).astype(np.float32)
+            
+        try:
+            # Ensure texts is a list
+            if isinstance(texts, str):
+                texts = [texts]
+                
+            # Tokenize the texts
+            inputs = self.tokenizer(
+                texts,
+                padding=True,
+                truncation=True,
+                max_length=self.config["max_length"],
+                return_tensors="pt"
+            ).to(self.device)
+            
+            # Generate embeddings
+            with torch.no_grad():
+                outputs = self.model(**inputs)
+                
+                # Use [CLS] token embedding as the sentence embedding
+                embeddings = outputs.last_hidden_state[:, 0, :]
+                
+                # Normalize embeddings if required
+                if normalize:
+                    embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
+            
+            # Convert to numpy array
+            embeddings_np = embeddings.cpu().numpy()
+            
+            return embeddings_np
+            
+        except Exception as e:
+            logger.error(f"Error during BGE embedding generation: {str(e)}")
+            # Return random embeddings in case of error
+            num_texts = len(texts) if isinstance(texts, list) else 1
+            return np.random.randn(num_texts, 1024).astype(np.float32)
+    
+    def execute(self, requests):
+        """
+        Process inference requests.
+        """
+        responses = []
+        
+        for request in requests:
+            # Get input tensors
+            input_tensor = pb_utils.get_input_tensor_by_name(request, "text_input")
+            
+            # Get texts from input tensor
+            if input_tensor is not None:
+                input_data = input_tensor.as_numpy()
+                texts = []
+                
+                # Convert bytes to strings
+                for i in range(len(input_data)):
+                    if input_data[i].dtype.type is np.bytes_:
+                        texts.append(input_data[i].decode('utf-8'))
+                    else:
+                        texts.append(str(input_data[i]))
+                
+                # Generate embeddings
+                embeddings = self._generate_embeddings(texts)
+                
+                # Create output tensor
+                output_tensor = pb_utils.Tensor(
+                    "embedding_output",
+                    embeddings.astype(np.float32)
+                )
+            else:
+                # If no input is provided, return empty tensor
+                output_tensor = pb_utils.Tensor(
+                    "embedding_output",
+                    np.array([], dtype=np.float32)
+                )
+            
+            # Create inference response
+            inference_response = pb_utils.InferenceResponse(output_tensors=[output_tensor])
+            responses.append(inference_response)
+        
+        return responses
+    
+    def finalize(self):
+        """
+        Clean up resources when the model is unloaded.
+        """
+        if self.model is not None:
+            self.model = None
+            self.tokenizer = None
+            self._loaded = False
+            
+            # Force garbage collection
+            import gc
+            gc.collect()
+            
+            if self.device == "cuda":
+                torch.cuda.empty_cache()
+                
+        logger.info("BGE embedding model unloaded") 

isa_model/inference/engine/triton/model_repository/gemma/1/model.py

@@ -0,0 +1,250 @@
+import os
+import json
+import numpy as np
+import triton_python_backend_utils as pb_utils
+import sys
+import base64
+import torch
+from PIL import Image
+import io
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("gemma_triton_model")
+
+class TritonPythonModel:
+    """
+    Python model for Gemma Vision LLM (VLM).
+    """
+    
+    def initialize(self, args):
+        """
+        Initialize the model.
+        """
+        self.model_config = json.loads(args['model_config'])
+        
+        # Try different possible model paths
+        possible_paths = [
+            "/models/Gemma3-4B",  # Original path
+            "/models/gemma",      # Alternative path
+        ]
+        
+        # Find the first path that exists
+        self.model_path = None
+        for path in possible_paths:
+            if os.path.exists(path):
+                self.model_path = path
+                logger.info(f"Found model at path: {path}")
+                break
+                
+        if self.model_path is None:
+            logger.error("Could not find model path!")
+            self.model_path = "/models/Gemma3-4B"  # Default, will fail later
+            
+        self.device = "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
+        }
+        
+        logger.info(f"Initializing Gemma Vision model at {self.model_path} on {self.device}")
+        self._load_model()
+        
+        if self._loaded:
+            logger.info("Gemma Vision model initialized successfully")
+        else:
+            logger.error("Failed to initialize Gemma Vision model")
+    
+    def _load_model(self):
+        """Load the Gemma model and tokenizer"""
+        if self._loaded:
+            return
+            
+        try:
+            from transformers import AutoModelForCausalLM, AutoTokenizer
+            
+            # Log environment information
+            logger.info(f"Possible model paths:")
+            logger.info(f"Current dir: {os.getcwd()}")
+            logger.info(f"Model path exists: {os.path.exists(self.model_path)}")
+            logger.info(f"Directory listing of /models:")
+            if os.path.exists("/models"):
+                logger.info(", ".join(os.listdir("/models")))
+            
+            # Load tokenizer
+            logger.info(f"Loading Gemma tokenizer from {self.model_path}")
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
+            
+            # Load model
+            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
+            logger.info("Gemma model loaded successfully")
+            
+        except Exception as e:
+            logger.error(f"Failed to load Gemma model: {str(e)}")
+            # Fall back to dummy model for testing if model loading fails
+            self._loaded = False
+    
+    def _process_image(self, image_data):
+        """Process base64 image data"""
+        try:
+            # Extract the base64 part if it's a data URL
+            if isinstance(image_data, str) and image_data.startswith("data:image"):
+                # Extract the base64 part
+                image_data = image_data.split(",")[1]
+                
+            # Decode base64
+            image_bytes = base64.b64decode(image_data)
+            
+            # Open as PIL Image
+            image = Image.open(io.BytesIO(image_bytes))
+            
+            # Process for model input if needed
+            # For now, we're just returning the image for text description
+            return image
+            
+        except Exception as e:
+            logger.error(f"Error processing image: {str(e)}")
+            return None
+    
+    def _generate_text(self, prompt, system_prompt=None, generation_config=None):
+        """Generate text using the Gemma model"""
+        if not self._loaded:
+            return "Model not loaded. Using fallback response: I can see an image but cannot analyze it properly as the vision model is not loaded."
+        
+        try:
+            # Get generation config
+            config = self.default_config.copy()
+            if generation_config:
+                config.update(generation_config)
+            
+            # 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:
+            logger.error(f"Error during Gemma text generation: {str(e)}")
+            return f"Error generating response: {str(e)}"
+    
+    def execute(self, requests):
+        """
+        Process inference requests.
+        """
+        responses = []
+        
+        for request in requests:
+            # Get input tensors
+            input_tensor = pb_utils.get_input_tensor_by_name(request, "prompt")
+            
+            # Convert to string
+            if input_tensor is not None:
+                input_data = input_tensor.as_numpy()
+                if input_data.dtype.type is np.bytes_:
+                    input_text = input_data[0][0].decode('utf-8')
+                else:
+                    input_text = str(input_data[0][0])
+                
+                # Check if the input contains an image (base64)
+                if "data:image" in input_text:
+                    # Extract image description query
+                    query = "Describe this image in detail."
+                    if "?" in input_text:
+                        parts = input_text.split("?")
+                        query = parts[0] + "?"
+                    
+                    # For image inputs
+                    if self._loaded:
+                        response_text = self._generate_text(input_text)
+                    else:
+                        # Fallback if model not loaded
+                        response_text = "Model not loaded. Using fallback response: I can see an image but cannot analyze it properly as the vision model is not loaded."
+                else:
+                    # For text-only prompts
+                    system_prompt_tensor = pb_utils.get_input_tensor_by_name(request, "system_prompt")
+                    system_prompt = None
+                    if system_prompt_tensor is not None:
+                        system_prompt_data = system_prompt_tensor.as_numpy()
+                        if system_prompt_data.dtype.type is np.bytes_:
+                            system_prompt = system_prompt_data[0].decode('utf-8')
+                    
+                    response_text = self._generate_text(input_text, system_prompt)
+            else:
+                response_text = "No input provided"
+            
+            # Create output tensor
+            output_tensor = pb_utils.Tensor(
+                "text_output",
+                np.array([[response_text]], dtype=np.object_)
+            )
+            
+            # Create inference response
+            inference_response = pb_utils.InferenceResponse(output_tensors=[output_tensor])
+            responses.append(inference_response)
+        
+        return responses
+    
+    def finalize(self):
+        """
+        Clean up resources when the model is unloaded.
+        """
+        if self.model is not None:
+            self.model = None
+            self.tokenizer = None
+            self._loaded = False
+            
+            # Force garbage collection
+            import gc
+            gc.collect()
+            
+            if self.device == "cuda":
+                torch.cuda.empty_cache()
+                
+        logger.info("Gemma Vision model unloaded") 

isa_model/inference/engine/triton/model_repository/llama/1/model.py

@@ -0,0 +1,76 @@
+import json
+import numpy as np
+import triton_python_backend_utils as pb_utils
+import random
+
+class TritonPythonModel:
+    """
+    Simulated Llama model for testing.
+    """
+    
+    def initialize(self, args):
+        """
+        Initialize the model.
+        """
+        self.model_config = json.loads(args['model_config'])
+        self.responses = {
+            "artificial intelligence": "Artificial Intelligence (AI) refers to the simulation of human intelligence in machines that are programmed to think and learn like humans. AI encompasses various subfields including machine learning, natural language processing, computer vision, and robotics. Modern AI systems can perform tasks such as understanding natural language, recognizing images, making decisions, and solving complex problems.",
+            "language model": "A language model is a type of artificial intelligence model that's trained to understand and generate human language. Large Language Models (LLMs) like myself are trained on vast amounts of text data to predict the next word in a sequence, enabling them to generate coherent and contextually relevant text, answer questions, translate languages, and perform various text-based tasks.",
+            "machine learning": "Machine learning is a subset of artificial intelligence that enables systems to learn and improve from experience without being explicitly programmed. It focuses on developing algorithms that can access data, learn from it, and make predictions or decisions. Common types include supervised learning, unsupervised learning, and reinforcement learning.",
+            "default": "As an AI language model, I'm designed to process and generate human-like text based on the input I receive. I can assist with a wide range of tasks including answering questions, providing explanations, generating creative content, and engaging in conversations on various topics. How else can I help you today?"
+        }
+        print(f"Initialized Simulated Llama model")
+    
+    def execute(self, requests):
+        """
+        Process inference requests.
+        """
+        responses = []
+        
+        for request in requests:
+            # Get input tensors
+            prompt = pb_utils.get_input_tensor_by_name(request, "prompt")
+            prompt_data = prompt.as_numpy()
+            
+            # Decode prompt from bytes
+            if prompt_data.dtype == np.object_:
+                prompt_str = prompt_data[0][0].decode('utf-8')
+            else:
+                prompt_str = prompt_data[0][0]
+            
+            # Generate a relevant response based on the prompt
+            generated_text = self._generate_response(prompt_str)
+            
+            # Create output tensor
+            output_tensor = pb_utils.Tensor(
+                "text_output", 
+                np.array([[generated_text]], dtype=np.object_)
+            )
+            
+            # Create and append response
+            inference_response = pb_utils.InferenceResponse(
+                output_tensors=[output_tensor]
+            )
+            responses.append(inference_response)
+        
+        return responses
+    
+    def _generate_response(self, prompt):
+        """
+        Generate a response based on the prompt keywords.
+        """
+        prompt_lower = prompt.lower()
+        
+        # Check for keywords in the prompt
+        for key in self.responses:
+            if key in prompt_lower:
+                return self.responses[key]
+        
+        # If no keywords match, return the default response
+        return self.responses["default"]
+    
+    def finalize(self):
+        """
+        Clean up resources when the model is unloaded.
+        """
+        print("Simulated Llama model unloaded") 

isa_model/inference/engine/triton/model_repository/whisper/1/model.py

@@ -0,0 +1,195 @@
+import json
+import numpy as np
+import triton_python_backend_utils as pb_utils
+import os
+import sys
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("whisper_triton_model")
+
+class TritonPythonModel:
+    """
+    Python model for Whisper speech-to-text using a simplified approach.
+    """
+    
+    def initialize(self, args):
+        """
+        Initialize the model.
+        """
+        self.model_config = json.loads(args['model_config'])
+        
+        # Get model name from config
+        self.model_name = "/models/Whisper-tiny"
+        if 'parameters' in self.model_config:
+            parameters = self.model_config['parameters']
+            if 'model_name' in parameters:
+                self.model_name = parameters['model_name']['string_value']
+        
+        logger.info(f"Initializing simplified Whisper model: {self.model_name}")
+        
+        # This is a simple mock model for testing
+        # In production, you would use an actual Whisper model
+        self.languages = {
+            "en": "English",
+            "fr": "French",
+            "es": "Spanish",
+            "de": "German",
+            "zh": "Chinese",
+            "ja": "Japanese"
+        }
+        
+        # Try loading the model
+        try:
+            from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq
+            import torch
+            
+            logger.info(f"Attempting to load Whisper model from {self.model_name}")
+            self.processor = AutoProcessor.from_pretrained(self.model_name)
+            self.model = AutoModelForSpeechSeq2Seq.from_pretrained(
+                self.model_name,
+                torch_dtype=torch.float32,
+                device_map="cpu"
+            )
+            self.model.eval()
+            self.model_loaded = True
+            logger.info("Whisper model loaded successfully")
+        except Exception as e:
+            logger.warning(f"Failed to load Whisper model: {e}")
+            self.model_loaded = False
+            logger.info("Using fallback mock transcription")
+        
+        logger.info("Simplified Whisper model initialized successfully")
+    
+    def execute(self, requests):
+        """
+        Process inference requests.
+        """
+        responses = []
+        
+        for request in requests:
+            try:
+                # Get input tensors
+                audio_input = pb_utils.get_input_tensor_by_name(request, "audio_input")
+                language_input = pb_utils.get_input_tensor_by_name(request, "language")
+                
+                # Get language or use default
+                language = "en"
+                if language_input is not None:
+                    # Fix for decoding language input
+                    lang_np = language_input.as_numpy()
+                    if lang_np.dtype.type is np.bytes_:
+                        language = lang_np[0][0].decode('utf-8')
+                    elif lang_np.dtype.type is np.object_:
+                        language = str(lang_np[0][0])
+                    else:
+                        language = str(lang_np[0][0])
+                
+                # Process audio input
+                if audio_input is not None:
+                    audio_data = audio_input.as_numpy()
+                    
+                    # Handle input shape [1, -1]
+                    logger.info(f"Original audio data shape: {audio_data.shape}")
+                    
+                    # If the model is loaded, use it for transcription
+                    if hasattr(self, 'model_loaded') and self.model_loaded:
+                        try:
+                            import torch
+                            
+                            # Reshape if needed
+                            if len(audio_data.shape) > 2:  # [batch, 1, length]
+                                audio_data = audio_data.reshape(audio_data.shape[0], -1)
+                                
+                            # Process audio with transformers
+                            inputs = self.processor(
+                                audio_data,
+                                sampling_rate=16000,
+                                return_tensors="pt"
+                            )
+                            
+                            # Generate transcription
+                            with torch.no_grad():
+                                generated_ids = self.model.generate(
+                                    inputs.input_features,
+                                    language=language,
+                                    task="transcribe"
+                                )
+                                
+                            # Decode transcription
+                            transcription = self.processor.batch_decode(
+                                generated_ids,
+                                skip_special_tokens=True
+                            )[0]
+                            
+                            logger.info(f"Generated transcription using model: {transcription}")
+                        except Exception as e:
+                            logger.error(f"Error using loaded model: {e}")
+                            # Fall back to mock transcription
+                            audio_length = audio_data.size
+                            transcription = self._get_mock_transcription(audio_length, language)
+                    else:
+                        # Use mock transcription
+                        audio_length = audio_data.size
+                        transcription = self._get_mock_transcription(audio_length, language)
+                    
+                else:
+                    transcription = "No audio input provided."
+                
+                # Create output tensor
+                transcription_tensor = pb_utils.Tensor(
+                    "text_output", 
+                    np.array([transcription], dtype=np.object_)
+                )
+                
+                # Create and append response
+                inference_response = pb_utils.InferenceResponse(
+                    output_tensors=[transcription_tensor]
+                )
+                responses.append(inference_response)
+                
+            except Exception as e:
+                logger.error(f"Error processing request: {e}")
+                # Return error response
+                error_response = pb_utils.InferenceResponse(
+                    output_tensors=[
+                        pb_utils.Tensor(
+                            "text_output",
+                            np.array([f"Error: {str(e)}"], dtype=np.object_)
+                        )
+                    ]
+                )
+                responses.append(error_response)
+        
+        return responses
+    
+    def _get_mock_transcription(self, audio_length, language):
+        """Generate a mock transcription based on audio length"""
+        if audio_length > 100000:
+            return f"This is a test transcription in {self.languages.get(language, 'English')}. The audio is quite long with {audio_length} samples."
+        elif audio_length > 50000:
+            return f"This is a medium length test transcription in {self.languages.get(language, 'English')}."
+        else:
+            return f"Short test transcription in {self.languages.get(language, 'English')}."
+    
+    def finalize(self):
+        """
+        Clean up resources when the model is unloaded.
+        """
+        if hasattr(self, 'model') and self.model is not None:
+            self.model = None
+            self.processor = None
+            
+            # Force garbage collection
+            import gc
+            gc.collect()
+            
+            try:
+                import torch
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+            except ImportError:
+                pass
+                
+        logger.info("Whisper model unloaded")