Update app.py

app.py

@@ -43,226 +43,240 @@ CLASSIFICATION_MODELS = [
     }
 ]
 
-# グローバル変数でモデルとAPIクライアントを管理
-tokenizers = {}
-pipelines = {}
-api_clients = {}
+class ModelManager:
+    def __init__(self):
+        self.tokenizers = {}
+        self.pipelines = {}
+        self.api_clients = {}
+        self._initialize_api_clients()
+        self._preload_local_models()
 
-def initialize_api_clients():
-    """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']}")
-            api_clients[model["model_id"]] = InferenceClient(
-                model["model_id"],
-                token=True  # HFトークンを使用
-            )
-
-def preload_local_models():
-    """ローカルモデルを事前ロード"""
-    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}")
-                tokenizers[model_path] = AutoTokenizer.from_pretrained(model_path)
-                pipelines[model_path] = pipeline(
-                    "text-generation",
-                    model=model_path,
-                    tokenizer=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}")
-                tokenizers[model_path] = AutoTokenizer.from_pretrained(model_path)
-                pipelines[model_path] = pipeline(
-                    "text-classification",
-                    model=model_path,
-                    tokenizer=tokenizers[model_path],
-                    torch_dtype=torch.bfloat16,
-                    trust_remote_code=True,
-                    device_map="auto"
+    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トークンを使用
                 )
-                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(model_path, text):
-    """ローカルモデルでのテキスト生成"""
-    try:
-        logger.info(f"Running local text generation with {model_path}")
-        outputs = 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 _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)}")
 
-def generate_text_api(model_id, text):
-    """API経由でのテキスト生成"""
-    try:
-        logger.info(f"Running API text generation with {model_id}")
-        response = 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 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)}"
 
-@spaces.GPU
-def classify_text_local(model_path, text):
-    """ローカルモデルでのテキスト分類"""
-    try:
-        logger.info(f"Running local classification with {model_path}")
-        result = 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 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)}"
 
-def classify_text_api(model_id, text):
-    """API経由でのテキスト分類"""
-    try:
-        logger.info(f"Running API classification with {model_id}")
-        response = 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)}"
+    @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 handle_invoke(text, selected_types):
-    """選択されたタイプのモデルで分析を実行"""
-    results = []
-    
-    # テキスト生成モデルの実行
-    for model in TEXT_GENERATION_MODELS:
-        if model["type"] in selected_types:
-            if model["type"] == "local":
-                result = generate_text_local(model["model_path"], text)
-            else:  # api
-                result = 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 = classify_text_local(model["model_path"], text)
-            else:  # api
-                result = 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
+    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 create_ui():
-    """UIの作成"""
-    with gr.Blocks() as demo:
-        # ヘッダー
-        gr.Markdown("""
-        # Toxic Eye (Local + API Version)
-        This system evaluates the toxicity level of input text using both local models and Inference API.
-        """)
+    def run_models(self, text, selected_types):
+        """選択されたタイプのモデルで分析を実行"""
+        results = []
         
-        # 入力セクション
-        with gr.Row():
-            input_text = gr.Textbox(
-                label="Input Text",
-                placeholder="Enter text to analyze...",
-                lines=3
-            )
+        # テキスト生成モデルの実行
+        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}")
         
-        # フィルターセクション
-        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
-            )
+        # 分類モデルの実行
+        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}")
         
-        # 実行ボタン
-        with gr.Row():
-            invoke_button = gr.Button(
-                "Analyze Text",
-                variant="primary",
-                size="lg"
-            )
+        # 結果リストの長さを調整
+        while len(results) < len(TEXT_GENERATION_MODELS) + len(CLASSIFICATION_MODELS):
+            results.append("")
         
-        # モデル出力表示エリア
-        all_outputs = []
+        return results
+
+class UIManager:
+    def __init__(self, model_manager):
+        self.model_manager = model_manager
         
-        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)
+    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"
+                )
             
-            # 分類モデルのタブ
-            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)
+            # モデル出力表示エリア
+            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
+            )
         
-        # イベント接続
-        invoke_button.click(
-            fn=handle_invoke,
-            inputs=[input_text, filter_checkboxes],
-            outputs=all_outputs
-        )
+        return demo
     
-    return demo
+    def handle_invoke(self, text, selected_types):
+        """モデル実行をハンドリング"""
+        return self.model_manager.run_models(text, selected_types)
 
-def main():
-    # APIクライアントの初期化
-    initialize_api_clients()
+class ToxicityApp:
+    def __init__(self):
+        self.model_manager = ModelManager()
+        self.ui_manager = UIManager(self.model_manager)
     
-    # ローカルモデルを事前ロード
-    preload_local_models()
-    
-    # UIを作成して起動
-    demo = create_ui()
-    demo.launch()
+    def run(self):
+        """アプリを起動"""
+        demo = self.ui_manager.create_ui()
+        demo.launch()
+
+def main():
+    app = ToxicityApp()
+    app.run()
 
 if __name__ == "__main__":
     main()