Upload app.py with huggingface_hub

app.py

@@ -1,163 +1,176 @@
-from fastapi import FastAPI, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
-from fastapi.responses import FileResponse
-from fastapi.staticfiles import StaticFiles
-import os
-import re
-
-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.")
-
-app = FastAPI()
-
-# CORS configuration to allow the frontend to communicate with this API
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],  # Allows all origins for local development
-    allow_credentials=True,
-    allow_methods=["*"],  # Allows all methods
-    allow_headers=["*"],  # Allows all headers
-)
-
-class TextRequest(BaseModel):
-    text: str
-
-class MapResponse(BaseModel):
-    markdown: str
-
-# ---------------------------------------------------------
-# AI Model Initialization (Phase 5)
-# ---------------------------------------------------------
-
-if USE_LLAMA_CPP:
-    GGUF_PATH = "./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. あなたの推測や外部知識を一切混ぜず、入力文章内の事実のみを構造化してください。"""
-
-@app.post("/generate", response_model=MapResponse)
-def generate_mindmap(request: TextRequest):
-    input_text = request.text.strip()
-    if not input_text:
-        raise HTTPException(status_code=400, detail="Text is empty")
-
-    print(f"APIリクエストを受信しました（文字数: {len(input_text)}文字）")
-    print("AIが推論（Markdown構造）を生成中...")
-
-    # Qwen用のチャットプロンプト構築
-    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, # Greedy Decoding (推測・創造を完全に排除)
-            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,     # Greedy Decoding を有効化
-                repetition_penalty=1.1,
-                pad_token_id=tokenizer.pad_token_id,
-                eos_token_id=tokenizer.eos_token_id
-            )
-        
-        # 4. 入力プロンプト部分を切り落として、新しく生成された部分だけを取り出す
-        input_length = inputs["input_ids"].shape[1]
-        generated_tokens = outputs[0][input_length:]
-        
-        # 5. 数値（トークン）を人間が読めるMarkdownテキストにデコード
-        generated_markdown = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
-    
-    # 【超重要】T5トークナイザーが改行(\n)を消してしまう問題の対策
-    # 見出し記号（#）の前に強制的に改行を挿入して、Markmapが認識できるようにする
-    generated_markdown = re.sub(r'\s*(#+ )', r'\n\1', generated_markdown).strip()
-    
-    # Markmapは必ず「#（大見出し）」から始まる必要があるため、存在しない場合は強制付与
-    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- ')
-
-    print("生成が完了しました！ (コンソール出力時の文字化けエラーを防ぐため内容の表示を省略します)")
-    
-    return MapResponse(markdown=generated_markdown)
-
-# フロントエンドの静的ファイルをマウント
-app.mount("/assets", StaticFiles(directory="frontend"), name="assets")
-
-@app.get("/")
-async def root():
-    # ルートURLにアクセスされたらフロントエンドのindex.htmlを返す
-    return FileResponse("frontend/index.html")
-
-@app.get("/{filename}")
-async def get_frontend_file(filename: str):
-    # app.js や style.css などのファイルを返す
-    file_path = os.path.join("frontend", filename)
-    if os.path.exists(file_path):
-        return FileResponse(file_path)
-    raise HTTPException(status_code=404, detail="File not found")
-
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from fastapi.responses import FileResponse
+from fastapi.staticfiles import StaticFiles
+import os
+import re
+
+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.")
+
+app = FastAPI()
+
+# CORS configuration to allow the frontend to communicate with this API
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Allows all origins for local development
+    allow_credentials=True,
+    allow_methods=["*"],  # Allows all methods
+    allow_headers=["*"],  # Allows all headers
+)
+
+class TextRequest(BaseModel):
+    text: str
+
+class MapResponse(BaseModel):
+    markdown: str
+
+# ---------------------------------------------------------
+# AI Model Initialization (Phase 5)
+# ---------------------------------------------------------
+
+if USE_LLAMA_CPP:
+    import os
+    
+    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=GGUF_FILENAME)
+    
+    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. あなたの推測や外部知識を一切混ぜず、入力文章内の事実のみを構造化してください。"""
+
+@app.post("/generate", response_model=MapResponse)
+def generate_mindmap(request: TextRequest):
+    input_text = request.text.strip()
+    if not input_text:
+        raise HTTPException(status_code=400, detail="Text is empty")
+
+    print(f"APIリクエストを受信しました（文字数: {len(input_text)}文字）")
+    print("AIが推論（Markdown構造）を生成中...")
+
+    # Qwen用のチャットプロンプト構築
+    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, # Greedy Decoding (推測・創造を完全に排除)
+            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,     # Greedy Decoding を有効化
+                repetition_penalty=1.1,
+                pad_token_id=tokenizer.pad_token_id,
+                eos_token_id=tokenizer.eos_token_id
+            )
+        
+        # 4. 入力プロンプト部分を切り落として、新しく生成された部分だけを取り出す
+        input_length = inputs["input_ids"].shape[1]
+        generated_tokens = outputs[0][input_length:]
+        
+        # 5. 数値（トークン）を人間が読めるMarkdownテキストにデコード
+        generated_markdown = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
+    
+    # 【超重要】T5トークナイザーが改行(\n)を消してしまう問題の対策
+    # 見出し記号（#）の前に強制的に改行を挿入して、Markmapが認識できるようにする
+    generated_markdown = re.sub(r'\s*(#+ )', r'\n\1', generated_markdown).strip()
+    
+    # Markmapは必ず「#（大見出し）」から始まる必要があるため、存在しない場合は強制付与
+    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- ')
+
+    print("生成が完了しました！ (コンソール出力時の文字化けエラーを防ぐため内容の表示を省略します)")
+    
+    return MapResponse(markdown=generated_markdown)
+
+# フロントエンドの静的ファイルをマウント
+app.mount("/assets", StaticFiles(directory="frontend"), name="assets")
+
+@app.get("/")
+async def root():
+    # ルートURLにアクセスされたらフロントエンドのindex.htmlを返す
+    return FileResponse("frontend/index.html")
+
+@app.get("/{filename}")
+async def get_frontend_file(filename: str):
+    # app.js や style.css などのファイルを返す
+    file_path = os.path.join("frontend", filename)
+    if os.path.exists(file_path):
+        return FileResponse(file_path)
+    raise HTTPException(status_code=404, detail="File not found")
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)