chat template

app.py

@@ -116,17 +116,11 @@ def build_llm_prompt(system: str, context: list[str], user_question: str) -> str
     conversation.append({"role": "user", "content": user_question.strip()})
 
     # 套用 LLaMA-style prompt 格式
-    prompt = ""
-    for turn in conversation:
-        role = turn["role"]
-        content = turn["content"]
-        if role == "system":
-            prompt += f"<<SYS>>\n{content}\n<</SYS>>\n\n"
-        elif role == "user":
-            prompt += f"[INST] {content} [/INST]\n"
-        elif role == "assistant":
-            prompt += f"{content}\n"
-    return prompt
+    return tokenizer.apply_chat_template(
+        conversation,
+        tokenize=False,
+        add_generation_prompt=False
+    )
 
 # ---------- 4. Gradio playground (same UI as before) --------------------------
 def store_doc(doc_text: str,user_id="demo",chunk_size=DEFAULT_CHUNK_SIZE,chunk_overlap=DEFAULT_CHUNK_OVERLAP):
@@ -198,7 +192,6 @@ def answer(system: str, context: str, question: str,
         full   = tokenizer.decode(output[0], skip_special_tokens=True)
         reply  = full.split("<|im_start|>assistant")[-1].strip()
         return reply
-
           
     except Exception as e:
         tb = traceback.format_exc()
@@ -215,7 +208,6 @@ def clear_kb(user_id="demo"):
         return f"User ID '{user_id}' not found."
 
 # ---- UI layout (feel free to tweak cosmetics) -----------------------------
-# ---- UI layout -----------------------------------------------------
 with gr.Blocks() as demo:
     gr.Markdown(
         "### Tiny-RAG playground – 1) paste a passage → **Store** · "
@@ -278,7 +270,6 @@ with gr.Blocks() as demo:
         outputs=answer_box
     )
 
-
 # ---------- 3. FastAPI layer --------------------------------------------------
 class IngestReq(BaseModel):
     user_id:str