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 "