Update app.py

app.py

@@ -5,11 +5,10 @@
 import os, glob, requests
 from langchain.docstore.document import Document
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.chains import RetrievalQA
-from langchain_huggingface import HuggingFaceEmbeddings, HuggingFaceEndpoint
+from langchain_community.vectorstores import FAISS
+from langchain_huggingface import HuggingFaceEmbeddings
 from docx import Document as DocxDocument
 import gradio as gr
-from langchain_community.vectorstores import FAISS
 
 # -------------------------------
 # 2. 環境變數與資料路徑
@@ -17,7 +16,7 @@ from langchain_community.vectorstores import FAISS
 TXT_FOLDER = "./out_texts"
 DB_PATH = "./faiss_db"
 os.makedirs(DB_PATH, exist_ok=True)
-os.makedirs(TXT_FOLDER, exist_ok=True)  # 避免沒有 txt 檔時錯誤
+os.makedirs(TXT_FOLDER, exist_ok=True)
 
 HF_TOKEN = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
 if not HF_TOKEN:
@@ -51,53 +50,47 @@ else:
 retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 5})
 
 # -------------------------------
-# 4. LLM 設定（Hugging Face Endpoint）
+# 4. 定義 REST API 呼叫函數
 # -------------------------------
-llm = HuggingFaceEndpoint(
-    repo_id="google/flan-t5-large",
-    task="text2text-generation",
-    huggingfacehub_api_token=HF_TOKEN,
-    temperature=0.7,
-    max_new_tokens=512,
-)
+INFERENCE_MODEL = "google/flan-t5-large"
+API_URL = f"https://api-inference.huggingface.co/models/{INFERENCE_MODEL}"
+HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"}
 
-qa_chain = RetrievalQA.from_chain_type(
-    llm=llm,
-    retriever=retriever,
-    return_source_documents=True
-)
-
-# -------------------------------
-# 5. 檢查 Hugging Face Token 權限
-# -------------------------------
-def check_hf_token_permissions():
-    """確認 Token 是否可呼叫 Inference Endpoint"""
-    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
+def call_hf_inference(prompt, max_new_tokens=512):
+    payload = {
+        "inputs": prompt,
+        "parameters": {"max_new_tokens": max_new_tokens}
+    }
+    try:
+        response = requests.post(API_URL, headers=HEADERS, json=payload, timeout=60)
+        response.raise_for_status()
+        data = response.json()
+        if isinstance(data, list) and "generated_text" in data[0]:
+            return data[0]["generated_text"]
+        elif isinstance(data, dict) and "error" in data:
+            return f"（生成失敗：{data['error']}）"
+        else:
+            return str(data)
+    except Exception as e:
+        return f"（生成失敗：{e}）"
+
+# -------------------------------
+# 5. 查詢 API 剩餘額度
+# -------------------------------
+def get_hf_rate_limit():
     try:
-        r = requests.get("https://huggingface.co/api/whoami-v2", headers=headers)
+        r = requests.get("https://huggingface.co/api/whoami", headers=HEADERS)
         r.raise_for_status()
         data = r.json()
-        if "allow_inference" in data and data["allow_inference"]:
-            return True
-        return False
+        remaining = data.get("rate_limit", {}).get("remaining", "未知")
+        return f"本小時剩餘 API 次數：約 {remaining}"
     except Exception:
-        return False
-
-token_valid = check_hf_token_permissions()
-if not token_valid:
-    print("⚠ 警告：Hugging Face API Token 權限不足，無法呼叫模型。")
-
+        return "無法取得 API 速率資訊"
 
 # -------------------------------
-# 6. 生成文章（修正版，支援進度顯示）
+# 6. 生成文章（即時進度）
 # -------------------------------
-def generate_article_with_progress(query, segments=5):
-    if not token_valid:
-        # Token 權限不足，直接返回訊息
-        yield "⚠ API Token 權限不足，請檢查 Token 是否允許呼叫 Inference Endpoint。", None
-        return
-
-    import time
+def generate_article_progress(query, segments=5):
     docx_file = "/tmp/generated_article.docx"
     doc = DocxDocument()
     doc.add_heading(query, level=1)
@@ -106,33 +99,24 @@ def generate_article_with_progress(query, segments=5):
     prompt = f"請依據下列主題生成段落：{query}\n\n每段約150-200字。"
 
     for i in range(int(segments)):
-        try:
-            result = qa_chain({"query": prompt})
-            paragraph = result.get("result", "").strip()
-            if not paragraph:
-                paragraph = "（本段生成失敗，請嘗試減少段落或改用較小模型。）"
-        except Exception as e:
-            paragraph = f"（本段生成失敗：{e}）"
-
+        paragraph = call_hf_inference(prompt)
         all_text.append(paragraph)
         doc.add_paragraph(paragraph)
         prompt = f"請接續上一段生成下一段：\n{paragraph}\n\n下一段："
 
+        # yield 即時更新 Textbox
         yield "\n\n".join(all_text), None
-        time.sleep(0.1)
 
     doc.save(docx_file)
     rate_info = get_hf_rate_limit()
-    final_text = f"{rate_info}\n\n" + "\n\n".join(all_text)
-    yield final_text, docx_file
-
+    yield f"{rate_info}\n\n" + "\n\n".join(all_text), docx_file
 
 # -------------------------------
-# 7. Gradio 介面（修正版）
+# 7. Gradio 介面
 # -------------------------------
 with gr.Blocks() as demo:
     gr.Markdown("# 佛教經論 RAG 系統 (HF API)")
-    gr.Markdown("使用 Hugging Face Endpoint LLM + FAISS RAG，生成文章並提示 API 剩餘額度。")
+    gr.Markdown("使用 Hugging Face REST API + FAISS RAG，生成文章並提示 API 剩餘額度。")
 
     query_input = gr.Textbox(lines=2, placeholder="請輸入文章主題", label="文章主題")
     segments_input = gr.Slider(minimum=1, maximum=10, step=1, value=5, label="段落數")
@@ -140,17 +124,14 @@ with gr.Blocks() as demo:
     output_file = gr.File(label="下載 DOCX")
 
     btn = gr.Button("生成文章")
-
-    # 使用 .click() 搭配 generator
     btn.click(
-        fn=generate_article_with_progress,
+        generate_article_progress,
         inputs=[query_input, segments_input],
         outputs=[output_text, output_file]
     )
 
 # -------------------------------
-# 8. 啟動 Gradio（HF Space 適用）
+# 8. 啟動 Gradio（Hugging Face Space 適用）
 # -------------------------------
 if __name__ == "__main__":
     demo.launch()
-