app.py

import os
import gradio as gr
from llama_cpp import Llama
import spaces
import time

# モデルパスの設定
MODEL_PATH = "teenemo-reasoning-v2-q4_k_m.gguf"
MODEL_URL = "https://huggingface.co/YUGOROU/TeenEmo-Reasoning-v2-Q4_K_M-GGUF/resolve/main/teenemo-reasoning-v2-q4_k_m.gguf"

# システムプロンプト - モデルの特性を活かす設定
SYSTEM_PROMPT = """
あなたは10代の若者の悩みや感情を理解し、推論能力を発揮して的確なアドバイスや説明を提供するアシスタントです。
中学生や高校生が抱える勉強、友人関係、将来の不安などについて共感的に対応し、
論理的思考に基づいた解決策を提案することができます。
質問に対しては正確で分かりやすい説明を心がけ、ユーザーの年齢や理解度に合わせた回答をしてください。
"""

# モデルが存在しない場合はダウンロード
def download_model_if_needed():
    if not os.path.exists(MODEL_PATH):
        try:
            import urllib.request
            print(f"モデルをダウンロード中: {MODEL_URL}")
            urllib.request.urlretrieve(MODEL_URL, MODEL_PATH)
            print("ダウンロード完了")
            return True
        except Exception as e:
            print(f"モデルのダウンロード中にエラーが発生しました: {e}")
            return False
    return True

# モデルの設定とロード
@spaces.GPU(duration=120)  # ZeroGPUを使用、最大120秒
def load_model():
    if not download_model_if_needed():
        raise Exception("モデルのダウンロードに失敗しました")

    return Llama(
        model_path=MODEL_PATH,
        n_ctx=2048,  # コンテキストウィンドウサイズ
        n_gpu_layers=-1,  # 全レイヤーをGPUにロード
        n_threads=8,  # CPUスレッド数
        verbose=False,  # 詳細なログを非表示
    )

# チャットの履歴フォーマット関数
def format_prompt(message, history):
    prompt = SYSTEM_PROMPT + "\n\n"

    # 過去の会話履歴を追加
    for user_msg, bot_msg in history:
        prompt += f"User: {user_msg}\nAssistant: {bot_msg}\n\n"

    # 新しいメッセージを追加
    prompt += f"User: {message}\nAssistant: "

    return prompt

# チャット処理関数（ストリーミング対応）
def chat_function(message, history):
    # モデルが既にロードされているか確認し、ロードされていなければロード
    if not hasattr(chat_function, "llm"):
        try:
            chat_function.llm = load_model()
        except Exception as e:
            return f"モデルのロードに失敗しました: {e}"

    # プロンプトを作成
    prompt = format_prompt(message, history)

    # 推論パラメータ
    generation_config = {
        "max_tokens": 1024,
        "temperature": 0.7,
        "top_p": 0.95,
        "top_k": 40,
        "repeat_penalty": 1.1,
        "stop": ["User:", "\n\n"]
    }

    try:
        # 推論実行（ストリーミングなし）
        output = chat_function.llm.create_completion(
            prompt=prompt,
            **generation_config
        )

        # 生成されたテキストを取得
        generated_text = output["choices"][0]["text"]

        return generated_text
    except Exception as e:
        return f"推論中にエラーが発生しました: {e}"

# ストリーミング対応チャット関数（必要に応じて使用）
def chat_function_streaming(message, history):
    # モデルが既にロードされているか確認し、ロードされていなければロード
    if not hasattr(chat_function_streaming, "llm"):
        try:
            chat_function_streaming.llm = load_model()
        except Exception as e:
            yield f"モデルのロードに失敗しました: {e}"
            return

    # プロンプトを作成
    prompt = format_prompt(message, history)

    # 推論パラメータ
    generation_config = {
        "max_tokens": 1024,
        "temperature": 0.7,
        "top_p": 0.95,
        "top_k": 40,
        "repeat_penalty": 1.1,
        "stop": ["User:", "\n\n"],
        "stream": True
    }

    try:
        # ストリーミング推論の開始
        partial_text = ""
        for chunk in chat_function_streaming.llm.create_completion(
            prompt=prompt,
            **generation_config
        ):
            if chunk["choices"][0]["text"]:
                partial_text += chunk["choices"][0]["text"]
                yield partial_text
    except Exception as e:
        yield f"推論中にエラーが発生しました: {e}"

# Gradioインターフェースの作成
def create_demo():
    with gr.Blocks(theme="soft") as demo:
        gr.Markdown("# TeenEmo-Reasoning-v2 Chatbot")
        gr.Markdown("""
        10代の感情と推論能力に焦点を当てたチャットボットです。
        Qwen/Qwen3-4B-Baseベースのモデルを使用し、ZeroGPUで高速推論を実現しています。

        質問や悩み事を日本語で入力してください。
        """)

        chatbot = gr.Chatbot(
            height=500,
            show_copy_button=True,
            layout="panel"
        )

        msg = gr.Textbox(
            placeholder="メッセージを入力してください...",
            container=False,
            scale=7
        )

        with gr.Row():
            submit_btn = gr.Button("送信", variant="primary")
            clear_btn = gr.Button("クリア")

        # 例示メッセージ
        example_queries = [
            "学校でクラスメイトに無視されて悲しい気持ちです。どうしたら良いですか？",
            "中学校の数学の宿題に困っています。二次方程式の解き方を教えてもらえますか？",
            "友達との関係について悩んでいます。アドバイスをください。",
            "プログラミングを始めたいと思っています。何から始めるべきですか？",
            "テストが近いのに集中できません。どうすれば勉強に集中できますか？",
            "将来どんな職業に就くべきか悩んでいます。アドバイスください。"
        ]

        gr.Examples(
            examples=example_queries,
            inputs=msg
        )

        # イベントの設定
        msg.submit(chat_function, [msg, chatbot], [chatbot, msg]).then(
            lambda: "", None, [msg]
        )

        submit_btn.click(chat_function, [msg, chatbot], [chatbot, msg]).then(
            lambda: "", None, [msg]
        )

        clear_btn.click(lambda: None, None, [chatbot], queue=False)

        # 初期メッセージを表示
        demo.load(
            lambda: [[None, "こんにちは！何について話したいですか？あなたの質問や悩みに答えます。"]],
            None, 
            [chatbot]
        )

    return demo

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