API 使用次數提示

app.py

@@ -1,114 +1,124 @@
-import os, glob, time, requests
-from langchain_community.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.vectorstores import FAISS
+# app.py
+# -------------------------------
+# 套件
+# -------------------------------
+import os, glob, requests
 from langchain.docstore.document import Document
-from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.vectorstores import FAISS
 from langchain.chains import RetrievalQA
-from langchain_huggingface import HuggingFaceHub
+from langchain.llms import HuggingFaceHub
 from docx import Document as DocxDocument
 import gradio as gr
 
 # -------------------------------
-# 1. Hugging Face API Key
-# -------------------------------
-HF_API_TOKEN = os.environ.get("HF_API_TOKEN")  # 或直接在 Space Secrets 設定 HF_API_TOKEN
-
-# -------------------------------
-# 2. 資料路徑
+# 資料路徑
 # -------------------------------
-txt_folder = "./out_texts"
-db_path = "./faiss_db"
-os.makedirs(db_path, exist_ok=True)
+TXT_FOLDER = "./out_texts"
+DB_PATH = "./faiss_db"
+os.makedirs(DB_PATH, exist_ok=True)
 
 # -------------------------------
-# 3. Embeddings
+# 向量資料庫
 # -------------------------------
-embeddings_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+EMBEDDINGS_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
+from langchain.embeddings import HuggingFaceEmbeddings
+embeddings_model = HuggingFaceEmbeddings(model_name=EMBEDDINGS_MODEL_NAME)
 
-# -------------------------------
-# 4. 載入或建立向量資料庫
-# -------------------------------
-if os.path.exists(os.path.join(db_path, "index.faiss")):
-    print("載入現有向量資料庫...")
-    db = FAISS.load_local(db_path, embeddings_model, allow_dangerous_deserialization=True)
+if os.path.exists(os.path.join(DB_PATH, "index.faiss")):
+    db = FAISS.load_local(DB_PATH, embeddings_model, allow_dangerous_deserialization=True)
 else:
-    print("建立新向量資料庫...")
-    txt_files = glob.glob(f"{txt_folder}/*.txt")
+    txt_files = glob.glob(f"{TXT_FOLDER}/*.txt")
     docs = []
-    for fp in txt_files:
-        with open(fp, "r", encoding="utf-8") as f:
-            docs.append(Document(page_content=f.read(), metadata={"source": os.path.basename(fp)}))
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
-    split_docs = text_splitter.split_documents(docs)
+    for filepath in txt_files:
+        with open(filepath, "r", encoding="utf-8") as f:
+            docs.append(Document(page_content=f.read(), metadata={"source": os.path.basename(filepath)}))
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
+    split_docs = splitter.split_documents(docs)
     db = FAISS.from_documents(split_docs, embeddings_model)
-    db.save_local(db_path)
-    print("向量資料庫建立完成。")
+    db.save_local(DB_PATH)
 
 retriever = db.as_retriever(search_type="similarity", search_kwargs={"k":5})
 
 # -------------------------------
-# 5. 模型選擇
+# Hugging Face Hub LLM
 # -------------------------------
-MODEL_DICT = {
-    "google/flan-t5-base": "text2text-generation",
-    "google/flan-t5-large": "text2text-generation",
-    "google/flan-t5-xl": "text2text-generation"
-}
+HUGGINGFACEHUB_API_TOKEN = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
+
+llm = HuggingFaceHub(
+    repo_id="google/flan-t5-large",
+    model_kwargs={"temperature":0.7, "max_new_tokens":512},
+    huggingfacehub_api_token=HUGGINGFACEHUB_API_TOKEN
+)
+
+qa_chain = RetrievalQA.from_chain_type(
+    llm=llm,
+    retriever=retriever,
+    return_source_documents=True
+)
 
-def load_hf_llm(model_name):
-    return HuggingFaceHub(
-        repo_id=model_name,
-        model_kwargs={"temperature":0.7, "max_new_tokens":512},
-        huggingfacehub_api_token=HF_API_TOKEN
-    )
+# -------------------------------
+# 查剩餘額度
+# -------------------------------
+def get_hf_rate_limit():
+    headers = {"Authorization": f"Bearer {HUGGINGFACEHUB_API_TOKEN}"}
+    try:
+        r = requests.get("https://huggingface.co/api/whoami", headers=headers)
+        r.raise_for_status()
+        data = r.json()
+        # free plan 每小時 300 次
+        used = data.get("rate_limit", {}).get("used", 0)
+        remaining = 300 - used if used is not None else "未知"
+        return f"本小時剩餘 API 次數：約 {remaining}"
+    except:
+        return "無法取得 API 速率資訊"
 
 # -------------------------------
-# 6. RAG 生成文章
+# 文章生成
 # -------------------------------
-def rag_generate_hf(query, model_name, segments=5):
-    llm = load_hf_llm(model_name)
-    qa_chain = RetrievalQA.from_chain_type(llm=llm, retriever=retriever, return_source_documents=True)
-    
-    docx_file = "./generated_article.docx"
+def generate_article_with_rate(query, segments=5):
+    docx_file = "/tmp/generated_article.docx"
     doc = DocxDocument()
     doc.add_heading(query, level=1)
     
     all_text = []
-    prompt = f"請依據下列主題生成段落：{query}\n每段約150-200字。"
-
+    prompt = f"請依據下列主題生成段落：{query}\n\n每段約150-200字。"
+    
     for i in range(int(segments)):
         try:
             result = qa_chain({"query": prompt})
             paragraph = result["result"].strip()
+            if not paragraph:
+                paragraph = "（本段生成失敗，請嘗試減少段落或改用較小模型。）"
         except Exception as e:
-            paragraph = f"(本段生成失敗: {e})"
+            paragraph = f"（本段生成失敗：{e}）"
         all_text.append(paragraph)
         doc.add_paragraph(paragraph)
-        prompt = f"請接續上一段生成下一段：\n{paragraph}\n下一段："
-        time.sleep(0.5)  # 避免 API 速率過快
-
+        prompt = f"請接續上一段生成下一段：\n{paragraph}\n\n下一段："
+    
     doc.save(docx_file)
     full_text = "\n\n".join(all_text)
-
-    # 顯示 Hugging Face API 限額
-    headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
-    usage = requests.get("https://api-inference.huggingface.co/usage", headers=headers).json()
-    quota = usage.get("model_card", "無法取得額度")
-    return full_text + f"\n\n[API 使用額度: {quota}]", docx_file
+    
+    # 取得 API 剩餘次數
+    rate_info = get_hf_rate_limit()
+    
+    return f"{rate_info}\n\n{full_text}", docx_file
 
 # -------------------------------
-# 7. Gradio 介面
+# Gradio 介面
 # -------------------------------
 iface = gr.Interface(
-    fn=rag_generate_hf,
+    fn=generate_article_with_rate,
     inputs=[
-        gr.Textbox(lines=2, placeholder="請輸入文章主題"),
-        gr.Dropdown(list(MODEL_DICT.keys()), value="google/flan-t5-base", label="選擇模型"),
-        gr.Slider(minimum=1, maximum=10, value=5, step=1, label="段落數")
+        gr.Textbox(lines=2, placeholder="請輸入文章主題", label="文章主題"),
+        gr.Slider(minimum=1, maximum=10, step=1, value=5, label="段落數")
+    ],
+    outputs=[
+        gr.Textbox(label="生成文章 + API 剩餘次數"),
+        gr.File(label="下載 DOCX")
     ],
-    outputs=[gr.Textbox(label="生成文章"), gr.File(label="下載 DOCX")],
-    title="佛教經論 RAG 系統 (Hugging Face API)",
-    description="使用 Hugging Face API 大模型生成文章，可選模型與段落數，生成完成可下載 DOCX"
+    title="佛教經論 RAG 系統 (HF API)",
+    description="使用 Hugging Face Hub LLM + FAISS RAG，生成文章並提示 API 剩餘額度。"
 )
 
 iface.launch()