isa-model 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

isa_model/inference/services/llm/gemma_service.py

@@ -1,143 +0,0 @@
-import json
-import logging
-import asyncio
-from typing import Dict, List, Any, Optional, Union
-
-from isa_model.inference.services.base_service import BaseService
-from isa_model.inference.backends.triton_client import TritonClient
-
-logger = logging.getLogger(__name__)
-
-
-class GemmaService(BaseService):
-    """
-    Service for Gemma LLM using Triton Inference Server.
-    """
-    
-    def __init__(self, triton_url: str = "localhost:8001", model_name: str = "gemma"):
-        """
-        Initialize the Gemma service.
-        
-        Args:
-            triton_url: URL of the Triton Inference Server
-            model_name: Name of the model in Triton
-        """
-        super().__init__()
-        self.triton_url = triton_url
-        self.model_name = model_name
-        self.client = None
-        
-        # 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
-    
-    async def load(self) -> None:
-        """
-        Load the client connection to Triton.
-        """
-        if self.is_loaded():
-            return
-        
-        try:
-            # Create Triton client
-            self.logger.info(f"Connecting to Triton server at {self.triton_url}")
-            self.client = TritonClient(self.triton_url)
-            
-            # Check if model is ready
-            if not await self.client.is_model_ready(self.model_name):
-                self.logger.error(f"Model {self.model_name} is not ready on Triton server")
-                raise RuntimeError(f"Model {self.model_name} is not ready on Triton server")
-            
-            self._loaded = True
-            self.logger.info(f"Connected to Triton for model {self.model_name}")
-            
-        except Exception as e:
-            self.logger.error(f"Failed to connect to Triton: {str(e)}")
-            raise
-    
-    async def unload(self) -> None:
-        """
-        Unload the client connection.
-        """
-        if not self.is_loaded():
-            return
-        
-        self.client = None
-        self._loaded = False
-        self.logger.info("Triton client connection closed")
-    
-    async def generate(self, 
-                      prompt: str, 
-                      system_prompt: Optional[str] = None,
-                      generation_config: Optional[Dict[str, Any]] = None) -> str:
-        """
-        Generate text from a prompt using Triton.
-        
-        Args:
-            prompt: User prompt
-            system_prompt: System prompt to control model behavior
-            generation_config: Configuration for text generation
-            
-        Returns:
-            Generated text
-        """
-        if not self.is_loaded():
-            await self.load()
-        
-        # Get configuration
-        merged_config = self.default_config.copy()
-        if generation_config:
-            merged_config.update(generation_config)
-        
-        try:
-            # Prepare inputs
-            inputs = {
-                "prompt": [prompt],
-            }
-            
-            # Add optional inputs
-            if system_prompt:
-                inputs["system_prompt"] = [system_prompt]
-            
-            if merged_config:
-                inputs["generation_config"] = [json.dumps(merged_config)]
-            
-            # Run inference
-            result = await self.client.infer(
-                model_name=self.model_name,
-                inputs=inputs,
-                outputs=["text_output"]
-            )
-            
-            # Extract generated text
-            generated_text = result["text_output"][0].decode('utf-8')
-            
-            return generated_text
-            
-        except Exception as e:
-            self.logger.error(f"Error during 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": self.model_name,
-            "type": "llm",
-            "backend": "triton",
-            "url": self.triton_url,
-            "loaded": self.is_loaded(),
-            "config": self.default_config
-        } 

isa_model/inference/services/llm/llama_service.py

@@ -1,143 +0,0 @@
-import json
-import logging
-import asyncio
-from typing import Dict, List, Any, Optional, Union
-
-from isa_model.inference.services.base_service import BaseService
-from isa_model.inference.backends.triton_client import TritonClient
-
-logger = logging.getLogger(__name__)
-
-
-class LlamaService(BaseService):
-    """
-    Service for Llama LLM using Triton Inference Server.
-    """
-    
-    def __init__(self, triton_url: str = "localhost:8001", model_name: str = "llama"):
-        """
-        Initialize the Llama service.
-        
-        Args:
-            triton_url: URL of the Triton Inference Server
-            model_name: Name of the model in Triton
-        """
-        super().__init__()
-        self.triton_url = triton_url
-        self.model_name = model_name
-        self.client = None
-        
-        # 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
-    
-    async def load(self) -> None:
-        """
-        Load the client connection to Triton.
-        """
-        if self.is_loaded():
-            return
-        
-        try:
-            # Create Triton client
-            self.logger.info(f"Connecting to Triton server at {self.triton_url}")
-            self.client = TritonClient(self.triton_url)
-            
-            # Check if model is ready
-            if not await self.client.is_model_ready(self.model_name):
-                self.logger.error(f"Model {self.model_name} is not ready on Triton server")
-                raise RuntimeError(f"Model {self.model_name} is not ready on Triton server")
-            
-            self._loaded = True
-            self.logger.info(f"Connected to Triton for model {self.model_name}")
-            
-        except Exception as e:
-            self.logger.error(f"Failed to connect to Triton: {str(e)}")
-            raise
-    
-    async def unload(self) -> None:
-        """
-        Unload the client connection.
-        """
-        if not self.is_loaded():
-            return
-        
-        self.client = None
-        self._loaded = False
-        self.logger.info("Triton client connection closed")
-    
-    async def generate(self, 
-                      prompt: str, 
-                      system_prompt: Optional[str] = None,
-                      generation_config: Optional[Dict[str, Any]] = None) -> str:
-        """
-        Generate text from a prompt using Triton.
-        
-        Args:
-            prompt: User prompt
-            system_prompt: System prompt to control model behavior
-            generation_config: Configuration for text generation
-            
-        Returns:
-            Generated text
-        """
-        if not self.is_loaded():
-            await self.load()
-        
-        # Get configuration
-        merged_config = self.default_config.copy()
-        if generation_config:
-            merged_config.update(generation_config)
-        
-        try:
-            # Prepare inputs
-            inputs = {
-                "prompt": [prompt],
-            }
-            
-            # Add optional inputs
-            if system_prompt:
-                inputs["system_prompt"] = [system_prompt]
-            
-            if merged_config:
-                inputs["generation_config"] = [json.dumps(merged_config)]
-            
-            # Run inference
-            result = await self.client.infer(
-                model_name=self.model_name,
-                inputs=inputs,
-                outputs=["text_output"]
-            )
-            
-            # Extract generated text
-            generated_text = result["text_output"][0].decode('utf-8')
-            
-            return generated_text
-            
-        except Exception as e:
-            self.logger.error(f"Error during 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": self.model_name,
-            "type": "llm",
-            "backend": "triton",
-            "url": self.triton_url,
-            "loaded": self.is_loaded(),
-            "config": self.default_config
-        } 

isa_model/inference/services/llm/replicate_llm_service.py

@@ -1,179 +0,0 @@
-import logging
-from typing import Dict, Any, List, Union, AsyncGenerator, Optional
-from isa_model.inference.services.base_service import BaseLLMService
-from isa_model.inference.providers.base_provider import BaseProvider
-from isa_model.inference.backends.replicate_client import ReplicateBackendClient
-
-logger = logging.getLogger(__name__)
-
-class ReplicateLLMService(BaseLLMService):
-    """Replicate LLM service implementation"""
-    
-    def __init__(self, provider: 'BaseProvider', model_name: str = "meta/llama-3-8b-instruct", backend: Optional[ReplicateBackendClient] = None):
-        super().__init__(provider, model_name)
-        
-        # Use provided backend or create new one
-        if backend:
-            self.backend = backend
-        else:
-            api_token = self.config.get("api_token", "")
-            self.backend = ReplicateBackendClient(api_token)
-            
-        # Parse model name for Replicate format (owner/model)
-        if "/" not in model_name:
-            logger.warning(f"Model name {model_name} is not in Replicate format (owner/model). Using as-is.")
-        
-        # Store version separately if provided
-        self.model_version = None
-        if ":" in model_name:
-            self.model_name, self.model_version = model_name.split(":", 1)
-            
-        self.last_token_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
-        logger.info(f"Initialized ReplicateLLMService with model {model_name}")
-    
-    async def ainvoke(self, prompt: Union[str, List[Dict[str, str]], Any]):
-        """Universal invocation method"""
-        if isinstance(prompt, str):
-            return await self.acompletion(prompt)
-        elif isinstance(prompt, list):
-            return await self.achat(prompt)
-        else:
-            raise ValueError("Prompt must be string or list of messages")
-    
-    async def achat(self, messages: List[Dict[str, str]]):
-        """Chat completion method"""
-        try:
-            temperature = self.config.get("temperature", 0.7)
-            max_tokens = self.config.get("max_tokens", 1024)
-            
-            # Convert to Replicate format
-            prompt = self._convert_messages_to_prompt(messages)
-            
-            # Prepare input data
-            input_data = {
-                "prompt": prompt,
-                "temperature": temperature,
-                "max_new_tokens": max_tokens,
-                "system_prompt": self._extract_system_prompt(messages)
-            }
-            
-            # Call Replicate API
-            prediction = await self.backend.create_prediction(
-                self.model_name, 
-                self.model_version, 
-                input_data
-            )
-            
-            # Get output - could be a list of strings or a single string
-            output = prediction.get("output", "")
-            if isinstance(output, list):
-                output = "".join(output)
-            
-            # Approximate token usage - Replicate doesn't provide token counts
-            approx_prompt_tokens = len(prompt) // 4  # Very rough approximation
-            approx_completion_tokens = len(output) // 4
-            self.last_token_usage = {
-                "prompt_tokens": approx_prompt_tokens,
-                "completion_tokens": approx_completion_tokens,
-                "total_tokens": approx_prompt_tokens + approx_completion_tokens
-            }
-            
-            return output
-            
-        except Exception as e:
-            logger.error(f"Error in chat completion: {e}")
-            raise
-    
-    def _convert_messages_to_prompt(self, messages: List[Dict[str, str]]) -> str:
-        """Convert chat messages to Replicate prompt format"""
-        prompt = ""
-        
-        for msg in messages:
-            role = msg.get("role", "user")
-            content = msg.get("content", "")
-            
-            # Skip system prompts - handled separately
-            if role == "system":
-                continue
-                
-            if role == "user":
-                prompt += f"Human: {content}\n\n"
-            elif role == "assistant":
-                prompt += f"Assistant: {content}\n\n"
-            else:
-                # Default to user for unknown roles
-                prompt += f"Human: {content}\n\n"
-        
-        # Add final assistant prefix for the model to continue
-        prompt += "Assistant: "
-        
-        return prompt
-    
-    def _extract_system_prompt(self, messages: List[Dict[str, str]]) -> str:
-        """Extract system prompt from messages"""
-        for msg in messages:
-            if msg.get("role") == "system":
-                return msg.get("content", "")
-        return ""
-    
-    async def acompletion(self, prompt: str):
-        """Text completion method"""
-        try:
-            # For simple completion, use chat format with a single user message
-            messages = [{"role": "user", "content": prompt}]
-            return await self.achat(messages)
-        except Exception as e:
-            logger.error(f"Error in text completion: {e}")
-            raise
-    
-    async def agenerate(self, messages: List[Dict[str, str]], n: int = 1) -> List[str]:
-        """Generate multiple completions"""
-        # Replicate doesn't support multiple outputs in one call,
-        # so we make multiple calls
-        results = []
-        for _ in range(n):
-            result = await self.achat(messages)
-            results.append(result)
-        return results
-    
-    async def astream_chat(self, messages: List[Dict[str, str]]) -> AsyncGenerator[str, None]:
-        """Stream chat responses"""
-        try:
-            temperature = self.config.get("temperature", 0.7)
-            max_tokens = self.config.get("max_tokens", 1024)
-            
-            # Convert to Replicate format
-            prompt = self._convert_messages_to_prompt(messages)
-            
-            # Prepare input data
-            input_data = {
-                "prompt": prompt,
-                "temperature": temperature,
-                "max_new_tokens": max_tokens,
-                "system_prompt": self._extract_system_prompt(messages),
-                "stream": True
-            }
-            
-            # Call Replicate API with streaming
-            async for chunk in self.backend.stream_prediction(
-                self.model_name, 
-                self.model_version, 
-                input_data
-            ):
-                yield chunk
-                
-        except Exception as e:
-            logger.error(f"Error in stream chat: {e}")
-            raise
-    
-    def get_token_usage(self):
-        """Get total token usage statistics"""
-        return self.last_token_usage
-    
-    def get_last_token_usage(self) -> Dict[str, int]:
-        """Get token usage from last request"""
-        return self.last_token_usage
-        
-    async def close(self):
-        """Close the backend client"""
-        await self.backend.close() 

isa_model/inference/services/llm/triton_llm_service.py

@@ -1,230 +0,0 @@
-import json
-import logging
-import asyncio
-from typing import Dict, List, Any, AsyncGenerator, Optional, Union
-
-import numpy as np
-
-from isa_model.inference.services.base_service import BaseLLMService
-from isa_model.inference.providers.triton_provider import TritonProvider
-
-logger = logging.getLogger(__name__)
-
-
-class TritonLLMService(BaseLLMService):
-    """
-    LLM service that uses Triton Inference Server to run inference.
-    """
-    
-    def __init__(self, provider: TritonProvider, model_name: str):
-        """
-        Initialize the Triton LLM service.
-        
-        Args:
-            provider: The Triton provider
-            model_name: Name of the model in Triton (e.g., "Llama3-8B")
-        """
-        super().__init__(provider, model_name)
-        self.client = None
-        self.tokenizer = None
-        self.token_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
-        self.last_token_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
-        
-    async def _initialize_client(self):
-        """Initialize the Triton client"""
-        if self.client is None:
-            self.client = self.provider.create_client()
-            
-            # Check if model is ready
-            if not self.provider.is_model_ready(self.model_name):
-                logger.error(f"Model {self.model_name} is not ready on Triton server")
-                raise RuntimeError(f"Model {self.model_name} is not ready on Triton server")
-                
-            logger.info(f"Initialized Triton client for model: {self.model_name}")
-    
-    async def ainvoke(self, prompt: Union[str, List[Dict[str, str]], Any]) -> str:
-        """
-        Universal invocation method.
-        
-        Args:
-            prompt: Text prompt or chat messages
-        
-        Returns:
-            Generated text
-        """
-        if isinstance(prompt, str):
-            return await self.acompletion(prompt)
-        elif isinstance(prompt, list) and all(isinstance(m, dict) for m in prompt):
-            return await self.achat(prompt)
-        else:
-            raise ValueError("Prompt must be either a string or a list of message dictionaries")
-    
-    async def achat(self, messages: List[Dict[str, str]]) -> str:
-        """
-        Chat completion method.
-        
-        Args:
-            messages: List of message dictionaries with 'role' and 'content'
-            
-        Returns:
-            Generated chat response
-        """
-        # Format chat messages into a single prompt
-        formatted_prompt = self._format_chat_messages(messages)
-        return await self.acompletion(formatted_prompt)
-    
-    async def acompletion(self, prompt: str) -> str:
-        """
-        Text completion method.
-        
-        Args:
-            prompt: Text prompt
-            
-        Returns:
-            Generated text completion
-        """
-        await self._initialize_client()
-        
-        try:
-            import tritonclient.http as httpclient
-            
-            # Create input tensors
-            input_text = np.array([prompt], dtype=np.object_)
-            inputs = [httpclient.InferInput("TEXT", input_text.shape, "BYTES")]
-            inputs[0].set_data_from_numpy(input_text)
-            
-            # Default parameters
-            generation_params = {
-                "max_new_tokens": self.config.get("max_new_tokens", 512),
-                "temperature": self.config.get("temperature", 0.7),
-                "top_p": self.config.get("top_p", 0.9),
-                "do_sample": self.config.get("temperature", 0.7) > 0
-            }
-            
-            # Add parameters as input tensor
-            param_json = json.dumps(generation_params)
-            param_data = np.array([param_json], dtype=np.object_)
-            param_input = httpclient.InferInput("PARAMETERS", param_data.shape, "BYTES")
-            param_input.set_data_from_numpy(param_data)
-            inputs.append(param_input)
-            
-            # Create output tensor
-            outputs = [httpclient.InferRequestedOutput("TEXT")]
-            
-            # Send the request
-            response = await asyncio.to_thread(
-                self.client.infer,
-                self.model_name,
-                inputs,
-                outputs=outputs
-            )
-            
-            # Process the response
-            output = response.as_numpy("TEXT")
-            response_text = output[0].decode('utf-8')
-            
-            # Update token usage (estimated since we don't have actual token counts)
-            prompt_tokens = len(prompt) // 4  # Rough estimate
-            completion_tokens = len(response_text) // 4  # Rough estimate
-            total_tokens = prompt_tokens + completion_tokens
-            
-            self.last_token_usage = {
-                "prompt_tokens": prompt_tokens,
-                "completion_tokens": completion_tokens,
-                "total_tokens": total_tokens
-            }
-            
-            # Update total token usage
-            self.token_usage["prompt_tokens"] += prompt_tokens
-            self.token_usage["completion_tokens"] += completion_tokens
-            self.token_usage["total_tokens"] += total_tokens
-            
-            return response_text
-            
-        except Exception as e:
-            logger.error(f"Error during Triton inference: {str(e)}")
-            raise
-    
-    async def agenerate(self, messages: List[Dict[str, str]], n: int = 1) -> List[str]:
-        """
-        Generate multiple completions.
-        
-        Args:
-            messages: List of message dictionaries
-            n: Number of completions to generate
-            
-        Returns:
-            List of generated completions
-        """
-        results = []
-        for _ in range(n):
-            result = await self.achat(messages)
-            results.append(result)
-        return results
-    
-    async def astream_chat(self, messages: List[Dict[str, str]]) -> AsyncGenerator[str, None]:
-        """
-        Stream chat responses.
-        
-        Args:
-            messages: List of message dictionaries
-            
-        Yields:
-            Generated text chunks
-        """
-        # For Triton, we don't have true streaming, so we generate the full response
-        # and then simulate streaming
-        full_response = await self.achat(messages)
-        
-        # Simulate streaming by yielding words
-        words = full_response.split()
-        for i in range(len(words)):
-            chunk = ' '.join(words[:i+1])
-            yield chunk
-            await asyncio.sleep(0.05)  # Small delay to simulate streaming
-    
-    def get_token_usage(self) -> Dict[str, int]:
-        """
-        Get total token usage statistics.
-        
-        Returns:
-            Dictionary with token usage statistics
-        """
-        return self.token_usage
-    
-    def get_last_token_usage(self) -> Dict[str, int]:
-        """
-        Get token usage from last request.
-        
-        Returns:
-            Dictionary with token usage statistics from last request
-        """
-        return self.last_token_usage
-    
-    def _format_chat_messages(self, messages: List[Dict[str, str]]) -> str:
-        """
-        Format chat messages into a single prompt for models that don't support chat natively.
-        
-        Args:
-            messages: List of message dictionaries
-            
-        Returns:
-            Formatted prompt
-        """
-        formatted_prompt = ""
-        
-        for message in messages:
-            role = message.get("role", "user").lower()
-            content = message.get("content", "")
-            
-            if role == "system":
-                formatted_prompt += f"System: {content}\n\n"
-            elif role == "user":
-                formatted_prompt += f"User: {content}\n\n"
-            elif role == "assistant":
-                formatted_prompt += f"Assistant: {content}\n\n"
-            else:
-                formatted_prompt += f"{role.capitalize()}: {content}\n\n"
-        
-        formatted_prompt += "Assistant: "
-        return formatted_prompt