使用 Hugging Face API 大模型生成文章

app.py

@@ -1,135 +1,114 @@
-# -------------------------------
-# 1. 匯入套件
-# -------------------------------
-import os, glob, time
+import os, glob, time, requests
 from langchain_community.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
-from langchain_core.documents import Document
-from langchain_community.chat_models import ChatHuggingFaceHub
+from langchain.docstore.document import Document
+from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.chains import RetrievalQA
-
+from langchain_huggingface 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. 設定路徑
+# 2. 資料路徑
 # -------------------------------
-txt_folder = "out_texts"     # 放你的 .txt 檔
-db_path = "faiss_db"
+txt_folder = "./out_texts"
+db_path = "./faiss_db"
 os.makedirs(db_path, exist_ok=True)
 
 # -------------------------------
-# 3. 建立 embeddings
+# 3. Embeddings
 # -------------------------------
 embeddings_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
 
 # -------------------------------
-# 4. 建立或載入向量資料庫
+# 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)
 else:
-    print("沒有資料庫，開始建立新向量資料庫...")
+    print("建立新向量資料庫...")
     txt_files = glob.glob(f"{txt_folder}/*.txt")
     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)}))
+    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)
-    print("產生向量嵌入中...")
     db = FAISS.from_documents(split_docs, embeddings_model)
     db.save_local(db_path)
     print("向量資料庫建立完成。")
 
+retriever = db.as_retriever(search_type="similarity", search_kwargs={"k":5})
+
 # -------------------------------
-# 5. Hugging Face 模型設定
+# 5. 模型選擇
 # -------------------------------
-HUGGINGFACE_API_TOKEN = os.getenv("HF_TOKEN")  # 建議在 Spaces Secrets 設定
-
 MODEL_DICT = {
-    "google/flan-t5-large": 512,
-    "tiiuae/falcon-7b-instruct": 512
+    "google/flan-t5-base": "text2text-generation",
+    "google/flan-t5-large": "text2text-generation",
+    "google/flan-t5-xl": "text2text-generation"
 }
 
-MAX_HOURLY_REQUESTS = 50
-request_count = 0
-last_reset_time = time.time()
-
-# -------------------------------
-# 6. RAG 主函式
-# -------------------------------
-def rag_generate_hfapi(query, model_name, segments=5, max_words=1500):
-    global request_count, last_reset_time
-    if time.time() - last_reset_time > 3600:
-        request_count = 0
-        last_reset_time = time.time()
-
-    if request_count >= MAX_HOURLY_REQUESTS:
-        return f"本小時生成次數已達上限 ({MAX_HOURLY_REQUESTS})，請稍後再試。", None
-
-    llm = ChatHuggingFaceHub(
+def load_hf_llm(model_name):
+    return HuggingFaceHub(
         repo_id=model_name,
-        model_kwargs={"temperature": 0.7, "max_new_tokens": MODEL_DICT[model_name]},
-        huggingfacehub_api_token=HUGGINGFACE_API_TOKEN
-    )
-
-    qa_chain = RetrievalQA.from_chain_type(
-        llm=llm,
-        retriever=db.as_retriever(search_type="similarity", search_kwargs={"k": 5}),
-        return_source_documents=True
+        model_kwargs={"temperature":0.7, "max_new_tokens":512},
+        huggingfacehub_api_token=HF_API_TOKEN
     )
 
-    prompt = f"""請依據下列主題生成一篇文章：
-主題：{query}
-需求：
-- 總共 {segments} 段
-- 每段約 {max_words // segments} 字
-- 總字數請控制在 {max_words} 字以內
-- 請自動分段輸出
-"""
-
-    try:
-        result = qa_chain({"query": prompt})
-        full_text = result["result"].strip()
-        if not full_text:
-            full_text = "(生成失敗，請改用其他模型或調整段落數)"
-    except Exception as e:
-        return f"(生成失敗：{str(e)})", None
-
-    request_count += 1
-
-    paragraphs = [p.strip() for p in full_text.split("\n") if p.strip()]
-
-    docx_file = "generated_article.docx"
+# -------------------------------
+# 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"
     doc = DocxDocument()
     doc.add_heading(query, level=1)
-    for p in paragraphs:
-        doc.add_paragraph(p)
+    
+    all_text = []
+    prompt = f"請依據下列主題生成段落：{query}\n每段約150-200字。"
+
+    for i in range(int(segments)):
+        try:
+            result = qa_chain({"query": prompt})
+            paragraph = result["result"].strip()
+        except Exception as e:
+            paragraph = f"(本段生成失敗: {e})"
+        all_text.append(paragraph)
+        doc.add_paragraph(paragraph)
+        prompt = f"請接續上一段生成下一段：\n{paragraph}\n下一段："
+        time.sleep(0.5)  # 避免 API 速率過快
+
     doc.save(docx_file)
+    full_text = "\n\n".join(all_text)
 
-    return "\n\n".join(paragraphs), docx_file
+    # 顯示 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
 
 # -------------------------------
 # 7. Gradio 介面
 # -------------------------------
 iface = gr.Interface(
-    fn=rag_generate_hfapi,
+    fn=rag_generate_hf,
     inputs=[
         gr.Textbox(lines=2, placeholder="請輸入文章主題"),
-        gr.Dropdown(list(MODEL_DICT.keys()), value="google/flan-t5-large", label="選擇模型"),
-        gr.Slider(minimum=1, maximum=10, value=5, step=1, label="段落數"),
-        gr.Slider(minimum=500, maximum=3000, value=1500, step=100, label="文章字數上限")
-    ],
-    outputs=[
-        gr.Textbox(label="生成文章"),
-        gr.File(label="下載 DOCX")
+        gr.Dropdown(list(MODEL_DICT.keys()), value="google/flan-t5-base", label="選擇模型"),
+        gr.Slider(minimum=1, maximum=10, value=5, step=1, label="段落數")
     ],
+    outputs=[gr.Textbox(label="生成文章"), gr.File(label="下載 DOCX")],
     title="佛教經論 RAG 系統 (Hugging Face API)",
-    description="使用 Hugging Face API 生成文章，可選大模型，分段生成並下載 DOCX，每小時生成次數有限制"
+    description="使用 Hugging Face API 大模型生成文章，可選模型與段落數，生成完成可下載 DOCX"
 )
 
 iface.launch()