isa-model 0.0.1__py3-none-any.whl

isa_model/inference/services/others/table_transformer_service.py

@@ -0,0 +1,61 @@
+from typing import Dict, Any, List, Union, Optional
+from ...base_service import BaseService
+from ...base_provider import BaseProvider
+from transformers import TableTransformerForObjectDetection, DetrImageProcessor
+import torch
+from PIL import Image
+import numpy as np
+
+class TableTransformerService(BaseService):
+    """Table detection service using Microsoft's Table Transformer"""
+    
+    def __init__(self, provider: 'BaseProvider', model_name: str = "microsoft/table-transformer-detection"):
+        super().__init__(provider, model_name)
+        self.processor = DetrImageProcessor.from_pretrained(model_name)
+        self.model = TableTransformerForObjectDetection.from_pretrained(model_name)
+        if torch.cuda.is_available():
+            self.model = self.model.cuda()
+    
+    async def detect_tables(self, image_path: str) -> Dict[str, Any]:
+        """Detect tables in image"""
+        try:
+            # Load and process image
+            image = Image.open(image_path)
+            inputs = self.processor(images=image, return_tensors="pt")
+            
+            if torch.cuda.is_available():
+                inputs = {k: v.cuda() for k, v in inputs.items()}
+            
+            # Run inference
+            outputs = self.model(**inputs)
+            
+            # Convert outputs to image size
+            target_sizes = torch.tensor([image.size[::-1]])
+            results = self.processor.post_process_object_detection(
+                outputs, target_sizes=target_sizes, threshold=0.7
+            )[0]
+            
+            tables = []
+            for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+                if label == 1:  # Table class
+                    tables.append({
+                        "confidence": score.item(),
+                        "bbox": box.tolist(),
+                        "type": "table"
+                    })
+            
+            return {
+                "tables": tables,
+                "image_size": image.size
+            }
+            
+        except Exception as e:
+            raise RuntimeError(f"Table detection failed: {e}")
+    
+    async def close(self):
+        """Cleanup resources"""
+        if hasattr(self, 'model'):
+            del self.model
+        if hasattr(self, 'processor'):
+            del self.processor
+        torch.cuda.empty_cache() 

isa_model/inference/services/vision/__init__.py

@@ -0,0 +1,12 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+"""
+Vision服务包
+包含所有视觉相关服务模块
+"""
+
+# 导出ReplicateVisionService
+from isa_model.inference.services.vision.replicate_image_gen_service import ReplicateVisionService
+
+__all__ = ["ReplicateVisionService"] 

isa_model/inference/services/vision/helpers/image_utils.py

@@ -0,0 +1,58 @@
+from io import BytesIO
+from PIL import Image
+from typing import Union
+import base64
+from app.config.config_manager import config_manager
+
+logger = config_manager.get_logger(__name__)
+
+def compress_image(image_data: Union[bytes, BytesIO], max_size: int = 1024) -> bytes:
+    """压缩图片以减小大小
+    
+    Args:
+        image_data: 图片数据，可以是 bytes 或 BytesIO
+        max_size: 最大尺寸（像素）
+        
+    Returns:
+        bytes: 压缩后的图片数据
+    """
+    try:
+        # 如果输入是 bytes，转换为 BytesIO
+        if isinstance(image_data, bytes):
+            image_data = BytesIO(image_data)
+            
+        img = Image.open(image_data)
+        
+        # 转换为 RGB 模式（如果需要）
+        if img.mode in ('RGBA', 'P'):
+            img = img.convert('RGB')
+        
+        # 计算新尺寸，保持宽高比
+        ratio = max_size / max(img.size)
+        if ratio < 1:
+            new_size = tuple(int(dim * ratio) for dim in img.size)
+            img = img.resize(new_size, Image.Resampling.LANCZOS)
+        
+        # 保存压缩后的图片
+        output = BytesIO()
+        img.save(output, format='JPEG', quality=85, optimize=True)
+        return output.getvalue()
+        
+    except Exception as e:
+        logger.error(f"Error compressing image: {e}")
+        raise
+
+def encode_image_to_base64(image_data: bytes) -> str:
+    """将图片数据编码为 base64 字符串
+    
+    Args:
+        image_data: 图片二进制数据
+        
+    Returns:
+        str: base64 编码的字符串
+    """
+    try:
+        return base64.b64encode(image_data).decode('utf-8')
+    except Exception as e:
+        logger.error(f"Error encoding image to base64: {e}")
+        raise 

isa_model/inference/services/vision/helpers/text_splitter.py

@@ -0,0 +1,46 @@
+from typing import Dict, List, Optional
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from uuid import uuid4
+
+class TextChunkHelper:
+    """Text splitting and chunking helper"""
+    
+    def __init__(self, 
+                 chunk_size: int = 512,
+                 chunk_overlap: int = 50,
+                 min_chunk_size: int = 50):
+        self.text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=chunk_size,
+            chunk_overlap=chunk_overlap,
+            length_function=len,
+            separators=["\n\n", "\n", ". ", ", ", " "]
+        )
+        self.min_chunk_size = min_chunk_size
+
+    def create_chunks(self, text: str, metadata: Optional[Dict] = None) -> List[Dict]:
+        """Create text chunks with metadata"""
+        if not text or not isinstance(text, str):
+            raise ValueError("Text must be a non-empty string")
+            
+        chunks = self.text_splitter.split_text(text)
+        valid_chunks = [
+            chunk for chunk in chunks 
+            if len(chunk) >= self.min_chunk_size
+        ]
+        
+        results = []
+        for i, chunk in enumerate(valid_chunks):
+            chunk_data = {
+                "chunk_id": f"chunk_{uuid4().hex[:8]}",
+                "content": chunk,
+                "token_count": len(chunk),
+                "metadata": {
+                    **(metadata or {}),
+                    "position": i,
+                    "start_idx": text.find(chunk),
+                    "end_idx": text.find(chunk) + len(chunk)
+                }
+            }
+            results.append(chunk_data)
+            
+        return results 

isa_model/inference/services/vision/ollama_vision_service.py

@@ -0,0 +1,60 @@
+import os
+import json
+import base64
+import ollama
+from typing import Dict, Any, Union
+from tenacity import retry, stop_after_attempt, wait_exponential
+from isa_model.inference.services.base_service import BaseService
+from isa_model.inference.providers.base_provider import BaseProvider
+import logging
+
+logger = logging.getLogger(__name__)
+
+class OllamaVisionService(BaseService):
+    """Vision model service wrapper for Ollama using base64 encoded images"""
+    
+    def __init__(self, provider: 'BaseProvider', model_name: str = 'gemma3:4b'):
+        super().__init__(provider, model_name)
+        self.max_tokens = self.config.get('max_tokens', 1000)
+        self.temperature = self.config.get('temperature', 0.7)
+    
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        reraise=True
+    )
+    async def analyze_image(self, image_data: Union[bytes, str], query: str) -> str:
+        """分析图片并返回结果
+        
+        Args:
+            image_data: 图片数据，可以是 bytes 或图片路径字符串
+            query: 查询文本
+            
+        Returns:
+            str: 分析结果
+        """
+        try:
+            # 如果是文件路径，读取文件内容
+            if isinstance(image_data, str):
+                with open(image_data, 'rb') as f:
+                    image_data = f.read()
+            
+            # 转换为base64
+            image_base64 = base64.b64encode(image_data).decode('utf-8')
+            
+            # 使用 ollama 库直接调用
+            response = ollama.chat(
+                model=self.model_name,
+                messages=[{
+                    'role': 'user',
+                    'content': query,
+                    'images': [image_base64]
+                }]
+            )
+            
+            return response['message']['content']
+            
+        except Exception as e:
+            logger.error(f"Error in image analysis: {e}")
+            raise
+

isa_model/inference/services/vision/openai_vision_service.py

@@ -0,0 +1,80 @@
+from typing import Dict, Any, Union
+from openai import AsyncOpenAI
+from tenacity import retry, stop_after_attempt, wait_exponential
+from isa_model.inference.services.base_service import BaseService
+from isa_model.inference.providers.base_provider import BaseProvider
+from .helpers.image_utils import compress_image, encode_image_to_base64
+import logging
+
+logger = logging.getLogger(__name__)
+
+class OpenAIVisionService(BaseService):
+    """Vision 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.max_tokens = self.config.get('max_tokens', 1000)
+        self.temperature = self.config.get('temperature', 0.7)
+    
+    @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 analyze_image(self, image_data: Union[bytes, str], query: str) -> str:
+        """分析图片并返回结果
+        
+        Args:
+            image_data: 图片数据，可以是 bytes 或已编码的 base64 字符串
+            query: 查询文本
+            
+        Returns:
+            str: 分析结果
+        """
+        try:
+            # 处理图片数据
+            if isinstance(image_data, bytes):
+                # 压缩并编码图片
+                compressed_image = compress_image(image_data)
+                image_b64 = encode_image_to_base64(compressed_image)
+            else:
+                image_b64 = image_data
+                
+            # 移除可能存在的 base64 前缀
+            if 'base64,' in image_b64:
+                image_b64 = image_b64.split('base64,')[1]
+                
+            # 使用 AsyncOpenAI 客户端创建请求
+            response = await self._client.chat.completions.create(
+                model=self.model_name,
+                messages=[{
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": query},
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": f"data:image/jpeg;base64,{image_b64}"
+                            }
+                        }
+                    ]
+                }],
+                max_tokens=self.max_tokens,
+                temperature=self.temperature
+            )
+            
+            return response.choices[0].message.content
+            
+        except Exception as e:
+            logger.error(f"Error in image analysis: {e}")
+            raise

isa_model/inference/services/vision/replicate_image_gen_service.py

@@ -0,0 +1,185 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+"""
+Replicate Vision服务
+用于与Replicate API交互，支持图像生成和图像分析
+"""
+
+import os
+import time
+import uuid
+import logging
+from typing import Dict, Any, List, Optional, Union, Tuple
+import asyncio
+import aiohttp
+import replicate # 导入 replicate 库
+from PIL import Image
+from io import BytesIO
+
+# 调整 BaseService 的导入路径以匹配您的项目结构
+from isa_model.inference.services.base_service import BaseService
+from isa_model.inference.providers.base_provider import BaseProvider
+from isa_model.inference.base import ModelType
+
+# 设置日志记录
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class ReplicateVisionService(BaseService):
+    """
+    Replicate Vision服务，用于处理图像生成和分析。
+    经过调整，使用原生异步调用并优化了文件处理。
+    """
+    
+    def __init__(self, provider: BaseProvider, model_name: str):
+        """
+        初始化Replicate Vision服务
+        """
+        super().__init__(provider, model_name)
+        # 从 provider 或环境变量获取 API token
+        self.api_token = self.provider.config.get("api_token", os.environ.get("REPLICATE_API_TOKEN"))
+        self.model_type = ModelType.VISION
+        
+        # 可选的默认配置
+        self.guidance_scale = self.provider.config.get("guidance_scale", 7.5)
+        self.num_inference_steps = self.provider.config.get("num_inference_steps", 30)
+        
+        # 生成的图像存储目录
+        self.output_dir = "generated_images"
+        os.makedirs(self.output_dir, exist_ok=True)
+        
+        # ★ 调整点: 为 replicate 库设置 API token
+        if self.api_token:
+            # replicate 库会自动从环境变量读取，我们确保它被设置
+            os.environ["REPLICATE_API_TOKEN"] = self.api_token
+        else:
+            logger.warning("Replicate API token 未找到。服务可能无法正常工作。")
+
+    async def _prepare_input_files(self, input_data: Dict[str, Any]) -> Tuple[Dict[str, Any], List[Any]]:
+        """
+        ★ 新增辅助函数: 准备输入数据，将本地文件路径转换为文件对象。
+        这使得服务能统一处理本地文件和URL。
+        """
+        prepared_input = input_data.copy()
+        files_to_close = []
+        for key, value in prepared_input.items():
+            # 如果值是字符串，且看起来像一个存在的本地文件路径
+            if isinstance(value, str) and not value.startswith(('http://', 'https://')) and os.path.exists(value):
+                logger.info(f"检测到本地文件路径 '{value}'，准备打开文件。")
+                try:
+                    file_handle = open(value, "rb")
+                    prepared_input[key] = file_handle
+                    files_to_close.append(file_handle)
+                except Exception as e:
+                    logger.error(f"打开文件失败 '{value}': {e}")
+                    raise
+        return prepared_input, files_to_close
+
+    async def generate_image(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        使用Replicate模型生成图像 (已优化为原生异步)
+        """
+        prepared_input, files_to_close = await self._prepare_input_files(input_data)
+        try:
+            # 设置默认参数
+            if "guidance_scale" not in prepared_input:
+                prepared_input["guidance_scale"] = self.guidance_scale
+            if "num_inference_steps" not in prepared_input:
+                prepared_input["num_inference_steps"] = self.num_inference_steps
+            
+            logger.info(f"开始使用模型 {self.model_name} 生成图像 (原生异步)")
+            
+            # ★ 调整点: 使用原生异步的 replicate.async_run
+            output = await replicate.async_run(self.model_name, input=prepared_input)
+            
+            # 将结果转换为标准格式 (此部分逻辑无需改变)
+            if isinstance(output, list):
+                urls = output
+            else:
+                urls = [output]
+
+            result = {
+                "urls": urls,
+                "metadata": {
+                    "model": self.model_name,
+                    "input": input_data # 返回原始输入以供参考
+                }
+            }
+            logger.info(f"图像生成完成: {result['urls']}")
+            return result
+        except Exception as e:
+            logger.error(f"图像生成失败: {e}")
+            raise
+        finally:
+            # ★ 新增: 确保所有打开的文件都被关闭
+            for f in files_to_close:
+                f.close()
+    
+    async def analyze_image(self, image_path: str, prompt: str) -> Dict[str, Any]:
+        """
+        分析图像 (已优化为原生异步)
+        """
+        input_data = {"image": image_path, "prompt": prompt}
+        prepared_input, files_to_close = await self._prepare_input_files(input_data)
+        try:
+            logger.info(f"开始使用模型 {self.model_name} 分析图像 (原生异步)")
+            # ★ 调整点: 使用原生异步的 replicate.async_run
+            output = await replicate.async_run(self.model_name, input=prepared_input)
+            
+            result = {
+                "text": "".join(output) if isinstance(output, list) else output,
+                "metadata": {
+                    "model": self.model_name,
+                    "input": input_data
+                }
+            }
+            logger.info(f"图像分析完成")
+            return result
+        except Exception as e:
+            logger.error(f"图像分析失败: {e}")
+            raise
+        finally:
+            # ★ 新增: 确保所有打开的文件都被关闭
+            for f in files_to_close:
+                f.close()
+    
+    async def generate_and_save(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
+        """生成图像并保存到本地 (此方法无需修改)"""
+        result = await self.generate_image(input_data)
+        saved_paths = []
+        for i, url in enumerate(result["urls"]):
+            timestamp = int(time.time())
+            file_name = f"{self.output_dir}/{timestamp}_{uuid.uuid4().hex[:8]}_{i+1}.png"
+            try:
+                # Convert FileOutput object to string if necessary
+                url_str = str(url) if hasattr(url, "__str__") else url
+                await self._download_image(url_str, file_name)
+                saved_paths.append(file_name)
+                logger.info(f"图像已保存至: {file_name}")
+            except Exception as e:
+                logger.error(f"保存图像失败: {e}")
+        result["saved_paths"] = saved_paths
+        return result
+    
+    async def _download_image(self, url: str, save_path: str) -> None:
+        """异步下载图像并保存 (此方法无需修改)"""
+        try:
+            async with aiohttp.ClientSession() as session:
+                async with session.get(url) as response:
+                    response.raise_for_status()
+                    content = await response.read()
+                    with Image.open(BytesIO(content)) as img:
+                        img.save(save_path)
+        except Exception as e:
+            logger.error(f"下载图像时出错: {url}, {e}")
+            raise
+
+    # `load` 和 `unload` 方法在Replicate API场景下通常是轻量级的
+    async def load(self) -> None:
+        if not self.api_token:
+            raise ValueError("缺少Replicate API令牌，请设置REPLICATE_API_TOKEN环境变量或在provider配置中提供")
+        logger.info(f"Replicate Vision服务已准备就绪，使用模型: {self.model_name}")
+
+    async def unload(self) -> None:
+        logger.info(f"卸载Replicate Vision服务: {self.model_name}")

isa_model/inference/utils/conversion/bge_rerank_convert.py

@@ -0,0 +1,73 @@
+import os
+import torch
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from pathlib import Path
+
+def convert_bge_to_onnx(save_dir: str):
+    """Convert BGE reranker to ONNX format"""
+    try:
+        # Create save directory if it doesn't exist
+        save_dir = Path(save_dir).resolve()  # Get absolute path
+        save_dir.mkdir(parents=True, exist_ok=True)
+        
+        model_name = "BAAI/bge-reranker-v2-m3"
+        save_path = str(save_dir / "model.onnx")  # Convert to string for absolute path
+        
+        print(f"Loading model {model_name}...")
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForSequenceClassification.from_pretrained(model_name)
+        model.eval()
+        
+        # Save tokenizer for later use
+        print("Saving tokenizer...")
+        tokenizer.save_pretrained(save_dir)
+        
+        # Create dummy input
+        print("Creating dummy input...")
+        dummy_input = tokenizer(
+            [["what is panda?", "The giant panda is a bear species."]], 
+            padding=True, 
+            truncation=True, 
+            return_tensors='pt',
+            max_length=512
+        )
+        
+        # Export to ONNX with external data storage
+        print(f"Exporting to ONNX: {save_path}")
+        torch.onnx.export(
+            model,
+            (dummy_input['input_ids'], dummy_input['attention_mask']),
+            save_path,  # Using string absolute path
+            input_names=['input_ids', 'attention_mask'],
+            output_names=['logits'],
+            dynamic_axes={
+                'input_ids': {0: 'batch', 1: 'sequence'},
+                'attention_mask': {0: 'batch', 1: 'sequence'},
+                'logits': {0: 'batch'}
+            },
+            opset_version=16,
+            export_params=True,  # Export the trained parameter weights
+            do_constant_folding=True,  # Optimize constant-folding
+            verbose=True,
+            use_external_data_format=True  # Enable external data storage
+        )
+        print("Conversion completed successfully!")
+        return True
+        
+    except Exception as e:
+        print(f"Error during conversion: {e}")
+        return False
+
+if __name__ == "__main__":
+    # Get the absolute path to the model directory
+    current_dir = Path(__file__).parent.parent
+    model_dir = current_dir / "model_converted" / "bge-reranker-v2-m3"
+    
+    success = convert_bge_to_onnx(str(model_dir))
+    if success:
+        print(f"Model saved to: {model_dir}")
+        print("Files created:")
+        for file in model_dir.glob('*'):
+            print(f"- {file.name}")
+    else:
+        print("Conversion failed!")