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"
    }
]

# グローバル変数でモデルとAPIクライアントを管理
tokenizers = {}
pipelines = {}
api_clients = {}

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)}")

@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 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_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 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, 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 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.
        """)
        
        # 入力セクション
        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=handle_invoke,
            inputs=[input_text, filter_checkboxes],
            outputs=all_outputs
        )
    
    return demo

def main():
    # APIクライアントの初期化
    initialize_api_clients()
    
    # ローカルモデルを事前ロード
    preload_local_models()
    
    # UIを作成して起動
    demo = create_ui()
    demo.launch()

if __name__ == "__main__":
    main()