Update app.py

app.py

@@ -43,240 +43,226 @@ CLASSIFICATION_MODELS = [
     }
 ]
 
-class ModelManager:
-    def __init__(self):
-        self.tokenizers = {}
-        self.pipelines = {}
-        self.api_clients = {}
-        self._initialize_api_clients()
-        self._preload_local_models()
+# グローバル変数でモデルとAPIクライアントを管理
+tokenizers = {}
+pipelines = {}
+api_clients = {}
 
-    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 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"
                 )
+                logger.info(f"Model preloaded successfully: {model_path}")
+            except Exception as e:
+                logger.error(f"Error preloading model {model_path}: {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)}")
+@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)}"
 
-    @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(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)}"
 
-    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(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)}"
 
-    @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(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)}"
 
-    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 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 run_models(self, text, selected_types):
-        """選択されたタイプのモデルで分析を実行"""
-        results = []
+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.
+        """)
         
-        # テキスト生成モデルの実行
-        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():
+            input_text = gr.Textbox(
+                label="Input Text",
+                placeholder="Enter text to analyze...",
+                lines=3
+            )
         
-        # 分類モデルの実行
-        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():
+            filter_checkboxes = gr.CheckboxGroup(
+                choices=["local", "api"],
+                value=["local", "api"],
+                label="Filter Models",
+                info="Choose which types of models to use",
+                interactive=True
+            )
         
-        # 結果リストの長さを調整
-        while len(results) < len(TEXT_GENERATION_MODELS) + len(CLASSIFICATION_MODELS):
-            results.append("")
+        # 実行ボタン
+        with gr.Row():
+            invoke_button = gr.Button(
+                "Analyze Text",
+                variant="primary",
+                size="lg"
+            )
         
-        return results
-
-class UIManager:
-    def __init__(self, model_manager):
-        self.model_manager = model_manager
+        # モデル出力表示エリア
+        all_outputs = []
         
-    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.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)
             
-            # モデル出力表示エリア
-            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
-            )
+            # 分類モデルのタブ
+            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)
         
-        return demo
+        # イベント接続
+        invoke_button.click(
+            fn=handle_invoke,
+            inputs=[input_text, filter_checkboxes],
+            outputs=all_outputs
+        )
     
-    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()
+    return demo
 
 def main():
-    app = ToxicityApp()
-    app.run()
+    # APIクライアントの初期化
+    initialize_api_clients()
+    
+    # ローカルモデルを事前ロード
+    preload_local_models()
+    
+    # UIを作成して起動
+    demo = create_ui()
+    demo.launch()
 
 if __name__ == "__main__":
     main()