app.py

import os
import re
import gradio as gr
import spaces

try:
    from llama_cpp import Llama
    USE_LLAMA_CPP = True
    print("llama-cpp-python is installed. Will use GGUF inference.")
except ImportError:
    USE_LLAMA_CPP = False
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer
    print("llama-cpp-python is not installed. Falling back to transformers inference.")

# ---------------------------------------------------------
# AI Model Initialization
# ---------------------------------------------------------
if USE_LLAMA_CPP:
    GGUF_FILENAME = "production_mindmap_model.gguf"
    if os.path.exists(f"./{GGUF_FILENAME}"):
        GGUF_PATH = f"./{GGUF_FILENAME}"
        print(f"ローカルのモデルを使用します: {GGUF_PATH}")
    else:
        from huggingface_hub import hf_hub_download
        MODEL_REPO_ID = os.environ.get("MODEL_REPO_ID", "kazutab/mindmap-studio-model")
        if not MODEL_REPO_ID:
            raise ValueError("ローカルにモデルが見つからず、MODEL_REPO_ID環境変数も設定されていません。")
        print(f"Hugging Face ({MODEL_REPO_ID}) からモデルをダウンロード中...")
        GGUF_PATH = hf_hub_download(repo_id=MODEL_REPO_ID, filename="backend/production_mindmap_model.gguf")
    
    print("Loading AI Model (GGUF) into memory...")
    model = Llama(
        model_path=GGUF_PATH,
        n_ctx=2048,
        n_gpu_layers=-1 # Use all GPU layers if available
    )
    print("AIモデル(GGUF)の起動完了！")
else:
    MERGED_MODEL_PATH = "./production_mindmap_model_merged"
    print("Loading AI Model (Transformers FP16) into memory...")
    tokenizer = AutoTokenizer.from_pretrained(MERGED_MODEL_PATH)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    model = AutoModelForCausalLM.from_pretrained(
        MERGED_MODEL_PATH, 
        torch_dtype=torch.float16, 
        device_map="auto"
    )
    model.eval()
    print("AIモデル(Transformers)の起動完了！")

STRICT_SYSTEM_PROMPT = """あなたは極めて優秀で厳密な情報抽出アシスタントです。入力文章の論理構造を正確に読み取り、Markdown形式の目次（マインドマップ）を出力してください。



【厳格なルール】

1. 企業名、人物名、数値などの固有名詞や事実関係は、入力文章から「正確に」抽出すること。絶対に捏造したり、他の主語と混同したりしてはなりません。

2. ただし、情報を階層化して分かりやすく整理するための「一般的なカテゴリ名（例：背景、特徴、課題、概要など）」を親見出しとして補足することは許可します。

3. あなたの推測や外部知識を一切混ぜず、入力文章内の事実のみを構造化してください。"""

@spaces.GPU(duration=120)
def generate_mindmap(input_text: str):
    input_text = input_text.strip()
    if not input_text:
        return "<div style='color: red; padding: 20px;'>文章を入力してください。</div>"

    print(f"推論を開始します（文字数: {len(input_text)}文字）")

    USER_PROMPT = f"""以下の文章から論理構造を抽出し、Markdown形式の目次（マインドマップ）を出力してください。



【出力時の厳守ルール（違反厳禁）】

1. 否定表現の厳守：「〜しない」「過度に依存しない」などの否定表現を絶対に見落とさず、意味を逆転させないこと。

2. 創作の禁止：記事に明記されていない具体的な行動や予定（例：「〜への参加」「〜の強化を目指す」など）を勝手に推測して付け足さないこと。

3. 事実の完全一致：抽出した内容が、元の文章の事実と完全に一致していることのみを出力すること。



入力文章:

{input_text}"""
    messages = [
        {"role": "system", "content": STRICT_SYSTEM_PROMPT},
        {"role": "user", "content": USER_PROMPT}
    ]
    
    if USE_LLAMA_CPP:
        response = model.create_chat_completion(
            messages=messages,
            max_tokens=1024,
            temperature=0.0,
            repeat_penalty=1.1
        )
        generated_markdown = response['choices'][0]['message']['content'].strip()
    else:
        prompt = tokenizer.apply_chat_template(
            messages, 
            tokenize=False, 
            add_generation_prompt=True
        )
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=1024,
                do_sample=False,
                repetition_penalty=1.1,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id
            )
        input_length = inputs["input_ids"].shape[1]
        generated_tokens = outputs[0][input_length:]
        generated_markdown = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
    
    generated_markdown = re.sub(r'\s*(#+ )', r'\n\1', generated_markdown).strip()
    
    if not generated_markdown.startswith('#'):
        if '##' in generated_markdown or '###' in generated_markdown:
            generated_markdown = "# マインドマップ\n" + generated_markdown
        else:
            generated_markdown = "# マインドマップ\n## 抽出結果\n- " + generated_markdown.replace('\n', '\n- ')

    # Markmap 用のHTMLを動的に構築して返す
    html_output = f"""

    <div class="markmap" style="width: 100%; height: 100%; min-height: 500px;">

      <script type="text/template">

      {generated_markdown}

      </script>

    </div>

    <script>

      // GradioのHTML更新後にMarkmapを強制再レンダリングする

      setTimeout(() => {{

          if (window.markmap && window.markmap.autoLoader) {{

              window.markmap.autoLoader.renderAll();

          }}

      }}, 100);

    </script>

    """
    return html_output

# ---------------------------------------------------------
# UI Construction (Native Gradio)
# ---------------------------------------------------------
# オリジナルのCSSを読み込み、Gradio用のオーバーライドを追記する
with open("frontend/style.css", "r", encoding="utf-8") as f:
    base_css = f.read()

gradio_overrides = """

/* Gradio特有の不要な余白や枠線を完全に無効化 */

.gradio-container {

    max-width: 100% !important;

    padding: 0 !important;

    margin: 0 !important;

    border: none !important;

    background: transparent !important;

}

footer { display: none !important; }

#root { padding: 0 !important; }



/* オリジナルのレイアウトをGradioのRow/Columnに強制適用 */

.app-layout {

    gap: 0 !important;

    flex-wrap: nowrap !important;

    margin: 0 !important;

}

.sidebar {

    min-width: 360px !important;

    max-width: 360px !important;

    padding: 0 !important;

    border-radius: 0 !important;

    gap: 0 !important;

    border-right: 1px solid var(--border-color) !important;

}

.sidebar-content {

    padding: 24px !important;

    gap: 20px !important;

}

.canvas-area {

    border-radius: 0 !important;

    border: none !important;

    padding: 0 !important;

    margin: 0 !important;

}



/* Gradioのコンポーネントが勝手に作る枠線を消す */

.form { border: none !important; background: transparent !important; box-shadow: none !important; }

.block { padding: 0 !important; margin: 0 !important; border: none !important; box-shadow: none !important; background: transparent !important; }



/* テキストエリアとボタンにオリジナルのスタイルを適用 */

#text-input textarea {

    border-radius: 8px !important;

    height: 100% !important;

    min-height: 200px !important;

    border: 1px solid var(--border-color) !important;

    padding: 16px !important;

    font-size: 14px !important;

    box-shadow: 0 1px 2px rgba(0,0,0,0.02) !important;

}



button.btn-primary {

    background-color: #000 !important;

    color: #fff !important;

    border-radius: 6px !important;

    padding: 12px 16px !important;

    font-size: 14px !important;

    font-weight: 500 !important;

    display: flex !important;

    align-items: center !important;

    justify-content: center !important;

    gap: 8px !important;

}

button.btn-primary::before {

    content: url('data:image/svg+xml;utf8,<svg width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="white" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><polygon points="5 3 19 12 5 21 5 3"></polygon></svg>');

    display: inline-block;

    width: 16px;

    height: 16px;

    margin-right: 4px;

}

button.btn-primary:hover {

    background-color: #333 !important;

}

"""

custom_css = base_css + "\n" + gradio_overrides

head_scripts = """

<script src="https://cdn.jsdelivr.net/npm/d3@7"></script>

<script src="https://cdn.jsdelivr.net/npm/markmap-view"></script>

<script src="https://cdn.jsdelivr.net/npm/markmap-autoloader"></script>

"""

with gr.Blocks(css=custom_css, head=head_scripts, title="MindMap Studio") as demo:
    with gr.Row(elem_classes="app-layout"):
        with gr.Column(elem_classes="sidebar"):
            # オリジナルのサイドバーヘッダー（SVGロゴ）を復元
            gr.HTML("""

            <div class="sidebar-header" style="padding: 24px;">

                <div class="logo" style="display: flex; align-items: center; gap: 8px; font-weight: 700; font-size: 18px;">

                    <svg width="24" height="24" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M21 16V8a2 2 0 0 0-1-1.73l-7-4a2 2 0 0 0-2 0l-7 4A2 2 0 0 0 3 8v8a2 2 0 0 0 1 1.73l7 4a2 2 0 0 0 2 0l7-4A2 2 0 0 0 21 16z"></path><polyline points="3.27 6.96 12 12.01 20.73 6.96"></polyline><line x1="12" y1="22.08" x2="12" y2="12"></line></svg>

                    <span>MindMap Studio</span>

                </div>

            </div>

            """)
            
            with gr.Column(elem_classes="sidebar-content"):
                gr.HTML('<div class="input-group"><label style="font-size: 13px; font-weight: 500; color: #666666;">Source Text</label></div>')
                
                text_input = gr.Textbox(
                    lines=10, 
                    placeholder="議事録や講義のテキストをペーストしてください...",
                    show_label=False,
                    container=False,
                    elem_id="text-input"
                )
                
                submit_btn = gr.Button("マップを生成", elem_classes="btn-primary")
            
            # オリジナルのフッターを復元
            gr.HTML("""

            <div class="sidebar-footer" style="padding: 16px 24px; border-top: 1px solid #eaeaea; font-size: 12px; color: #a1a1aa; margin-top: auto;">

                <p>Powered by edha 1.0 3B</p>

            </div>

            """)
            
        with gr.Column(elem_classes="canvas-area"):
            map_output = gr.HTML(
                """

                <div style="width:100%; height:100%; display:flex; align-items:center; justify-content:center;">

                    <div class="disclaimer" style="position: absolute; bottom: 16px; left: 50%; transform: translateX(-50%); font-size: 11px; color: #a1a1aa; text-align: center; pointer-events: none; z-index: 1000; width: 100%;">

                        MindMap Studioの回答は正しいとは限らないので、重要な情報は必ず見直してください。

                    </div>

                </div>

                """
            )

    submit_btn.click(
        fn=generate_mindmap, 
        inputs=text_input, 
        outputs=map_output,
        api_name="generate"
    )

demo.launch()