Update app.py

app.py

@@ -1,7 +1,8 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 import torch
 from urllib.parse import quote
+from threading import Thread
 
 # --- 模型配置 (CPU TinyLlama) ---
 MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
@@ -10,25 +11,7 @@ model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 print("模型與分詞器載入完成。")
 
-# --- 核心功能函數 ---
-
-def llm_generate(prompt, max_new_tokens=256, temperature=0.7, top_p=0.9):
-    """基礎的語言模型文字生成函數"""
-    chat = [{"role": "user", "content": prompt}]
-    formatted_prompt = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
-    inputs = tokenizer(formatted_prompt, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=int(max_new_tokens),
-            do_sample=True,
-            temperature=float(temperature),
-            top_p=float(top_p),
-            eos_token_id=tokenizer.eos_token_id
-        )
-    response = tokenizer.decode(outputs[0, inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
-    return response.strip()
-
+# --- 地圖連結函數 (保持不變) ---
 def get_map_link(location_query):
     """根據地點名稱生成 Google 地圖的 URL"""
     if not location_query:
@@ -37,45 +20,84 @@ def get_map_link(location_query):
     map_url = base_url + quote(location_query)
     return f"點擊這裡查看 **{location_query}** 的地圖：\n[在 Google 地圖中開啟]({map_url})"
 
+# --- 核心處理函數 (修改為支援 Streaming) ---
 def unified_processor(query):
     """
-    統一處理函數：接收使用者輸入，同時生成地點描述和地圖連結。
+    統一處理函數：
+    1. 以串流方式生成地點描述。
+    2. 生成完畢後，一次性生成推薦內容。
+    3. 最後顯示地圖。
     """
     if not query:
-        return "請輸入一個地點或一段描述。", ""
+        yield "請輸入一個地點或一段描述。", "", "狀態：待機中"
+        return
 
-    print(f"正在處理查詢：'{query}'")
+    # --- 階段一：串流生成地點描述 ---
+    yield "", "", f"狀態：正在為「{query}」生成介紹..."
 
-    # 步驟 1: 讓模型生成關於這個地點的描述
-    description_prompt = f"請用繁體中文，生動地介紹一下「{query}」這個地方的特色、歷史或是有趣的景點。"
-    generated_description = llm_generate(description_prompt)
+    # 設置 streamer
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
 
-    # 步驟 2: 產生該地點的地圖連結
-    map_link_markdown = get_map_link(query)
+    # 準備生成參數
+    prompt = f"請用繁體中文，生動地介紹一下「{query}」這個地方的特色、歷史或是有趣的景點。"
+    chat = [{"role": "user", "content": prompt}]
+    formatted_prompt = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(formatted_prompt, return_tensors="pt")
+
+    generation_kwargs = dict(
+        inputs,
+        streamer=streamer,
+        max_new_tokens=256,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9
+    )
+
+    # 使用多執行緒來運行 blocking 的 generate 方法
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    # 即時從 streamer 讀取並更新輸出
+    description_output = "### 關於「" + query + "」\n"
+    for new_text in streamer:
+        description_output += new_text
+        yield description_output, "", f"狀態：正在為「{query}」生成介紹..."
+
+    # --- 階段二：生成推薦內容 ---
+    yield description_output, "", f"狀態：正在為「{query}」生成附近推薦..."
 
-    # 步驟 3: 讓模型推薦附近的景點
     recommendation_prompt = f"我在「{query}」這個地方，請用條列的方式，推薦3個附近的必去景點或必吃美食。"
-    generated_recommendations = llm_generate(recommendation_prompt)
+    chat_reco = [{"role": "user", "content": recommendation_prompt}]
+    formatted_reco_prompt = tokenizer.apply_chat_template(chat_reco, tokenize=False, add_generation_prompt=True)
+    inputs_reco = tokenizer(formatted_reco_prompt, return_tensors="pt")
+    
+    with torch.no_grad():
+        outputs_reco = model.generate(**inputs_reco, max_new_tokens=150)
+    recommendations = tokenizer.decode(outputs_reco[0, inputs_reco["input_ids"].shape[-1]:], skip_special_tokens=True)
 
-    # 組合最終的文字輸出
     final_text_output = (
-        f"### 關於「{query}」\n"
-        f"{generated_description}\n\n"
+        f"{description_output}\n\n"
         f"### 附近推薦\n"
-        f"{generated_recommendations}"
+        f"{recommendations.strip()}"
     )
 
-    return final_text_output, map_link_markdown
+    # --- 階段三：生成地圖 ---
+    map_link_markdown = get_map_link(query)
+    
+    yield final_text_output, map_link_markdown, "狀態：導覽完成！"
 
-# --- Gradio 介面 (整合版) ---
+
+# --- Gradio 介面 ---
 with gr.Blocks(theme=gr.themes.Default()) as demo:
     gr.Markdown(
         """
         # 🗺️ AI 智慧導遊 ✨
-        輸入一個地點，AI 將為您生成生動的介紹、推薦附近景點，並附上地圖！
+        輸入一個地點，AI 將為您即時生成生動的介紹、推薦附近景點，並附上地圖！
         """
     )
 
+    status_display = gr.Markdown("狀態：待機中")
+
     with gr.Row():
         query_input = gr.Textbox(
             label="請輸入地點名稱或描述",
@@ -90,10 +112,11 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
         with gr.Column(scale=1):
             map_output = gr.Markdown(label="地圖連結")
 
+    # Gradio 會自動處理 yield 函數的串流輸出
     process_button.click(
         fn=unified_processor,
         inputs=query_input,
-        outputs=[text_output, map_output]
+        outputs=[text_output, map_output, status_display]
     )
 
     gr.Examples(