""" app.py – Tiny-RAG (Gradio playground) + REST API (/ingest, /query) """ # ---------- 1. imports & global helpers ------------- import os os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" import 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 = "meta-llama/Llama-3.1-8B-Instruct" EMB_MODEL_ID = "mixedbread-ai/mxbai-embed-large-v1" MAX_PROMPT_TOKENS = 8192 # ---------- new defaults & helper ------------------ DEFAULT_TEMP = 0.7 DEFAULT_TOP_P = 0.9 DEFAULT_TOP_K_TOK = 40 # token-level sampling DEFAULT_CHUNK_SIZE = 512 # characters DEFAULT_CHUNK_OVERLAP = 128 def chunk_text(text: str, size: int, overlap: int): """Yield sliding-window chunks of *text* with character overlap.""" for start in range(0, len(text), size - overlap): yield text[start : start + size] # --- 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: load_embedder() with torch.no_grad(): inputs = emb_tokenizer(text, return_tensors="pt", truncation=True).to(emb_model.device) vec = emb_model(**inputs).last_hidden_state[:, 0] return F.normalize(vec, dim=-1).cpu() # ---------- 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],chunk_size: int = DEFAULT_CHUNK_SIZE,chunk_overlap: int = DEFAULT_CHUNK_OVERLAP) -> int: # ---------- NEW ---------- chunks = [] for d in docs: chunks.extend(chunk_text(d, chunk_size, chunk_overlap)) docs = [c for c in chunks if c.strip()] load_embedder() # lazy-load once new_vecs = torch.stack([embed(t) for t in docs]).cpu() 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) # ----- Qwen-chat prompt helper --------------------------------------------- def build_llm_prompt(system: str, context: list[str], user_question: str) -> str: """ 建立適用於 LLaMA/Qwen 等模型的 prompt，支援多段 context， 並強化 system prompt 限制模型僅輸出回應內容。 """ load_chat() # 確保 tokenizer 載入 # 強化指令：防止解釋與步驟 system_prompt = ( f"{system.strip()}\n" "Do not include any explanations, steps, or analysis. " "Only output the final reply content." ) conversation = [ {"role": "system", "content": system_prompt} ] # 多段 context 當作 user 發言 for ctx in context: if ctx.strip(): # 忽略空內容 conversation.append({"role": "user", "content": ctx.strip()}) # 最後加入使用者問題 conversation.append({"role": "user", "content": user_question.strip()}) # 套用 LLaMA-style 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): try: n = add_docs(user_id, [doc_text], chunk_size, chunk_overlap) if n == 0: return "Nothing stored (empty input)." return f"Stored — KB now has {len(kb[user_id]['texts'])} passage(s)." except Exception as e: return f"Error during storing: {e}" import traceback def answer(system: str, context: str, question: str, user_id="demo", history="None", temperature=DEFAULT_TEMP, top_p=DEFAULT_TOP_P, top_k_tok=DEFAULT_TOP_K_TOK): """UI callback: retrieve, build prompt with Qwen tags, generate answer.""" try: if not question.strip(): return "Please ask a question." if history != "None" and not kb[user_id]["texts"]: return "No reference passage yet. Add one first." context_list = [context] # 1. Retrieve top-k similar passages if history == "Some": q_vec = embed(question).view(-1).cpu() store = kb[user_id] vecs = store["vecs"] if vecs is None or vecs.size(0) == 0: return "Knowledge base is empty or corrupted." sims = torch.matmul(vecs, q_vec) # [N] if sims.dim() > 1: sims = sims.squeeze(1) k = min(4, sims.size(0)) idxs = torch.topk(sims, k=k, dim=0).indices.tolist() context_list += [store["texts"][i] for i in idxs] elif history == "All": store = kb[user_id] context_list += store["texts"] # 2. Build a Qwen-chat prompt (helper defined earlier) prompt = build_llm_prompt(system, context_list, question) # 3. Tokenise & cap load_chat() tokens = tokenizer( prompt, return_tensors="pt", add_special_tokens=False, # we built the chat template ourselves ) if tokens["input_ids"].size(1) > MAX_PROMPT_TOKENS: tokens = {k: v[:, -MAX_PROMPT_TOKENS:] for k, v in tokens.items()} tokens = {k: v.to(chat_model.device) for k, v in tokens.items()} # --- generate ------------------------------------------------------ output = chat_model.generate( **tokens, max_new_tokens=512, max_length=MAX_PROMPT_TOKENS + 512, do_sample=True, temperature=temperature, top_p=top_p, top_k=top_k_tok ) full = tokenizer.decode(output[0], skip_special_tokens=False) start = "<|start_header_id|>assistant<|end_header_id|>\n\n" startwithoutend = "<|start_header_id|>assistant" end = "<|eot_id|>" if start in full: reply = full.split(start)[-1].split(end)[0].strip() elif startwithoutend in full: reply = full.split(startwithoutend)[-1].split(end)[0].strip() else: reply = full return reply except Exception as e: tb = traceback.format_exc() return f"Error in app.py: {tb}, k={k}, sims.numel()={sims.numel()}, sims.shape={sims.shape if 'q_vec' in locals() else 'N/A'}" finally: torch.cuda.empty_cache() def clear_kb(user_id="demo"): if user_id in kb: kb[user_id]["texts"].clear() kb[user_id]["vecs"] = None return f"Cleared KB for user '{user_id}'." else: return f"User ID '{user_id}' not found." # ---- UI layout (feel free to tweak cosmetics) ----------------------------- with gr.Blocks() as demo: gr.Markdown("### Tiny-RAG playground …") # ---- passage ingestion ---- with gr.Row(): passage_box = gr.Textbox(lines=6, label="Reference passage") user_id_box = gr.Textbox(value="demo", label="User ID") chunk_box = gr.Slider(128, 2048, value=DEFAULT_CHUNK_SIZE, step=64, label="Chunk size (chars)") overlap_box = gr.Slider(0, 1024, value=DEFAULT_CHUNK_OVERLAP, step=32, label="Chunk overlap") store_btn = gr.Button("Store passage") clear_btn = gr.Button("Clear KB") status_box = gr.Markdown() # declare *before* wiring handlers # ---- wire handlers (each button exactly once) ---- store_btn.click( fn=store_doc, inputs=[passage_box, user_id_box, chunk_box, overlap_box], outputs=status_box ) clear_btn.click( fn=clear_kb, inputs=user_id_box, outputs=status_box ) # ---------- Q & A ---------- question_box = gr.Textbox(lines=2, label="Ask a question") history_cb = gr.Textbox(value="None", label="Use chat history") system_box = gr.Textbox(lines=2, label="System prompt") context_box = gr.Textbox(lines=6, label="Context passages") # NEW sampling sliders temp_box = gr.Slider(0.0, 1.5, value=DEFAULT_TEMP, step=0.05, label="Temperature") topp_box = gr.Slider(0.0, 1.0, value=DEFAULT_TOP_P, step=0.01, label="Top-p") topk_box = gr.Slider(1, 100, value=DEFAULT_TOP_K_TOK, step=1, label="Top-k (tokens)") answer_btn = gr.Button("Answer") answer_box = gr.Textbox(lines=6, label="Assistant reply") answer_btn.click( fn=answer, inputs=[system_box, context_box, question_box, user_id_box, history_cb, temp_box, topp_box, topk_box], outputs=answer_box ) # ---------- 3. FastAPI layer -------------------------------------------------- class IngestReq(BaseModel): user_id:str docs:list[str] class QueryReq(BaseModel): user_id:str question:str api = FastAPI() api = gr.mount_gradio_app(api, demo, path="/") # ---------- 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)