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

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/training/__init__.py

@@ -0,0 +1,44 @@
+"""
+ISA Model Training Framework
+
+This module provides unified interfaces for training various types of AI models:
+- LLM training with LlamaFactory
+- Image model training with Flux/LoRA
+- Model evaluation and benchmarking
+
+Usage:
+    from isa_model.training import TrainingFactory
+    
+    # Create training factory
+    factory = TrainingFactory()
+    
+    # Fine-tune Gemma 3:4B
+    model_path = factory.finetune_llm(
+        model_name="gemma:4b",
+        dataset_path="path/to/data.json",
+        training_type="sft"
+    )
+"""
+
+from .factory import TrainingFactory, finetune_gemma
+from .engine.llama_factory import (
+    LlamaFactory,
+    LlamaFactoryConfig,
+    SFTConfig,
+    RLConfig,
+    DPOConfig,
+    TrainingStrategy,
+    DatasetFormat
+)
+
+__all__ = [
+    "TrainingFactory",
+    "finetune_gemma",
+    "LlamaFactory", 
+    "LlamaFactoryConfig",
+    "SFTConfig",
+    "RLConfig", 
+    "DPOConfig",
+    "TrainingStrategy",
+    "DatasetFormat"
+] 

isa_model/training/factory.py

@@ -0,0 +1,393 @@
+"""
+Unified Training Factory for ISA Model Framework
+
+This factory provides a single interface for all training operations:
+- LLM fine-tuning (SFT, DPO, RLHF)
+- Image model training (Flux, LoRA)
+- Model evaluation and benchmarking
+"""
+
+import os
+import logging
+from typing import Optional, Dict, Any, Union, List
+from pathlib import Path
+import datetime
+
+from .engine.llama_factory import LlamaFactory, TrainingStrategy, DatasetFormat
+from .engine.llama_factory.config import SFTConfig, RLConfig, DPOConfig
+
+logger = logging.getLogger(__name__)
+
+
+class TrainingFactory:
+    """
+    Unified factory for all AI model training operations.
+    
+    This class provides simplified interfaces for:
+    - LLM training using LlamaFactory
+    - Image model training using Flux/LoRA
+    - Model evaluation and benchmarking
+    
+    Example usage for fine-tuning Gemma 3:4B:
+        ```python
+        from isa_model.training import TrainingFactory
+        
+        factory = TrainingFactory()
+        
+        # Fine-tune with your dataset
+        model_path = factory.finetune_llm(
+            model_name="google/gemma-2-4b-it",
+            dataset_path="path/to/your/data.json",
+            training_type="sft",
+            use_lora=True,
+            num_epochs=3,
+            batch_size=4,
+            learning_rate=2e-5
+        )
+        
+        # Train with DPO for preference optimization
+        dpo_model = factory.train_with_preferences(
+            model_path=model_path,
+            preference_data="path/to/preferences.json",
+            beta=0.1
+        )
+        ```
+    """
+    
+    def __init__(self, base_output_dir: Optional[str] = None):
+        """
+        Initialize the training factory.
+        
+        Args:
+            base_output_dir: Base directory for all training outputs
+        """
+        self.base_output_dir = base_output_dir or os.path.join(os.getcwd(), "training_outputs")
+        os.makedirs(self.base_output_dir, exist_ok=True)
+        
+        # Initialize sub-factories
+        self.llm_factory = LlamaFactory(base_output_dir=os.path.join(self.base_output_dir, "llm"))
+        
+        logger.info(f"TrainingFactory initialized with output dir: {self.base_output_dir}")
+    
+    def _get_output_dir(self, model_name: str, training_type: str) -> str:
+        """Generate timestamped output directory."""
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        safe_model_name = model_name.replace("/", "_").replace(":", "_")
+        return os.path.join(self.base_output_dir, f"{safe_model_name}_{training_type}_{timestamp}")
+    
+    # =================
+    # LLM Training Methods
+    # =================
+    
+    def finetune_llm(
+        self,
+        model_name: str,
+        dataset_path: str,
+        training_type: str = "sft",
+        output_dir: Optional[str] = None,
+        dataset_format: str = "alpaca",
+        use_lora: bool = True,
+        batch_size: int = 4,
+        num_epochs: int = 3,
+        learning_rate: float = 2e-5,
+        max_length: int = 1024,
+        lora_rank: int = 8,
+        lora_alpha: int = 16,
+        val_dataset_path: Optional[str] = None,
+        **kwargs
+    ) -> str:
+        """
+        Fine-tune an LLM model.
+        
+        Args:
+            model_name: Model identifier (e.g., "google/gemma-2-4b-it", "meta-llama/Llama-2-7b-hf")
+            dataset_path: Path to training dataset
+            training_type: Type of training ("sft", "dpo", "rlhf")
+            output_dir: Custom output directory
+            dataset_format: Dataset format ("alpaca", "sharegpt", "custom")
+            use_lora: Whether to use LoRA for efficient training
+            batch_size: Training batch size
+            num_epochs: Number of training epochs
+            learning_rate: Learning rate
+            max_length: Maximum sequence length
+            lora_rank: LoRA rank parameter
+            lora_alpha: LoRA alpha parameter
+            val_dataset_path: Path to validation dataset (optional)
+            **kwargs: Additional training parameters
+            
+        Returns:
+            Path to the trained model
+            
+        Example:
+            ```python
+            # Fine-tune Gemma 3:4B with your dataset
+            model_path = factory.finetune_llm(
+                model_name="google/gemma-2-4b-it",
+                dataset_path="my_training_data.json",
+                training_type="sft",
+                use_lora=True,
+                num_epochs=3,
+                batch_size=4
+            )
+            ```
+        """
+        if not output_dir:
+            output_dir = self._get_output_dir(model_name, training_type)
+        
+        # Convert format string to enum
+        format_map = {
+            "alpaca": DatasetFormat.ALPACA,
+            "sharegpt": DatasetFormat.SHAREGPT,
+            "custom": DatasetFormat.CUSTOM
+        }
+        dataset_format_enum = format_map.get(dataset_format, DatasetFormat.ALPACA)
+        
+        if training_type.lower() == "sft":
+            return self.llm_factory.finetune(
+                model_path=model_name,
+                train_data=dataset_path,
+                val_data=val_dataset_path,
+                output_dir=output_dir,
+                dataset_format=dataset_format_enum,
+                use_lora=use_lora,
+                batch_size=batch_size,
+                num_epochs=num_epochs,
+                learning_rate=learning_rate,
+                max_length=max_length,
+                lora_rank=lora_rank,
+                lora_alpha=lora_alpha,
+                **kwargs
+            )
+        else:
+            raise ValueError(f"Training type '{training_type}' not supported yet. Use 'sft' for now.")
+    
+    def train_with_preferences(
+        self,
+        model_path: str,
+        preference_data: str,
+        output_dir: Optional[str] = None,
+        reference_model: Optional[str] = None,
+        beta: float = 0.1,
+        use_lora: bool = True,
+        batch_size: int = 4,
+        num_epochs: int = 3,
+        learning_rate: float = 5e-6,
+        val_data: Optional[str] = None,
+        **kwargs
+    ) -> str:
+        """
+        Train model with preference data using DPO.
+        
+        Args:
+            model_path: Path to the base model
+            preference_data: Path to preference dataset
+            output_dir: Custom output directory
+            reference_model: Reference model for DPO (optional)
+            beta: DPO beta parameter
+            use_lora: Whether to use LoRA
+            batch_size: Training batch size
+            num_epochs: Number of epochs
+            learning_rate: Learning rate
+            val_data: Validation data path
+            **kwargs: Additional parameters
+            
+        Returns:
+            Path to the trained model
+        """
+        if not output_dir:
+            model_name = os.path.basename(model_path)
+            output_dir = self._get_output_dir(model_name, "dpo")
+        
+        return self.llm_factory.dpo(
+            model_path=model_path,
+            train_data=preference_data,
+            val_data=val_data,
+            reference_model=reference_model,
+            output_dir=output_dir,
+            use_lora=use_lora,
+            batch_size=batch_size,
+            num_epochs=num_epochs,
+            learning_rate=learning_rate,
+            beta=beta,
+            **kwargs
+        )
+    
+    def train_reward_model(
+        self,
+        model_path: str,
+        reward_data: str,
+        output_dir: Optional[str] = None,
+        use_lora: bool = True,
+        batch_size: int = 8,
+        num_epochs: int = 3,
+        learning_rate: float = 1e-5,
+        val_data: Optional[str] = None,
+        **kwargs
+    ) -> str:
+        """
+        Train a reward model for RLHF.
+        
+        Args:
+            model_path: Base model path
+            reward_data: Reward training data
+            output_dir: Output directory
+            use_lora: Whether to use LoRA
+            batch_size: Batch size
+            num_epochs: Number of epochs
+            learning_rate: Learning rate
+            val_data: Validation data
+            **kwargs: Additional parameters
+            
+        Returns:
+            Path to trained reward model
+        """
+        if not output_dir:
+            model_name = os.path.basename(model_path)
+            output_dir = self._get_output_dir(model_name, "reward")
+        
+        return self.llm_factory.train_reward_model(
+            model_path=model_path,
+            train_data=reward_data,
+            val_data=val_data,
+            output_dir=output_dir,
+            use_lora=use_lora,
+            batch_size=batch_size,
+            num_epochs=num_epochs,
+            learning_rate=learning_rate,
+            **kwargs
+        )
+    
+    # =================
+    # Image Model Training Methods
+    # =================
+    
+    def train_image_model(
+        self,
+        model_type: str = "flux",
+        training_images_dir: str = "",
+        output_dir: Optional[str] = None,
+        use_lora: bool = True,
+        num_epochs: int = 1000,
+        batch_size: int = 1,
+        learning_rate: float = 1e-4,
+        **kwargs
+    ) -> str:
+        """
+        Train an image generation model.
+        
+        Args:
+            model_type: Type of model ("flux", "lora")
+            training_images_dir: Directory containing training images
+            output_dir: Output directory
+            use_lora: Whether to use LoRA
+            num_epochs: Training epochs
+            batch_size: Batch size
+            learning_rate: Learning rate
+            **kwargs: Additional parameters
+            
+        Returns:
+            Path to trained model
+        """
+        if not output_dir:
+            output_dir = self._get_output_dir("image_model", model_type)
+        
+        # TODO: Implement image model training
+        logger.warning("Image model training not fully implemented yet")
+        return output_dir
+    
+    # =================
+    # Utility Methods
+    # =================
+    
+    def get_training_status(self, output_dir: str) -> Dict[str, Any]:
+        """
+        Get training status from output directory.
+        
+        Args:
+            output_dir: Training output directory
+            
+        Returns:
+            Dictionary with training status information
+        """
+        status = {
+            "output_dir": output_dir,
+            "exists": os.path.exists(output_dir),
+            "files": []
+        }
+        
+        if status["exists"]:
+            status["files"] = os.listdir(output_dir)
+            
+        return status
+    
+    def list_trained_models(self) -> List[Dict[str, Any]]:
+        """
+        List all trained models in the output directory.
+        
+        Returns:
+            List of model information dictionaries
+        """
+        models = []
+        
+        if os.path.exists(self.base_output_dir):
+            for item in os.listdir(self.base_output_dir):
+                item_path = os.path.join(self.base_output_dir, item)
+                if os.path.isdir(item_path):
+                    models.append({
+                        "name": item,
+                        "path": item_path,
+                        "created": datetime.datetime.fromtimestamp(
+                            os.path.getctime(item_path)
+                        ).isoformat()
+                    })
+        
+        return sorted(models, key=lambda x: x["created"], reverse=True)
+
+
+# Convenience functions for quick access
+def finetune_gemma(
+    dataset_path: str,
+    model_size: str = "4b",
+    output_dir: Optional[str] = None,
+    **kwargs
+) -> str:
+    """
+    Quick function to fine-tune Gemma models.
+    
+    Args:
+        dataset_path: Path to training dataset
+        model_size: Model size ("2b", "4b", "7b")
+        output_dir: Output directory
+        **kwargs: Additional training parameters
+        
+    Returns:
+        Path to fine-tuned model
+        
+    Example:
+        ```python
+        from isa_model.training import finetune_gemma
+        
+        model_path = finetune_gemma(
+            dataset_path="my_data.json",
+            model_size="4b",
+            num_epochs=3,
+            batch_size=4
+        )
+        ```
+    """
+    factory = TrainingFactory()
+    
+    model_map = {
+        "2b": "google/gemma-2-2b-it",
+        "4b": "google/gemma-2-4b-it", 
+        "7b": "google/gemma-2-7b-it"
+    }
+    
+    model_name = model_map.get(model_size, "google/gemma-2-4b-it")
+    
+    return factory.finetune_llm(
+        model_name=model_name,
+        dataset_path=dataset_path,
+        output_dir=output_dir,
+        **kwargs
+    )