isa-model 0.1.0__py3-none-any.whl

isa_model/inference/services/embedding/ollama_rerank_service.py

@@ -0,0 +1,118 @@
+from typing import Dict, Any, List, Optional
+from ollama import AsyncClient
+from ...base_service import BaseRerankService
+from ...base_provider import BaseProvider
+from app.config.config_manager import config_manager
+import httpx
+import asyncio
+from functools import wraps
+
+logger = config_manager.get_logger(__name__)
+
+def retry_on_connection_error(max_retries=3, delay=1):
+    """Decorator to retry on connection errors"""
+    def decorator(func):
+        @wraps(func)
+        async def wrapper(*args, **kwargs):
+            last_error = None
+            for attempt in range(max_retries):
+                try:
+                    return await func(*args, **kwargs)
+                except (httpx.RemoteProtocolError, httpx.ConnectError) as e:
+                    last_error = e
+                    if attempt < max_retries - 1:
+                        logger.warning(f"Connection error on attempt {attempt + 1}, retrying in {delay}s: {str(e)}")
+                        await asyncio.sleep(delay)
+                    continue
+            raise last_error
+        return wrapper
+    return decorator
+
+class OllamaRerankService(BaseRerankService):
+    """Reranking service wrapper around Ollama"""
+    
+    def __init__(self, provider: 'BaseProvider', model_name: str):
+        super().__init__(provider, model_name)
+        
+        # Initialize the Ollama client for reranking
+        self.client = AsyncClient(
+            host=self.config.get('base_url', 'http://localhost:11434')
+        )
+        self.model_name = model_name
+    
+    @retry_on_connection_error()
+    async def rerank(
+        self,
+        query: str,
+        documents: List[Dict],
+        top_k: int = 5
+    ) -> List[Dict]:
+        """Rerank documents based on query relevance"""
+        try:
+            if not query:
+                raise ValueError("Query cannot be empty")
+            if not documents:
+                return []
+                
+            results = []
+            for doc in documents:
+                if "content" not in doc:
+                    raise ValueError("Each document must have a 'content' field")
+                    
+                # Format prompt for relevance scoring
+                prompt = f"""Rate the relevance of the following text to the query on a scale of 0-100.
+Query: {query}
+Text: {doc['content']}
+Only respond with a number between 0 and 100."""
+
+                # Get relevance score using direct Ollama API
+                response = await self.client.generate(
+                    model=self.model_name,
+                    prompt=prompt,
+                    stream=False
+                )
+                try:
+                    score = float(response.response.strip())
+                    score = max(0.0, min(100.0, score)) / 100.0  # Normalize to 0-1
+                except ValueError:
+                    logger.warning(f"Could not parse score from response: {response.response}")
+                    score = 0.0
+                    
+                # Update document with rerank score
+                doc_copy = doc.copy()
+                doc_copy["rerank_score"] = score
+                doc_copy["final_score"] = doc.get("score", 1.0) * score
+                results.append(doc_copy)
+            
+            # Sort by final score in descending order
+            results.sort(key=lambda x: x["final_score"], reverse=True)
+            return results[:top_k]
+            
+        except Exception as e:
+            logger.error(f"Error in rerank: {e}")
+            raise
+    
+    @retry_on_connection_error()
+    async def rerank_texts(
+        self,
+        query: str,
+        texts: List[str]
+    ) -> List[Dict]:
+        """Rerank raw texts based on query relevance"""
+        try:
+            if not query:
+                raise ValueError("Query cannot be empty")
+            if not texts:
+                return []
+                
+            # Convert texts to document format
+            documents = [{"content": text, "score": 1.0} for text in texts]
+            return await self.rerank(query, documents)
+            
+        except Exception as e:
+            logger.error(f"Error in rerank_texts: {str(e)}")
+            raise
+    
+    async def close(self):
+        """Cleanup resources"""
+        await self.client.aclose()

isa_model/inference/services/embedding/onnx_rerank_service.py

@@ -0,0 +1,73 @@
+from typing import Dict, Any, List, Union, Optional
+from ...base_service import BaseService
+from ...base_provider import BaseProvider
+from transformers import AutoTokenizer
+import onnxruntime as ort
+import numpy as np
+import torch
+import os
+from pathlib import Path
+
+class ONNXRerankService(BaseService):
+    """ONNX Reranker service for BGE models"""
+    
+    def __init__(self, provider: 'BaseProvider', model_name: str):
+        super().__init__(provider, model_name)
+        self.model_path = self._get_model_path(model_name)
+        self.session = provider.get_session(self.model_path)
+        
+        # Initialize tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-v2-m3')
+        self.max_length = 512
+    
+    def _get_model_path(self, model_name: str) -> str:
+        """Get path to ONNX model file"""
+        base_dir = Path(__file__).parent
+        model_path = base_dir / "model_converted" / model_name / "model.onnx"
+        if not model_path.exists():
+            raise FileNotFoundError(f"ONNX model not found at {model_path}. Please run the conversion script first.")
+        return str(model_path)
+    
+    async def compute_score(self, 
+                          pairs: Union[List[str], List[List[str]]], 
+                          normalize: bool = False) -> Union[float, List[float]]:
+        """Compute reranking scores for query-passage pairs"""
+        try:
+            # Handle single pair case
+            if isinstance(pairs[0], str):
+                pairs = [pairs]
+            
+            # Tokenize inputs
+            inputs = self.tokenizer(
+                pairs,
+                padding=True,
+                truncation=True,
+                return_tensors='np',
+                max_length=self.max_length
+            )
+            
+            # Run inference
+            ort_inputs = {
+                'input_ids': inputs['input_ids'],
+                'attention_mask': inputs['attention_mask']
+            }
+            
+            scores = self.session.run(
+                None,  # output names, None means all
+                ort_inputs
+            )[0]
+            
+            # Convert to float and optionally normalize
+            scores = scores.flatten().tolist()
+            if normalize:
+                scores = [self._sigmoid(score) for score in scores]
+            
+            # Return single score for single pair
+            return scores[0] if len(scores) == 1 else scores
+            
+        except Exception as e:
+            raise RuntimeError(f"ONNX reranking failed: {e}")
+    
+    def _sigmoid(self, x: float) -> float:
+        """Apply sigmoid function to score"""
+        return 1 / (1 + np.exp(-x))

isa_model/inference/services/llm/__init__.py

@@ -0,0 +1,16 @@
+"""
+LLM Services - Business logic services for Language Models
+"""
+
+# Import LLM services here when created
+from .ollama_llm_service import OllamaLLMService
+from .triton_llm_service import TritonLLMService
+from .openai_llm_service import OpenAILLMService
+from .replicate_llm_service import ReplicateLLMService
+
+__all__ = [
+    "OllamaLLMService",
+    "TritonLLMService",
+    "OpenAILLMService",
+    "ReplicateLLMService",
+] 

isa_model/inference/services/llm/gemma_service.py

@@ -0,0 +1,143 @@
+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

@@ -0,0 +1,143 @@
+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/ollama_llm_service.py

@@ -0,0 +1,108 @@
+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.local_services import OllamaBackendClient
+
+logger = logging.getLogger(__name__)
+
+class OllamaLLMService(BaseLLMService):
+    """Ollama LLM service using backend client"""
+    
+    def __init__(self, provider: 'BaseProvider', model_name: str = "llama3.1", backend: Optional[OllamaBackendClient] = None):
+        super().__init__(provider, model_name)
+        
+        # Use provided backend or create new one
+        if backend:
+            self.backend = backend
+        else:
+            host = self.config.get("host", "localhost")
+            port = self.config.get("port", 11434)
+            self.backend = OllamaBackendClient(host, port)
+            
+        self.last_token_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
+        logger.info(f"Initialized OllamaLLMService 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:
+            payload = {
+                "model": self.model_name,
+                "messages": messages,
+                "stream": False
+            }
+            response = await self.backend.post("/api/chat", payload)
+            
+            # Update token usage if available
+            if "eval_count" in response:
+                self.last_token_usage = {
+                    "prompt_tokens": response.get("prompt_eval_count", 0),
+                    "completion_tokens": response.get("eval_count", 0),
+                    "total_tokens": response.get("prompt_eval_count", 0) + response.get("eval_count", 0)
+                }
+            
+            return response["message"]["content"]
+            
+        except Exception as e:
+            logger.error(f"Error in chat completion: {e}")
+            raise
+    
+    async def acompletion(self, prompt: str):
+        """Text completion method"""
+        try:
+            payload = {
+                "model": self.model_name,
+                "prompt": prompt,
+                "stream": False
+            }
+            response = await self.backend.post("/api/generate", payload)
+            
+            # Update token usage if available
+            if "eval_count" in response:
+                self.last_token_usage = {
+                    "prompt_tokens": response.get("prompt_eval_count", 0),
+                    "completion_tokens": response.get("eval_count", 0),
+                    "total_tokens": response.get("prompt_eval_count", 0) + response.get("eval_count", 0)
+                }
+            
+            return response["response"]
+            
+        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"""
+        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"""
+        # Note: This would require modifying the backend to support streaming
+        # For now, return the full response
+        response = await self.achat(messages)
+        yield response
+    
+    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()