Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import torch
 from transformers import AutoTokenizer, pipeline
+from huggingface_hub import InferenceClient
 import logging
 import spaces
 
@@ -11,17 +12,25 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-# シンプルなモデル定義（3つのローカルモデル）
+# モデル定義（ローカルモデルと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"
     }
 ]
 
@@ -29,59 +38,73 @@ CLASSIFICATION_MODELS = [
     {
         "name": "Toxic-BERT",
         "description": "Fine-tuned for toxic content detection",
+        "type": "local",
         "model_path": "unitary/toxic-bert"
     }
 ]
 
-# グローバル変数でモデルとトークナイザを管理
+# グローバル変数でモデルとAPIクライアントを管理
 tokenizers = {}
 pipelines = {}
+api_clients = {}
 
-def preload_models():
-    """アプリケーション起動時にモデルを事前ロード"""
-    logger.info("Preloading models at application startup...")
+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:
-        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)}")
+        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:
-        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)}")
+        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)}")
 
 @spaces.GPU
-def generate_text(model_path, text):
-    """テキスト生成の実行"""
+def generate_text_local(model_path, text):
+    """ローカルモデルでのテキスト生成"""
     try:
-        logger.info(f"Running text generation with {model_path}")
+        logger.info(f"Running local text generation with {model_path}")
         outputs = pipelines[model_path](
             text,
             max_new_tokens=100,
@@ -90,35 +113,69 @@ def generate_text(model_path, text):
         )
         return outputs[0]["generated_text"]
     except Exception as e:
-        logger.error(f"Error in text generation with {model_path}: {str(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)}"
 
 @spaces.GPU
-def classify_text(model_path, text):
-    """テキスト分類の実行"""
+def classify_text_local(model_path, text):
+    """ローカルモデルでのテキスト分類"""
     try:
-        logger.info(f"Running classification with {model_path}")
+        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 classification with {model_path}: {str(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 handle_invoke(text):
-    """すべてのモデルで分析を実行"""
+def handle_invoke(text, selected_types):
+    """選択されたタイプのモデルで分析を実行"""
     results = []
     
     # テキスト生成モデルの実行
     for model in TEXT_GENERATION_MODELS:
-        model_path = model["model_path"]
-        result = generate_text(model_path, text)
-        results.append(result)
+        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:
-        model_path = model["model_path"]
-        result = classify_text(model_path, text)
-        results.append(result)
+        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
 
@@ -127,8 +184,8 @@ def create_ui():
     with gr.Blocks() as demo:
         # ヘッダー
         gr.Markdown("""
-        # Toxic Eye (3 Models Version)
-        This system evaluates the toxicity level of input text using 3 local models.
+        # Toxic Eye (Local + API Version)
+        This system evaluates the toxicity level of input text using both local models and Inference API.
         """)
         
         # 入力セクション
@@ -139,6 +196,16 @@ def create_ui():
                 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(
@@ -148,48 +215,50 @@ def create_ui():
             )
         
         # モデル出力表���エリア
-        gen_outputs = []
-        class_outputs = []
+        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']}")
+                        gr.Markdown(f"### {model['name']} ({model['type']})")
                         output = gr.Textbox(
                             label=f"{model['name']} Output",
                             lines=5,
                             interactive=False,
                             info=model["description"]
                         )
-                        gen_outputs.append(output)
+                        all_outputs.append(output)
             
             # 分類モデルのタブ
             with gr.Tab("Classification Models"):
                 for model in CLASSIFICATION_MODELS:
                     with gr.Group():
-                        gr.Markdown(f"### {model['name']}")
+                        gr.Markdown(f"### {model['name']} ({model['type']})")
                         output = gr.Textbox(
                             label=f"{model['name']} Output",
                             lines=5,
                             interactive=False,
                             info=model["description"]
                         )
-                        class_outputs.append(output)
+                        all_outputs.append(output)
         
         # イベント接続
         invoke_button.click(
             fn=handle_invoke,
-            inputs=[input_text],
-            outputs=gen_outputs + class_outputs
+            inputs=[input_text, filter_checkboxes],
+            outputs=all_outputs
         )
     
     return demo
 
 def main():
-    # モデルを事前ロード
-    preload_models()
+    # APIクライアントの初期化
+    initialize_api_clients()
+    
+    # ローカルモデルを事前ロード
+    preload_local_models()
     
     # UIを作成して起動
     demo = create_ui()