多模型下拉選單

app.py

@@ -52,17 +52,16 @@ retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 5})
 # -------------------------------
 # 4. 定義 REST API 呼叫函數
 # -------------------------------
-INFERENCE_MODEL = "google/flan-t5-xl"
-API_URL = f"https://api-inference.huggingface.co/models/{INFERENCE_MODEL}"
 HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"}
 
-def call_hf_inference(prompt, max_new_tokens=512):
+def call_hf_inference(model_name, prompt, max_new_tokens=512):
+    api_url = f"https://api-inference.huggingface.co/models/{model_name}"
     payload = {
         "inputs": prompt,
         "parameters": {"max_new_tokens": max_new_tokens}
     }
     try:
-        response = requests.post(API_URL, headers=HEADERS, json=payload, timeout=60)
+        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]:
@@ -90,7 +89,7 @@ def get_hf_rate_limit():
 # -------------------------------
 # 6. 生成文章（即時進度）
 # -------------------------------
-def generate_article_progress(query, segments=5):
+def generate_article_progress(query, model_name, segments=5):
     docx_file = "/tmp/generated_article.docx"
     doc = DocxDocument()
     doc.add_heading(query, level=1)
@@ -99,12 +98,11 @@ def generate_article_progress(query, segments=5):
     prompt = f"請依據下列主題生成段落：{query}\n\n每段約150-200字。"
 
     for i in range(int(segments)):
-        paragraph = call_hf_inference(prompt)
+        paragraph = call_hf_inference(model_name, prompt)
         all_text.append(paragraph)
         doc.add_paragraph(paragraph)
         prompt = f"請接續上一段生成下一段：\n{paragraph}\n\n下一段："
 
-        # yield 即時更新 Textbox
         yield "\n\n".join(all_text), None
 
     doc.save(docx_file)
@@ -119,6 +117,16 @@ with gr.Blocks() as demo:
     gr.Markdown("使用 Hugging Face REST API + FAISS RAG，生成文章並提示 API 剩餘額度。")
 
     query_input = gr.Textbox(lines=2, placeholder="請輸入文章主題", label="文章主題")
+    model_dropdown = gr.Dropdown(
+        choices=[
+            "gpt2",
+            "EleutherAI/gpt-neo-2.7B",
+            "EleutherAI/gpt-j-6B",
+            "facebook/bart-large-cnn"
+        ],
+        value="EleutherAI/gpt-neo-2.7B",
+        label="選擇生成模型"
+    )
     segments_input = gr.Slider(minimum=1, maximum=10, step=1, value=5, label="段落數")
     output_text = gr.Textbox(label="生成文章 + API 剩餘次數")
     output_file = gr.File(label="下載 DOCX")
@@ -126,7 +134,7 @@ with gr.Blocks() as demo:
     btn = gr.Button("生成文章")
     btn.click(
         generate_article_progress,
-        inputs=[query_input, segments_input],
+        inputs=[query_input, model_dropdown, segments_input],
         outputs=[output_text, output_file]
     )