app.py

import gradio as gr
import torch
from transformers import AutoTokenizer, pipeline
from huggingface_hub import InferenceClient
import logging
import spaces

# ロガーの設定
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# モデル定義（ローカルモデルとAPIモデルの両方）
TEXT_GENERATION_MODELS = [
    {
        "name": "Llama-2",
        "description": "Known for its robust performance in content analysis",
        "type": "local",
        "model_path": "meta-llama/Llama-2-7b-hf"
    },
    {
        "name": "Mistral-7B",
        "description": "Offers precise and detailed text evaluation",
        "type": "local",
        "model_path": "mistralai/Mistral-7B-v0.1"
    },
    {
        "name": "Zephyr-7B",
        "description": "Specialized in understanding context and nuance",
        "type": "api",
        "model_id": "HuggingFaceH4/zephyr-7b-beta"
    }
]

CLASSIFICATION_MODELS = [
    {
        "name": "Toxic-BERT",
        "description": "Fine-tuned for toxic content detection",
        "type": "local",
        "model_path": "unitary/toxic-bert"
    }
]

class ModelManager:
    def __init__(self):
        self.tokenizers = {}
        self.pipelines = {}
        self.api_clients = {}
        self._initialize_api_clients()
        self._preload_local_models()

    def _initialize_api_clients(self):
        """Inference APIクライアントの初期化"""
        for model in TEXT_GENERATION_MODELS + CLASSIFICATION_MODELS:
            if model["type"] == "api" and "model_id" in model:
                logger.info(f"Initializing API client for {model['name']}")
                self.api_clients[model["model_id"]] = InferenceClient(
                    model["model_id"],
                    token=True  # HFトークンを使用
                )

    def _preload_local_models(self):
        """ローカルモデルを事前ロード"""
        logger.info("Preloading local models at application startup...")
        
        # テキスト生成モデル
        for model in TEXT_GENERATION_MODELS:
            if model["type"] == "local" and "model_path" in model:
                model_path = model["model_path"]
                try:
                    logger.info(f"Preloading text generation model: {model_path}")
                    self.tokenizers[model_path] = AutoTokenizer.from_pretrained(model_path)
                    self.pipelines[model_path] = pipeline(
                        "text-generation",
                        model=model_path,
                        tokenizer=self.tokenizers[model_path],
                        torch_dtype=torch.bfloat16,
                        trust_remote_code=True,
                        device_map="auto"
                    )
                    logger.info(f"Model preloaded successfully: {model_path}")
                except Exception as e:
                    logger.error(f"Error preloading model {model_path}: {str(e)}")
        
        # 分類モデル
        for model in CLASSIFICATION_MODELS:
            if model["type"] == "local" and "model_path" in model:
                model_path = model["model_path"]
                try:
                    logger.info(f"Preloading classification model: {model_path}")
                    self.tokenizers[model_path] = AutoTokenizer.from_pretrained(model_path)
                    self.pipelines[model_path] = pipeline(
                        "text-classification",
                        model=model_path,
                        tokenizer=self.tokenizers[model_path],
                        torch_dtype=torch.bfloat16,
                        trust_remote_code=True,
                        device_map="auto"
                    )
                    logger.info(f"Model preloaded successfully: {model_path}")
                except Exception as e:
                    logger.error(f"Error preloading model {model_path}: {str(e)}")

    @spaces.GPU
    def generate_text_local(self, model_path, text):
        """ローカルモデルでのテキスト生成"""
        try:
            logger.info(f"Running local text generation with {model_path}")
            outputs = self.pipelines[model_path](
                text,
                max_new_tokens=100,
                do_sample=False,
                num_return_sequences=1
            )
            return outputs[0]["generated_text"]
        except Exception as e:
            logger.error(f"Error in local text generation with {model_path}: {str(e)}")
            return f"Error: {str(e)}"

    def generate_text_api(self, model_id, text):
        """API経由でのテキスト生成"""
        try:
            logger.info(f"Running API text generation with {model_id}")
            response = self.api_clients[model_id].text_generation(
                text, 
                max_new_tokens=100, 
                temperature=0.7
            )
            return response
        except Exception as e:
            logger.error(f"Error in API text generation with {model_id}: {str(e)}")
            return f"Error: {str(e)}"

    @spaces.GPU
    def classify_text_local(self, model_path, text):
        """ローカルモデルでのテキスト分類"""
        try:
            logger.info(f"Running local classification with {model_path}")
            result = self.pipelines[model_path](text)
            return str(result)
        except Exception as e:
            logger.error(f"Error in local classification with {model_path}: {str(e)}")
            return f"Error: {str(e)}"

    def classify_text_api(self, model_id, text):
        """API経由でのテキスト分類"""
        try:
            logger.info(f"Running API classification with {model_id}")
            response = self.api_clients[model_id].text_classification(text)
            return str(response)
        except Exception as e:
            logger.error(f"Error in API classification with {model_id}: {str(e)}")
            return f"Error: {str(e)}"

    def run_models(self, text, selected_types):
        """選択されたタイプのモデルで分析を実行"""
        results = []
        
        # テキスト生成モデルの実行
        for model in TEXT_GENERATION_MODELS:
            if model["type"] in selected_types:
                if model["type"] == "local":
                    result = self.generate_text_local(model["model_path"], text)
                else:  # api
                    result = self.generate_text_api(model["model_id"], text)
                results.append(f"{model['name']}: {result}")
        
        # 分類モデルの実行
        for model in CLASSIFICATION_MODELS:
            if model["type"] in selected_types:
                if model["type"] == "local":
                    result = self.classify_text_local(model["model_path"], text)
                else:  # api
                    result = self.classify_text_api(model["model_id"], text)
                results.append(f"{model['name']}: {result}")
        
        # 結果リストの長さを調整
        while len(results) < len(TEXT_GENERATION_MODELS) + len(CLASSIFICATION_MODELS):
            results.append("")
        
        return results

class UIManager:
    def __init__(self, model_manager):
        self.model_manager = model_manager
        
    def create_ui(self):
        """UIの作成"""
        with gr.Blocks() as demo:
            # ヘッダー
            gr.Markdown("""
            # Toxic Eye (Class-based Version)
            This system evaluates the toxicity level of input text using both local models and Inference API.
            """)
            
            # 入力セクション
            with gr.Row():
                input_text = gr.Textbox(
                    label="Input Text",
                    placeholder="Enter text to analyze...",
                    lines=3
                )
            
            # フィルターセクション
            with gr.Row():
                filter_checkboxes = gr.CheckboxGroup(
                    choices=["local", "api"],
                    value=["local", "api"],
                    label="Filter Models",
                    info="Choose which types of models to use",
                    interactive=True
                )
            
            # 実行ボタン
            with gr.Row():
                invoke_button = gr.Button(
                    "Analyze Text",
                    variant="primary",
                    size="lg"
                )
            
            # モデル出力表示エリア
            all_outputs = []
            
            with gr.Tabs():
                # テキスト生成モデルのタブ
                with gr.Tab("Text Generation Models"):
                    for model in TEXT_GENERATION_MODELS:
                        with gr.Group():
                            gr.Markdown(f"### {model['name']} ({model['type']})")
                            output = gr.Textbox(
                                label=f"{model['name']} Output",
                                lines=5,
                                interactive=False,
                                info=model["description"]
                            )
                            all_outputs.append(output)
                
                # 分類モデルのタブ
                with gr.Tab("Classification Models"):
                    for model in CLASSIFICATION_MODELS:
                        with gr.Group():
                            gr.Markdown(f"### {model['name']} ({model['type']})")
                            output = gr.Textbox(
                                label=f"{model['name']} Output",
                                lines=5,
                                interactive=False,
                                info=model["description"]
                            )
                            all_outputs.append(output)
            
            # イベント接続
            invoke_button.click(
                fn=self.handle_invoke,
                inputs=[input_text, filter_checkboxes],
                outputs=all_outputs
            )
        
        return demo
    
    def handle_invoke(self, text, selected_types):
        """モデル実行をハンドリング"""
        return self.model_manager.run_models(text, selected_types)

class ToxicityApp:
    def __init__(self):
        self.model_manager = ModelManager()
        self.ui_manager = UIManager(self.model_manager)
    
    def run(self):
        """アプリを起動"""
        demo = self.ui_manager.create_ui()
        demo.launch()

def main():
    app = ToxicityApp()
    app.run()

if __name__ == "__main__":
    main()