isa-model 0.1.0__py3-none-any.whl

isa_model/inference/services/audio/whisper_service.py

@@ -0,0 +1,186 @@
+import logging
+import asyncio
+import io
+import numpy as np
+from typing import Dict, Any, Optional, Union, BinaryIO
+
+from isa_model.inference.services.base_service import BaseService
+from isa_model.inference.backends.triton_client import TritonClient
+
+logger = logging.getLogger(__name__)
+
+
+class WhisperService(BaseService):
+    """
+    Service for Whisper speech-to-text using Triton Inference Server.
+    """
+    
+    def __init__(self, triton_url: str = "localhost:8001", model_name: str = "whisper"):
+        """
+        Initialize the Whisper 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 configuration
+        self.default_config = {
+            "language": "en",
+            "sampling_rate": 16000
+        }
+        
+        self.logger = logger
+    
+    async def load(self) -> None:
+        """
+        Load the client connection to Triton.
+        """
+        if self.is_loaded():
+            return
+        
+        try:
+            from tritonclient.http import InferenceServerClient
+            
+            # 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 transcribe(self, 
+                       audio: Union[str, BinaryIO, bytes, np.ndarray], 
+                       language: str = "en",
+                       config: Optional[Dict[str, Any]] = None) -> str:
+        """
+        Transcribe audio to text using Triton.
+        
+        Args:
+            audio: Audio input (file path, file-like object, bytes, or numpy array)
+            language: Language code (e.g., "en", "fr")
+            config: Additional configuration parameters
+            
+        Returns:
+            Transcribed text
+        """
+        if not self.is_loaded():
+            await self.load()
+        
+        # Process audio to get numpy array
+        audio_array = await self._process_audio_input(audio)
+        
+        # Get configuration
+        merged_config = self.default_config.copy()
+        if config:
+            merged_config.update(config)
+        
+        # Override language if provided
+        if language:
+            merged_config["language"] = language
+        
+        try:
+            # Prepare inputs
+            inputs = {
+                "audio_input": audio_array,
+                "language": np.array([merged_config["language"]], dtype=np.object_)
+            }
+            
+            # Run inference
+            result = await self.client.infer(
+                model_name=self.model_name,
+                inputs=inputs,
+                outputs=["text_output"]
+            )
+            
+            # Extract transcription
+            transcription = result["text_output"][0].decode('utf-8')
+            
+            return transcription
+            
+        except Exception as e:
+            self.logger.error(f"Error during Whisper transcription: {str(e)}")
+            raise
+    
+    async def _process_audio_input(self, audio: Union[str, BinaryIO, bytes, np.ndarray]) -> np.ndarray:
+        """
+        Process different types of audio inputs into a numpy array.
+        
+        Args:
+            audio: Audio input (file path, file-like object, bytes, or numpy array)
+            
+        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.default_config["sampling_rate"])
+                return y.astype(np.float32)
+            
+            elif isinstance(audio, (io.IOBase, BinaryIO)):
+                # File-like object
+                audio.seek(0)
+                y, sr = librosa.load(audio, sr=self.default_config["sampling_rate"])
+                return y.astype(np.float32)
+            
+            elif isinstance(audio, bytes):
+                # Bytes
+                with io.BytesIO(audio) as audio_bytes:
+                    y, sr = librosa.load(audio_bytes, sr=self.default_config["sampling_rate"])
+                return y.astype(np.float32)
+            
+            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
+    
+    def get_model_info(self) -> Dict[str, Any]:
+        """
+        Get information about the model.
+        
+        Returns:
+            Dictionary containing model information
+        """
+        return {
+            "name": self.model_name,
+            "type": "speech",
+            "backend": "triton",
+            "url": self.triton_url,
+            "loaded": self.is_loaded(),
+            "config": self.default_config
+        } 

isa_model/inference/services/audio/yyds_audio_service.py

@@ -0,0 +1,71 @@
+from typing import Dict, Any
+import tempfile
+import os
+from openai import AsyncOpenAI
+from tenacity import retry, stop_after_attempt, wait_exponential
+from ...base_service import BaseService
+from ...base_provider import BaseProvider
+from app.config.config_manager import config_manager
+
+logger = config_manager.get_logger(__name__)
+
+class YYDSAudioService(BaseService):
+    """Audio model service wrapper for YYDS"""
+    
+    def __init__(self, provider: 'BaseProvider', model_name: str):
+        super().__init__(provider, model_name)
+        # 初始化 AsyncOpenAI 客户端
+        self._client = AsyncOpenAI(
+            api_key=self.config.get('api_key'),
+            base_url=self.config.get('base_url')
+        )
+        self.language = self.config.get('language', None)
+    
+    @property
+    def client(self) -> AsyncOpenAI:
+        """获取底层的 OpenAI 客户端"""
+        return self._client
+    
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        reraise=True
+    )
+    async def transcribe(self, audio_data: bytes) -> Dict[str, Any]:
+        """转写音频数据
+        
+        Args:
+            audio_data: 音频二进制数据
+            
+        Returns:
+            Dict[str, Any]: 包含转写文本的字典
+        """
+        try:
+            # 创建临时文件存储音频数据
+            with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_file:
+                temp_file.write(audio_data)
+                temp_file.flush()
+                
+                # 以二进制模式打开文件用于 API 请求
+                with open(temp_file.name, 'rb') as audio_file:
+                    # 只在有效的 ISO-639-1 语言代码时包含 language 参数
+                    params = {
+                        'model': self.model_name,
+                        'file': audio_file,
+                    }
+                    if self.language and isinstance(self.language, str):
+                        params['language'] = self.language
+                        
+                    response = await self._client.audio.transcriptions.create(**params)
+                    
+                # 清理临时文件
+                os.unlink(temp_file.name)
+                
+                # 返回包含转写文本的字典
+                return {
+                    "text": response.text
+                }
+                
+        except Exception as e:
+            logger.error(f"Error in audio transcription: {e}")
+            raise

isa_model/inference/services/base_service.py

@@ -0,0 +1,106 @@
+from abc import ABC, abstractmethod
+from typing import Dict, Any, List, Union, AsyncGenerator, TypeVar, Optional
+from isa_model.inference.providers.base_provider import BaseProvider
+
+T = TypeVar('T')  # Generic type for responses
+
+class BaseService(ABC):
+    """Base class for all AI services"""
+    
+    def __init__(self, provider: 'BaseProvider', model_name: str):
+        self.provider = provider
+        self.model_name = model_name
+        self.config = provider.get_config()
+        
+    def __await__(self):
+        """Make the service awaitable"""
+        yield
+        return self
+
+class BaseLLMService(BaseService):
+    """Base class for LLM services"""
+    
+    @abstractmethod
+    async def ainvoke(self, prompt: Union[str, List[Dict[str, str]], Any]) -> T:
+        """Universal invocation method"""
+        pass
+    
+    @abstractmethod
+    async def achat(self, messages: List[Dict[str, str]]) -> T:
+        """Chat completion method"""
+        pass
+    
+    @abstractmethod
+    async def acompletion(self, prompt: str) -> T:
+        """Text completion method"""
+        pass
+    
+    @abstractmethod
+    async def agenerate(self, messages: List[Dict[str, str]], n: int = 1) -> List[T]:
+        """Generate multiple completions"""
+        pass
+    
+    @abstractmethod
+    async def astream_chat(self, messages: List[Dict[str, str]]) -> AsyncGenerator[str, None]:
+        """Stream chat responses"""
+        pass
+    
+    @abstractmethod
+    def get_token_usage(self) -> Any:
+        """Get total token usage statistics"""
+        pass
+    
+    @abstractmethod
+    def get_last_token_usage(self) -> Dict[str, int]:
+        """Get token usage from last request"""
+        pass
+
+class BaseEmbeddingService(BaseService):
+    """Base class for embedding services"""
+    
+    @abstractmethod
+    async def create_text_embedding(self, text: str) -> List[float]:
+        """Create embedding for single text"""
+        pass
+    
+    @abstractmethod
+    async def create_text_embeddings(self, texts: List[str]) -> List[List[float]]:
+        """Create embeddings for multiple texts"""
+        pass
+    
+    @abstractmethod
+    async def create_chunks(self, text: str, metadata: Optional[Dict] = None) -> List[Dict]:
+        """Create text chunks with embeddings"""
+        pass
+    
+    @abstractmethod
+    async def compute_similarity(self, embedding1: List[float], embedding2: List[float]) -> float:
+        """Compute similarity between two embeddings"""
+        pass
+    
+    @abstractmethod
+    async def close(self):
+        """Cleanup resources"""
+        pass
+
+class BaseRerankService(BaseService):
+    """Base class for reranking services"""
+    
+    @abstractmethod
+    async def rerank(
+        self,
+        query: str,
+        documents: List[Dict],
+        top_k: int = 5
+    ) -> List[Dict]:
+        """Rerank documents based on query relevance"""
+        pass
+    
+    @abstractmethod
+    async def rerank_texts(
+        self,
+        query: str,
+        texts: List[str]
+    ) -> List[Dict]:
+        """Rerank raw texts based on query relevance"""
+        pass

isa_model/inference/services/base_tts_service.py

@@ -0,0 +1,66 @@
+from abc import abstractmethod
+from typing import Dict, Any, Optional, Union, BinaryIO
+from .base_service import BaseService
+
+class BaseTTSService(BaseService):
+    """Base class for Text-to-Speech services"""
+    
+    @abstractmethod
+    async def generate_speech(
+        self, 
+        text: str, 
+        voice_id: Optional[str] = None,
+        language: Optional[str] = None,
+        speed: float = 1.0,
+        options: Optional[Dict[str, Any]] = None
+    ) -> bytes:
+        """
+        Generate speech from text
+        
+        Args:
+            text: The text to convert to speech
+            voice_id: Optional voice identifier
+            language: Optional language code
+            speed: Speech speed factor (1.0 is normal speed)
+            options: Additional model-specific options
+            
+        Returns:
+            Audio data as bytes
+        """
+        pass
+    
+    @abstractmethod
+    async def save_to_file(
+        self,
+        text: str,
+        output_file: Union[str, BinaryIO],
+        voice_id: Optional[str] = None,
+        language: Optional[str] = None,
+        speed: float = 1.0,
+        options: Optional[Dict[str, Any]] = None
+    ) -> str:
+        """
+        Generate speech and save to file
+        
+        Args:
+            text: The text to convert to speech
+            output_file: Path to output file or file-like object
+            voice_id: Optional voice identifier
+            language: Optional language code
+            speed: Speech speed factor (1.0 is normal speed)
+            options: Additional model-specific options
+            
+        Returns:
+            Path to the saved file
+        """
+        pass
+    
+    @abstractmethod
+    async def get_available_voices(self) -> Dict[str, Any]:
+        """
+        Get available voices for the TTS service
+        
+        Returns:
+            Dictionary of available voices with their details
+        """
+        pass 

isa_model/inference/services/embedding/bge_service.py

@@ -0,0 +1,183 @@
+import logging
+import asyncio
+import numpy as np
+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 BgeEmbeddingService(BaseService):
+    """
+    Service for BGE embedding using Triton Inference Server.
+    """
+    
+    def __init__(self, triton_url: str = "localhost:8001", model_name: str = "bge_embed"):
+        """
+        Initialize the BGE embedding 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 configuration
+        self.default_config = {
+            "normalize": 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 embed(self, 
+                   texts: Union[str, List[str]],
+                   normalize: Optional[bool] = None) -> np.ndarray:
+        """
+        Generate embeddings for texts using Triton.
+        
+        Args:
+            texts: Single text or list of texts to embed
+            normalize: Whether to normalize embeddings (if None, use default)
+            
+        Returns:
+            Numpy array of embeddings, shape [batch_size, embedding_dim]
+        """
+        if not self.is_loaded():
+            await self.load()
+        
+        # Handle single text input
+        if isinstance(texts, str):
+            texts = [texts]
+        
+        # Use default normalize setting if not specified
+        if normalize is None:
+            normalize = self.default_config["normalize"]
+        
+        try:
+            # Prepare inputs
+            inputs = {
+                "text_input": texts,
+                "normalize": np.array([normalize], dtype=bool)
+            }
+            
+            # Run inference
+            result = await self.client.infer(
+                model_name=self.model_name,
+                inputs=inputs,
+                outputs=["embedding"]
+            )
+            
+            # Extract embeddings
+            embeddings = result["embedding"]
+            
+            return embeddings
+            
+        except Exception as e:
+            self.logger.error(f"Error during embedding generation: {str(e)}")
+            raise
+    
+    async def similarity(self, 
+                        text1: str, 
+                        text2: str,
+                        normalize: Optional[bool] = None) -> float:
+        """
+        Calculate the similarity between two texts.
+        
+        Args:
+            text1: First text
+            text2: Second text
+            normalize: Whether to normalize embeddings (if None, use default)
+            
+        Returns:
+            Cosine similarity score (float between -1 and 1)
+        """
+        # Generate embeddings for both texts
+        embeddings = await self.embed([text1, text2], normalize=normalize)
+        
+        # Calculate cosine similarity
+        from sklearn.metrics.pairwise import cosine_similarity
+        similarity = cosine_similarity(embeddings[0:1], embeddings[1:2])[0][0]
+        
+        return float(similarity)
+    
+    async def batch_similarity(self, 
+                              queries: List[str], 
+                              documents: List[str],
+                              normalize: Optional[bool] = None) -> np.ndarray:
+        """
+        Calculate similarities between queries and documents.
+        
+        Args:
+            queries: List of query texts
+            documents: List of document texts
+            normalize: Whether to normalize embeddings (if None, use default)
+            
+        Returns:
+            Numpy array of similarity scores, shape [len(queries), len(documents)]
+        """
+        # Generate embeddings for queries and documents
+        query_embeddings = await self.embed(queries, normalize=normalize)
+        doc_embeddings = await self.embed(documents, normalize=normalize)
+        
+        # Calculate cosine similarities
+        from sklearn.metrics.pairwise import cosine_similarity
+        similarities = cosine_similarity(query_embeddings, doc_embeddings)
+        
+        return similarities
+    
+    def get_model_info(self) -> Dict[str, Any]:
+        """
+        Get information about the model.
+        
+        Returns:
+            Dictionary containing model information
+        """
+        return {
+            "name": self.model_name,
+            "type": "embedding",
+            "backend": "triton",
+            "url": self.triton_url,
+            "loaded": self.is_loaded(),
+            "embedding_dim": 1024,  # Typical for BGE models
+            "config": self.default_config
+        } 

isa_model/inference/services/embedding/ollama_embed_service.py

@@ -0,0 +1,85 @@
+import logging
+from typing import List, Dict, Any, Optional
+from isa_model.inference.services.base_service import BaseEmbeddingService
+from isa_model.inference.providers.base_provider import BaseProvider
+from isa_model.inference.backends.local_services import OllamaBackendClient
+
+logger = logging.getLogger(__name__)
+
+class OllamaEmbedService(BaseEmbeddingService):
+    """Ollama embedding service using backend client"""
+    
+    def __init__(self, provider: 'BaseProvider', model_name: str = "bge-m3", 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)
+            
+        logger.info(f"Initialized OllamaEmbedService with model {model_name}")
+    
+    async def create_text_embedding(self, text: str) -> List[float]:
+        """Create embedding for text"""
+        try:
+            payload = {
+                "model": self.model_name,
+                "prompt": text
+            }
+            response = await self.backend.post("/api/embeddings", payload)
+            return response["embedding"]
+            
+        except Exception as e:
+            logger.error(f"Error creating text embedding: {e}")
+            raise
+    
+    async def create_text_embeddings(self, texts: List[str]) -> List[List[float]]:
+        """Create embeddings for multiple texts"""
+        embeddings = []
+        for text in texts:
+            embedding = await self.create_text_embedding(text)
+            embeddings.append(embedding)
+        return embeddings
+    
+    async def create_chunks(self, text: str, metadata: Optional[Dict] = None) -> List[Dict]:
+        """Create text chunks with embeddings"""
+        # 简单实现：将文本分成固定大小的块
+        chunk_size = 200  # 单词数量
+        chunks = []
+        
+        # 按单词分割
+        words = text.split()
+        
+        # 分块
+        for i in range(0, len(words), chunk_size):
+            chunk_text = " ".join(words[i:i+chunk_size])
+            embedding = await self.create_text_embedding(chunk_text)
+            
+            chunk = {
+                "text": chunk_text,
+                "embedding": embedding,
+                "metadata": metadata or {}
+            }
+            chunks.append(chunk)
+            
+        return chunks
+    
+    async def compute_similarity(self, embedding1: List[float], embedding2: List[float]) -> float:
+        """计算两个嵌入向量之间的余弦相似度"""
+        # 余弦相似度简单实现
+        dot_product = sum(a * b for a, b in zip(embedding1, embedding2))
+        norm1 = sum(a * a for a in embedding1) ** 0.5
+        norm2 = sum(b * b for b in embedding2) ** 0.5
+        
+        if norm1 * norm2 == 0:
+            return 0.0
+            
+        return dot_product / (norm1 * norm2)
+    
+    async def close(self):
+        """Close the backend client"""
+        await self.backend.close()
+