app.py

import gradio as gr
from transformers import pipeline
import time

# 利用可能なモデルのリスト（公開されていて比較的小さいモデル）
AVAILABLE_MODELS = {
    "新 独自学習Qwen2-1.5B": "eyepyon/rc3qwen2-1.5b-finetuned",
    "新 独自学習Llama-3.2-1B": "eyepyon/rc3llama-3.2-1b-finetuned",
    "新 独自学習Gemma-2-2B": "eyepyon/rc3gemma-2-2b-finetuned",
    "新 独自学習Phi4-mini-reasoning": "eyepyon/rc3phi-4-mini-reasoning-finetuned",
    "新 独自学習DeepSeek-R1-1.5B": "eyepyon/rc3deepseek-r1-1.5b-finetuned",
    "旧 独自学習Qwen2-1.5B": "eyepyon/rcqwen2-merged-final",
    "旧 独自学習Llama-3.2-1B": "eyepyon/rclama32-merged-final",
    "旧 独自学習Gemma-2-9B": "eyepyon/rcgemma2-merged-final",
    "旧 独自学習Phi4-mini-reasoning": "eyepyon/rcrcPhi4miniR-merged-final",
    "GPT2-Small": "gpt2",
    "DistilGPT2": "distilgpt2",
}
#    "独自学習phi2": "eyepyon/rcphi2",
#      "BERT-Base": "bert-base-uncased",  # マスク補完用
#    "FLAN-T5-Small": "google/flan-t5-small",  # テキスト生成
#     "OPT-125M": "facebook/opt-125m",


# グローバル変数
current_model = None
current_model_name = None
generator = None

# モデルをロードする関数
def load_model(model_key):
    global generator, current_model, current_model_name

    if model_key not in AVAILABLE_MODELS:
        return f"エラー: モデル '{model_key}' は利用できません。"

    model_path = AVAILABLE_MODELS[model_key]

    try:
        # ロード開始時刻
        start_time = time.time()
        print(f"モデル {model_key} ({model_path}) をロード中...")

        # モデルタイプに基づいてパイプラインを選択
        if "t5" in model_path.lower():
            # T5モデルはtext2textタスク
            generator = pipeline("text2text-generation", model=model_path)
        elif "bert" in model_path.lower():
            # BERTモデルはfill-maskタスク
            generator = pipeline("fill-mask", model=model_path)
        else:
            # デフォルトはtext-generation
            generator = pipeline("text-generation", model=model_path)

        # 現在のモデル名を保存
        current_model = model_path
        current_model_name = model_key

        # ロード時間
        load_time = time.time() - start_time
        success_msg = f"✅ {model_key} モデルのロードが完了しました！（{load_time:.1f}秒）"
        print(success_msg)
        return success_msg

    except Exception as e:
        error_msg = f"❌ モデルのロード中にエラーが発生しました: {str(e)}"
        print(error_msg)
        return error_msg

# テキスト生成/処理関数
def process_text(prompt, model_key):
    global generator, current_model, current_model_name

    # 入力チェック
    if not prompt.strip():
        return "テキストを入力してください。"

    # モデルがロードされていない場合
    if generator is None:
        return "モデルがまだロードされていません。まずモデルを選択してロードしてください。"

    # 選択されたモデルと現在ロードされているモデルが異なる場合
    if current_model_name != model_key:
        return f"現在ロードされているのは '{current_model_name}' モデルです。'{model_key}' モデルを使用するには、まずロードしてください。"

    try:
        # モデルに応じた処理
        if "bert" in current_model.lower():
            # BERTのマスク補完
            # プロンプトに[MASK]がない場合は追加
            if "[MASK]" not in prompt:
                prompt += " [MASK]"

            result = generator(prompt)
            if isinstance(result, list):
                # トップ3の候補を表示
                output = f"入力: {prompt}\n\n候補:\n"
                for i, res in enumerate(result[:3]):
                    output += f"{i+1}. {res['sequence']} (確率: {res['score']:.2f})\n"
                return output
            else:
                return f"結果: {result['sequence']}"

        elif "t5" in current_model.lower():
            # T5形式のテキスト生成
            result = generator(prompt, max_length=100)
            return f"入力: {prompt}\n\n生成結果: {result[0]['generated_text']}"

        else:
            # 通常のテキスト生成
            result = generator(prompt, max_length=100, do_sample=True, temperature=0.7)
            return f"入力: {prompt}\n\n生成結果: {result[0]['generated_text']}"

    except Exception as e:
        return f"処理中にエラーが発生しました: {str(e)}"

# サンプルテキストをセットする関数
def set_gpt_sample1():
    return "こんにちは、私の名前は"

def set_gpt_sample2():
    return "人工知能の将来について"

def set_bert_sample1():
    return "私は[MASK]が好きです。"

def set_bert_sample2():
    return "今日の[MASK]はとても良いです。"

def set_t5_sample1():
    return "translate English to German: Hello"

def set_t5_sample2():
    return "summarize: AI models process text patterns"

# 単純化したGradioインターフェース
with gr.Blocks(title="モデルテスター") as demo:
    gr.Markdown("# モデルテスター")
    gr.Markdown("複数のモデルをロードしてテストできます。まずモデルを選択してロードし、テキスト処理を行ってください。")

    # モデル選択と処理部分
    with gr.Row():
        # 左側のカラム - モデル選択
        with gr.Column():
            model_dropdown = gr.Dropdown(
                choices=list(AVAILABLE_MODELS.keys()),
                label="テストするモデル",
                value="DistilGPT2"
            )

            load_button = gr.Button("モデルをロード")
            model_info = gr.Textbox(label="モデル情報", value="モデルが選択されていません")

            # モデルタイプごとのヒント
#             gr.Markdown("### モデル別ヒント")
#             gr.Markdown("**GPT2・DistilGPT2**: テキスト生成モデル")
#             gr.Markdown("**BERT**: [MASK]部分を補完するモデル")
#             gr.Markdown("**T5**: 翻訳・要約などのタスク向けモデル")

        # 右側のカラム - テキスト入力と処理
        with gr.Column():
            input_text = gr.Textbox(
                label="入力テキスト",
                placeholder="テキストを入力してください...",
                lines=3
            )

            # サンプルテキストボタン
#             gr.Markdown("### サンプルテキスト")
#             with gr.Row():
#                 gpt_btn1 = gr.Button("GPT用例1")
#                 gpt_btn2 = gr.Button("GPT用例2")
#
#             with gr.Row():
#                 bert_btn1 = gr.Button("BERT用例1")
#                 bert_btn2 = gr.Button("BERT用例2")
#
#             with gr.Row():
#                 t5_btn1 = gr.Button("T5用例1")
#                 t5_btn2 = gr.Button("T5用例2")

            process_button = gr.Button("実行")
            output_text = gr.Textbox(label="出力結果", lines=10)

    # モデルロードイベント
    load_button.click(
        fn=load_model,
        inputs=[model_dropdown],
        outputs=[model_info]
    )

    # テキスト処理イベント
    process_button.click(
        fn=process_text,
        inputs=[input_text, model_dropdown],
        outputs=[output_text]
    )

    # サンプルボタンイベント - GPT系
#     gpt_btn1.click(fn=set_gpt_sample1, inputs=None, outputs=input_text)
#     gpt_btn2.click(fn=set_gpt_sample2, inputs=None, outputs=input_text)

    # サンプルボタンイベント - BERT
#     bert_btn1.click(fn=set_bert_sample1, inputs=None, outputs=input_text)
#     bert_btn2.click(fn=set_bert_sample2, inputs=None, outputs=input_text)

    # サンプルボタンイベント - T5
#     t5_btn1.click(fn=set_t5_sample1, inputs=None, outputs=input_text)
#     t5_btn2.click(fn=set_t5_sample2, inputs=None, outputs=input_text)

# アプリケーション起動
if __name__ == "__main__":
    demo.launch()