.

app.py

@@ -1,11 +1,8 @@
 import gradio as gr
+import torch
+from transformers import AutoTokenizer, pipeline
 from huggingface_hub import InferenceClient
-from typing import List, Dict, Optional, Union
 import logging
-from dataclasses import dataclass
-from enum import Enum, auto
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification, pipeline
 import spaces
 
 # ロガーの設定
@@ -15,99 +12,67 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-# モデルの型定義
-class ModelType(Enum):
-    LOCAL = "local"
-    INFERENCE_API = "inference_api"
-
-@dataclass
-class ModelConfig:
-    name: str
-    description: str
-    type: ModelType
-    model_id: Optional[str] = None
-    model_path: Optional[str] = None
-
-# モデル定義を拡充
+# モデル定義（ローカルモデルとAPIモデルの両方）
 TEXT_GENERATION_MODELS = [
-    ModelConfig(
-        name="Zephyr-7B",
-        description="Specialized in understanding context and nuance",
-        type=ModelType.INFERENCE_API,
-        model_id="HuggingFaceH4/zephyr-7b-beta"
-    ),
-    ModelConfig(
-        name="Llama-2",
-        description="Known for its robust performance in content analysis",
-        type=ModelType.LOCAL,
-        model_path="meta-llama/Llama-2-7b-hf"
-    ),
-    ModelConfig(
-        name="Mistral-7B",
-        description="Offers precise and detailed text evaluation",
-        type=ModelType.LOCAL,
-        model_path="mistralai/Mistral-7B-v0.1"
-    )
+    {
+        "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 = [
-    ModelConfig(
-        name="Toxic-BERT",
-        description="Fine-tuned for toxic content detection",
-        type=ModelType.LOCAL,
-        model_path="unitary/toxic-bert"
-    )
+    {
+        "name": "Toxic-BERT",
+        "description": "Fine-tuned for toxic content detection",
+        "type": "local",
+        "model_path": "unitary/toxic-bert"
+    }
 ]
 
-class LocalModelManager:
+# GPU関連の装飾なしでクラスを定義
+class ModelManager:
     def __init__(self):
-        self.models = {}
         self.tokenizers = {}
         self.pipelines = {}
-        
-    def preload_models(self, model_paths, tasks=None):
-        """アプリケーション起動時にモデルを事前ロード"""
-        if tasks is None:
-            tasks = {}  # デフォルトは空の辞書
-            
-        logger.info("Preloading models at application startup...")
-        for model_path in model_paths:
-            task = tasks.get(model_path, "text-generation")
-            try:
-                logger.info(f"Preloading model: {model_path} for task: {task}")
-                self.tokenizers[model_path] = AutoTokenizer.from_pretrained(model_path)
-                
-                if task == "text-generation":
-                    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"
-                    )
-                else:  # classification
-                    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)}")
-                # 続行するが、エラーをログに記録
+        self.api_clients = {}
+        self._initialize_api_clients()
+        self._preload_local_models()
 
-    def load_model(self, model_path: str, task: str = "text-generation"):
-        """モデルが既にロードされているか確認し、なければロード"""
-        if model_path not in self.pipelines:
-            logger.info(f"Loading model on demand: {model_path}")
-            try:
-                self.tokenizers[model_path] = AutoTokenizer.from_pretrained(model_path)
-                
-                if task == "text-generation":
+    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,
@@ -116,7 +81,17 @@ class LocalModelManager:
                         trust_remote_code=True,
                         device_map="auto"
                     )
-                else:  # classification
+                    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,
@@ -125,277 +100,195 @@ class LocalModelManager:
                         trust_remote_code=True,
                         device_map="auto"
                     )
-                
-                logger.info(f"Model loaded successfully: {model_path}")
-            except Exception as e:
-                logger.error(f"Error loading model {model_path}: {str(e)}")
-                raise
-
-    @spaces.GPU
-    def _generate_text_sync(self, pipeline, text: str) -> str:
-        """同期的なテキスト生成の実行"""
-        outputs = pipeline(
-            text,
-            max_new_tokens=100,
-            do_sample=False,
-            num_return_sequences=1
-        )
-        return outputs[0]["generated_text"]
+                    logger.info(f"Model preloaded successfully: {model_path}")
+                except Exception as e:
+                    logger.error(f"Error preloading model {model_path}: {str(e)}")
 
-    def generate_text(self, model_path: str, text: str) -> str:
-        """テキスト生成の実行"""
-        if model_path not in self.pipelines:
-            self.load_model(model_path, "text-generation")
-        
+    def generate_text_local(self, model_path, text):
+        """ローカルモデルでのテキスト生成 (GPUデコレータはクラス外部で適用)"""
         try:
-            return self._generate_text_sync(self.pipelines[model_path], text)
+            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 text generation with {model_path}: {str(e)}")
-            raise
+            logger.error(f"Error in local text generation with {model_path}: {str(e)}")
+            return f"Error: {str(e)}"
 
-    @spaces.GPU
-    def _classify_text_sync(self, pipeline, text: str) -> str:
-        """同期的なテキスト分類の実行"""
-        result = pipeline(text)
-        return str(result)
-
-    def classify_text(self, model_path: str, text: str) -> str:
-        """テキスト分類の実行"""
-        if model_path not in self.pipelines:
-            self.load_model(model_path, "text-classification")
-        
+    def generate_text_api(self, model_id, text):
+        """API経由でのテキスト生成"""
         try:
-            return self._classify_text_sync(self.pipelines[model_path], text)
+            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 classification with {model_path}: {str(e)}")
-            raise
+            logger.error(f"Error in API text generation with {model_id}: {str(e)}")
+            return f"Error: {str(e)}"
 
-class ModelManager:
-    def __init__(self):
-        self.api_clients = {}
-        self.local_manager = LocalModelManager()
-        self._initialize_clients()
-        self._preload_local_models()
+    def classify_text_local(self, model_path, text):
+        """ローカルモデルでのテキスト分類 (GPUデコレータはクラス外部で適用)"""
+        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 _initialize_clients(self):
-        """Inference APIクライアントの初期化"""
-        for model in TEXT_GENERATION_MODELS + CLASSIFICATION_MODELS:
-            if model.type == ModelType.INFERENCE_API and model.model_id:
-                self.api_clients[model.model_id] = InferenceClient(
-                    model.model_id,
-                    token=True  # これによりHFトークンを使用
-                )
+    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 _preload_local_models(self):
-        """ローカルモデルの事前ロード"""
-        models_to_preload = []
-        tasks = {}
+    def run_models(self, text, selected_types):
+        """選択されたタイプのモデルで分析を実行"""
+        results = []
         
-        # テキスト生成モデルの追加
+        # テキスト生成モデルの実行
         for model in TEXT_GENERATION_MODELS:
-            if model.type == ModelType.LOCAL and model.model_path:
-                models_to_preload.append(model.model_path)
-                tasks[model.model_path] = "text-generation"
+            if model["type"] in selected_types:
+                if model["type"] == "local":
+                    # クラス外部でGPUデコレータが適用される前提
+                    result = gpu_wrapper_generate(self, 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 == ModelType.LOCAL and model.model_path:
-                models_to_preload.append(model.model_path)
-                tasks[model.model_path] = "text-classification"
+            if model["type"] in selected_types:
+                if model["type"] == "local":
+                    # クラス外部でGPUデコレータが適用される前提
+                    result = gpu_wrapper_classify(self, 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("")
         
-        # 事前ロード実行
-        self.local_manager.preload_models(models_to_preload, tasks)
-
-    def run_text_generation(self, text: str, selected_types: List[str]) -> List[str]:
-        """テキスト生成モデルの実行"""
-        results = []
-        for model in TEXT_GENERATION_MODELS:
-            if model.type.value in selected_types:
-                try:
-                    if model.type == ModelType.INFERENCE_API:
-                        logger.info(f"Running API text generation: {model.name}")
-                        response = self.api_clients[model.model_id].text_generation(
-                            text, max_new_tokens=100, temperature=0.7
-                        )
-                        results.append(f"{model.name}: {response}")
-                    else:
-                        logger.info(f"Running local text generation: {model.name}")
-                        response = self.local_manager.generate_text(model.model_path, text)
-                        results.append(f"{model.name}: {response}")
-                except Exception as e:
-                    logger.error(f"Error in {model.name}: {str(e)}")
-                    results.append(f"{model.name}: Error - {str(e)}")
-        return results
-
-    def run_classification(self, text: str, selected_types: List[str]) -> List[str]:
-        """分類モデルの実行"""
-        results = []
-        for model in CLASSIFICATION_MODELS:
-            if model.type.value in selected_types:
-                try:
-                    if model.type == ModelType.INFERENCE_API:
-                        logger.info(f"Running API classification: {model.name}")
-                        response = self.api_clients[model.model_id].text_classification(text)
-                        results.append(f"{model.name}: {response}")
-                    else:
-                        logger.info(f"Running local classification: {model.name}")
-                        response = self.local_manager.classify_text(model.model_path, text)
-                        results.append(f"{model.name}: {response}")
-                except Exception as e:
-                    logger.error(f"Error in {model.name}: {str(e)}")
-                    results.append(f"{model.name}: Error - {str(e)}")
         return results
 
-class UIComponents:
-    def __init__(self):
-        self.input_text = None
-        self.filter_checkboxes = None
-        self.invoke_button = None
-        self.gen_model_outputs = []
-        self.class_model_outputs = []
-        self.community_output = None
-
-    def create_header(self):
-        """ヘッダーセクションの作成"""
-        return gr.Markdown("""
-        # Toxic Eye
-        This system evaluates the toxicity level of input text using multiple approaches.
-        """)
-
-    def create_input_section(self):
-        """入力セクションの作成"""
-        with gr.Row():
-            self.input_text = gr.Textbox(
-                label="Input Text",
-                placeholder="Enter text to analyze...",
-                lines=3
-            )
-
-    def create_filter_section(self):
-        """フィルターセクションの作成"""
-        with gr.Row():
-            self.filter_checkboxes = gr.CheckboxGroup(
-                choices=[t.value for t in ModelType],
-                value=[t.value for t in ModelType],
-                label="Filter Models",
-                info="Choose which types of models to display",
-                interactive=True
-            )
-
-    def create_invoke_button(self):
-        """Invokeボタンの作成"""
-        with gr.Row():
-            self.invoke_button = gr.Button(
-                "Invoke Selected Models",
-                variant="primary",
-                size="lg"
-            )
+# クラスのメソッドをラップしてGPU装飾子を適用
+@spaces.GPU
+def gpu_wrapper_generate(manager, model_path, text):
+    return manager.generate_text_local(model_path, text)
 
-    def create_model_grid(self, models: List[ModelConfig]) -> List[Dict]:
-        """モデルグリッドの作成"""
-        outputs = []
-        with gr.Column() as container:
-            for i in range(0, len(models), 2):
-                with gr.Row() as row:
-                    for j in range(min(2, len(models) - i)):
-                        model = models[i + j]
-                        with gr.Column():
-                            with gr.Group() as group:
-                                gr.Markdown(f"### {model.name}")
-                                gr.Markdown(f"Type: {model.type.value}")
-                                output = gr.Textbox(
-                                    label="Model Output",
-                                    lines=5,
-                                    interactive=False,
-                                    info=model.description
-                                )
-                                outputs.append({
-                                    "type": model.type.value,
-                                    "name": model.name,
-                                    "output": output,
-                                    "group": group
-                                })
-        return outputs
+@spaces.GPU
+def gpu_wrapper_classify(manager, model_path, text):
+    return manager.classify_text_local(model_path, text)
 
-    def create_model_tabs(self):
-        """モデルタブの作成"""
-        with gr.Tabs():
-            with gr.Tab("Text Generation LLM"):
-                self.gen_model_outputs = self.create_model_grid(TEXT_GENERATION_MODELS)
-            with gr.Tab("Classification LLM"):
-                self.class_model_outputs = self.create_model_grid(CLASSIFICATION_MODELS)
-            with gr.Tab("Community (Not implemented)"):
-                with gr.Column():
-                    self.community_output = gr.Textbox(
-                        label="Related Community Topics",
-                        lines=5,
-                        interactive=False
-                    )
-
-class ToxicityApp:
-    def __init__(self):
-        self.ui = UIComponents()
-        self.model_manager = ModelManager()
-
-    def update_model_visibility(self, selected_types: List[str]) -> List[gr.update]:
-        """モデルの表示状態を更新"""
-        logger.info(f"Updating visibility for types: {selected_types}")
-        
-        updates = []
-        for outputs in [self.ui.gen_model_outputs, self.ui.class_model_outputs]:
-            for output in outputs:
-                visible = output["type"] in selected_types
-                logger.info(f"Model {output['name']} (type: {output['type']}): visible = {visible}")
-                updates.append(gr.update(visible=visible))
-        return updates
-
-    def handle_invoke(self, text: str, selected_types: List[str]) -> List[str]:
-        """Invokeボタンのハンドラ"""
-        gen_results = self.model_manager.run_text_generation(text, selected_types)
-        class_results = self.model_manager.run_classification(text, selected_types)
-        
-        # 結果リストの長さを調整
-        gen_results.extend([""] * (len(TEXT_GENERATION_MODELS) - len(gen_results)))
-        class_results.extend([""] * (len(CLASSIFICATION_MODELS) - len(class_results)))
+# UIの作成と管理のためのクラス
+class UIManager:
+    def __init__(self, model_manager):
+        self.model_manager = model_manager
         
-        return gen_results + class_results
-
     def create_ui(self):
         """UIの作成"""
         with gr.Blocks() as demo:
-            self.ui.create_header()
-            self.ui.create_input_section()
-            self.ui.create_filter_section()
-            self.ui.create_invoke_button()
-            self.ui.create_model_tabs()
-
-            # イベントハンドラの設定
-            self.ui.filter_checkboxes.change(
-                fn=self.update_model_visibility,
-                inputs=[self.ui.filter_checkboxes],
-                outputs=[
-                    output["group"]
-                    for outputs in [self.ui.gen_model_outputs, self.ui.class_model_outputs]
-                    for output in outputs
-                ]
-            )
-
-            self.ui.invoke_button.click(
+            # ヘッダー
+            gr.Markdown("""
+            # Toxic Eye (Class-based Version with GPU Wrappers)
+            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=[self.ui.input_text, self.ui.filter_checkboxes],
-                outputs=[
-                    output["output"]
-                    for outputs in [self.ui.gen_model_outputs, self.ui.class_model_outputs]
-                    for output in outputs
-                ]
+                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()
-    demo = app.create_ui()
-    demo.launch()
+    app.run()
 
 if __name__ == "__main__":
     main()