GPT生的鬼東西 試記憶+RAG 大概會炸 locally有備份

app.py

@@ -1,31 +1,189 @@
-import os
-import torch
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-# 全局變數初始化為 None
-tokenizer = None
-model = None
-
-def reply(prompt, model_id="QWen/Qwen1.5-7B-Chat", api_token=None):
-    try:
-        if api_token is None:
-            api_token = os.getenv("HF_token")
-    except Exception as e:
-        return f"無法取得 API token。\n錯誤訊息：{str(e)}"
-    
-    global tokenizer, model
-    try:
-        if tokenizer is None or model is None:
-            tokenizer = AutoTokenizer.from_pretrained(model_id, token=api_token)
-            model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.bfloat16, token=api_token)
-
-        inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
-        outputs = model.generate(**inputs, max_new_tokens=8192)
-        return tokenizer.decode(outputs[0], skip_special_tokens=True)
-    except RuntimeError as e:
-        return f"執行時錯誤：{str(e)}。"
-    except Exception as e:
-        return f"發生錯誤：{str(e)}"
-
-gr.Interface(fn=reply, inputs="text", outputs="text").launch()
+"""
+app.py  –  Tiny-RAG (Gradio playground)  +  REST API (/ingest, /query)
+"""
+
+# ---------- 1. imports & global helpers -------------
+import os, math, torch, uvicorn, gradio as gr
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from transformers import (
+    AutoTokenizer, AutoModelForCausalLM,
+    AutoTokenizer, AutoModel, AutoConfig
+)
+import torch.nn.functional as F
+from collections import defaultdict
+HF_TOKEN = os.getenv("HF_token")
+CHAT_MODEL_ID = "QWen/Qwen1.5-7B-Chat"
+EMB_MODEL_ID  = "mixedbread-ai/mxbai-embed-large-v1"
+
+# --- lazy loaders (unchanged) -------------------------------------------------
+tokenizer, chat_model = None, None
+emb_tokenizer, emb_model = None, None
+
+def load_chat():
+    global tokenizer, chat_model
+    if tokenizer is None:
+        tokenizer = AutoTokenizer.from_pretrained(CHAT_MODEL_ID, token=HF_TOKEN)
+        chat_model = AutoModelForCausalLM.from_pretrained(
+            CHAT_MODEL_ID, device_map="auto", torch_dtype=torch.bfloat16, token=HF_TOKEN
+        )
+
+def load_embedder():
+    global emb_tokenizer, emb_model
+    if emb_tokenizer is None:
+        emb_tokenizer = AutoTokenizer.from_pretrained(EMB_MODEL_ID, token=HF_TOKEN)
+        cfg = AutoConfig.from_pretrained(EMB_MODEL_ID, token=HF_TOKEN)
+        emb_model = AutoModel.from_pretrained(
+            EMB_MODEL_ID, device_map="auto", torch_dtype=torch.float16, config=cfg, token=HF_TOKEN
+        )
+        emb_model.eval()
+
+@torch.no_grad()
+def embed(text:str)->torch.Tensor:
+    """Return L2-normalised embedding vector."""
+    load_embedder()
+    inputs = emb_tokenizer(text, return_tensors="pt", truncation=True).to(emb_model.device)
+    vec = emb_model(**inputs).last_hidden_state[:, 0]  # CLS pooling
+    return F.normalize(vec, dim=-1).squeeze(0)
+
+# ---------- 2. tiny in-memory KB shared by Gradio & API ----------------------
+# ---------- 2. Tiny in-memory knowledge-base -------------------------------
+# One dict entry per user_id.
+# Each entry holds:
+#   • "texts": list[str]   – the raw passages we ingested
+#   • "vecs" : Tensor[N,d] – their embeddings stacked row-wise
+# --------------------------------------------------------------------------
+
+
+kb = defaultdict(lambda: {"texts": [], "vecs": None})
+
+def add_docs(user_id: str, docs: list[str]) -> int:
+    """Embed *docs* and append them to the KB for *user_id*.
+       Returns the number of docs actually stored."""
+    docs = [t for t in docs if t.strip()]          # skip blanks
+    if not docs:
+        return 0
+
+    load_embedder()                                # lazy-load once
+    new_vecs = torch.stack([embed(t) for t in docs])
+    store = kb[user_id]                            # auto-creates via defaultdict
+    store["texts"].extend(docs)
+    store["vecs"] = (
+        new_vecs if store["vecs"] is None
+        else torch.cat([store["vecs"], new_vecs])
+    )
+    return len(docs)
+
+# ---------- 3. FastAPI layer --------------------------------------------------
+class IngestReq(BaseModel):
+    user_id:str
+    docs:list[str]
+
+class QueryReq(BaseModel):
+    user_id:str
+    question:str
+
+api = FastAPI()
+
+@api.post("/ingest")
+def ingest(req:IngestReq):
+    load_embedder()
+    vecs = torch.stack([embed(t) for t in req.docs])
+    store = kb.setdefault(req.user_id, {"texts":[], "vecs":None})
+    store["texts"].extend(req.docs)
+    store["vecs"] = vecs if store["vecs"] is None else torch.cat([store["vecs"], vecs])
+    return {"added": len(req.docs)}
+
+@api.post("/query")
+def rag(req:QueryReq):
+    store = kb.get(req.user_id)
+    if not store:
+        raise HTTPException(404, "No knowledge ingested for this user.")
+    q_vec = embed(req.question)
+    sims  = torch.matmul(store["vecs"], q_vec)
+    topk  = torch.topk(sims, k=min(4, sims.size(0))).indices
+    context = "\n".join(store["texts"][i] for i in topk.tolist())
+
+    prompt = f"""You are an email assistant.
+Use the context to answer.
+Context:
+{context}
+
+User question: {req.question}
+Assistant:"""
+
+    load_chat()
+    inputs = tokenizer(prompt, return_tensors="pt").to(chat_model.device)
+    out = chat_model.generate(**inputs, max_new_tokens=512)
+    ans = tokenizer.decode(out[0], skip_special_tokens=True).split("Assistant:",1)[-1].strip()
+    return {"answer": ans}
+
+# ---------- 4. Gradio playground (same UI as before) --------------------------
+# ---------- 4. Gradio playground ------------------------------------------
+def store_doc(doc_text: str, user_id="demo"):
+    """UI callback: take the textbox content and shove it into the KB."""
+    n = add_docs(user_id, [doc_text])
+    if n == 0:
+        return "⚠️ Nothing stored (empty input)."
+    return f"📚 Stored ✅ — KB now has {len(kb[user_id]['texts'])} passage(s)."
+
+def answer(question: str, user_id="demo"):
+    """UI callback: retrieve, build prompt, generate answer."""
+    if not question.strip():
+        return "⚠️ Please ask a question."
+    if not kb[user_id]["texts"]:
+        return "⚠️ No reference passage yet. Add one first."
+
+    # 1️⃣  Retrieve top-k similar chunks (k ≤ #chunks)
+    q_vec  = embed(question)
+    store  = kb[user_id]
+    sims   = torch.matmul(store["vecs"], q_vec)                  # [N]
+    k      = min(4, sims.numel())
+    idxs   = torch.topk(sims, k=k).indices.tolist()
+    context = "\n".join(store["texts"][i] for i in idxs)
+
+    # 2️⃣  Build prompt
+    prompt = f"""You are an email assistant.
+Use ONLY the context below to answer.
+Context:
+{context}
+
+Question: {question}
+Answer:"""
+
+    # 3️⃣  Generate
+    load_chat()
+    inputs  = tokenizer(prompt, return_tensors="pt").to(chat_model.device)
+    output  = chat_model.generate(**inputs, max_new_tokens=512)
+    reply   = tokenizer.decode(output[0], skip_special_tokens=True)
+    return reply.split("Answer:", 1)[-1].strip()
+
+# ---- UI layout (feel free to tweak cosmetics) -----------------------------
+with gr.Blocks() as demo:
+    gr.Markdown("### 📥 Tiny-RAG playground  –  1) paste a passage → store   2) ask a question")
+
+    # ---- passage ingestion ----
+    with gr.Row():
+        passage_box = gr.Textbox(lines=6, label="Reference passage")
+        store_btn   = gr.Button("➕ Store passage")
+    status_box = gr.Markdown()
+    store_btn.click(fn=store_doc,
+                    inputs=passage_box,
+                    outputs=status_box)
+
+    # ---- Q & A ----
+    question_box = gr.Textbox(lines=2, label="Ask a question")
+    answer_btn   = gr.Button("🤖 Answer")
+    answer_box   = gr.Textbox(lines=6, label="Assistant reply")
+
+    answer_btn.click(fn=answer,
+                     inputs=question_box,
+                     outputs=answer_box)
+
+
+# ---------- 5. run both (FastAPI + Gradio) -----------------------------------
+if __name__ == "__main__":
+    # launch Gradio on a background thread
+    demo.queue().launch(share=False, prevent_thread_lock=True)
+    # then start FastAPI (uvicorn blocks main thread)
+    uvicorn.run(api, host="0.0.0.0", port=8000)