Update app.py

app.py

@@ -1,46 +1,130 @@
+# app.py — Ivy Chatbot (Free RAG Demo on CPU)
 import gradio as gr
-from transformers import BlenderbotSmallTokenizer, BlenderbotSmallForConditionalGeneration
+import numpy as np
+from transformers import (
+    BlenderbotSmallTokenizer,
+    BlenderbotSmallForConditionalGeneration
+)
+from sentence_transformers import SentenceTransformer
+
+# ---------------------------
+# 1) 模型与向量器（免费CPU可跑）
+# ---------------------------
+GEN_MODEL_NAME = "facebook/blenderbot_small-90M"
+EMB_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"  # 小、快、准
+
+tokenizer = BlenderbotSmallTokenizer.from_pretrained(GEN_MODEL_NAME)
+gen_model = BlenderbotSmallForConditionalGeneration.from_pretrained(GEN_MODEL_NAME)
+embedder = SentenceTransformer(EMB_MODEL_NAME)
 
-# 选用小模型，能在免费 CPU 上跑
-MODEL_NAME = "facebook/blenderbot_small-90M"
-tokenizer = BlenderbotSmallTokenizer.from_pretrained(MODEL_NAME)
-model = BlenderbotSmallForConditionalGeneration.from_pretrained(MODEL_NAME)
+# ---------------------------
+# 2) 你的“简历知识库”  —— 可继续补充
+#    规则：每一条是一个“可检索文档片段”
+# ---------------------------
+DOCS = [
+    # 教育
+    "Education: Columbia University, Master's in Statistics (Advanced Machine Learning Track), 2025–2027.",
+    "Education: Hong Kong Baptist University, B.Sc. Computer Science and Technology, First Class Honours, 2021–2025.",
+    "Exchange: University of Oxford, AI & Machine Learning programme, 2023.",
+    # 奖项
+    "Awards: First-class Scholarship (2021–2023), Second-class Scholarship (2024), MCM Honorable Mention (2024), Chinese Mathematical Competition Provincial 2nd Prize (2023).",
+    # 经验
+    "Experience: Social Media Analyst at StudyCo (part-time), Jan 2023–Present, Melbourne (Hybrid). Did audience/product data analysis with Python/Excel, wrote copy, edited short videos, increased Xiaohongshu followers from 3,000+ to 10,000+.",
+    "Experience: Data Science Intern at Wisers (May–Jul 2023), Shanghai (Hybrid). Weibo sentiment analysis; Chinese NLP preprocessing (jieba, custom dict, post classification); visualized with matplotlib, wordcloud; Tableau + MySQL real-time dashboards; SQL for data extraction.",
+    # 项目（挑重点）
+    "Project: Intelligent Retirement — Monte Carlo simulation; economic scenario generator; decision support for investment/consumption under uncertainty (2024).",
+    "Project: Predicting Intern Recruitment with Bayesian Networks — VE & Clique Tree inference; DAG with ~12 key nodes; interactive UI to explore outcomes (2024).",
+    "Project: LLM Check System — Flask + MySQL + SQLAlchemy; OOA/OOD; unit & integration testing; frontend HTML/CSS/JS (2024).",
+    "Project: Twitter Sentiment Analysis — NLP preprocessing; RNN/LSTM/BERT, best accuracy 89% (2024).",
+    "Project: Long-text abstraction & character relationship (T5 + PageRank + PySpark) — build character networks, accelerate top-20 character identification (2023).",
+    "Project: Online Restaurant Ordering System — ER with 10 entities/15 relationships; VIP, live order status, personalized recommendations; >60k users/products (2022).",
+    "Project: Path Planning for Mobile Robots — Q-learning vs DQN; DQN reached 94% in complex mazes (PyTorch) (2022).",
+    # 技能
+    "Skills: Python, Java, C, R, MATLAB; Web: HTML, CSS, JavaScript; Frameworks/Tools: PyTorch, Tableau, SQL, SPSS, Excel; Core: Machine Learning, Deep Learning, NLP, Data Analysis.",
+]
 
+# 预生成向量
+DOC_EMBS = embedder.encode(DOCS, normalize_embeddings=True)
+
+# ---------------------------
+# 3) Prompt 组件
+# ---------------------------
 SYSTEM_PROMPT = (
-    "You are Ivy (Jiexin Chen). Answer briefly and friendly in first person. "
-    "Focus on Ivy's background: HKBU CS (First Class Honours), Oxford exchange (AI/ML), "
-    "Columbia MS Statistics (Advanced ML Track, 2025–2027); Experience: StudyCo (social media analyst), "
-    "Wisers (data science intern, Weibo NLP sentiment). Skills: Python, Java, C, R, MATLAB, PyTorch, SQL, Tableau, Excel, HTML/CSS/JS; "
-    "Core: ML, DL, NLP, Data Analysis. If question is unrelated to Ivy, gently steer back."
+    "You are Ivy (Jiexin Chen). Answer in FIRST PERSON, concise, friendly, professional. "
+    "Only talk about Ivy's background, education, projects, skills, and experience. "
+    "If asked something unrelated, briefly steer back to career topics.\n"
 )
 
+STYLE_HINT = (
+    "Style: 2–5 short sentences. Use concrete facts from the provided context if relevant. "
+    "If the information is missing, say it briefly and suggest related areas I can talk about."
+)
+
+def build_prompt(question: str, retrieved_snippets: list[str]) -> str:
+    context_block = "\n".join(f"- {s}" for s in retrieved_snippets)
+    prompt = (
+        f"{SYSTEM_PROMPT}\n"
+        f"{STYLE_HINT}\n\n"
+        f"Context:\n{context_block}\n\n"
+        f"User: {question}\n"
+        f"Ivy:"
+    )
+    return prompt
+
+# ---------------------------
+# 4) 简���检索器（余弦相似度）
+# ---------------------------
+def retrieve(query: str, top_k: int = 4) -> list[str]:
+    q_emb = embedder.encode([query], normalize_embeddings=True)[0]
+    sims = (DOC_EMBS @ q_emb)  # 余弦相似度（已归一化）
+    idx = np.argsort(-sims)[:top_k]
+    return [DOCS[i] for i in idx]
+
+# ---------------------------
+# 5) 对话函数：检索 → 组 Prompt → 生成
+# ---------------------------
 def chat_fn(message, history):
-    # 把历史合并成一个上下文（小模型+免费CPU，尽量短）
-    context = SYSTEM_PROMPT + "\n\n"
-    for user, bot in history[-5:]:
-        context += f"User: {user}\nIvy: {bot}\n"
-    context += f"User: {message}\nIvy:"
-
-    inputs = tokenizer([context], return_tensors="pt")
-    reply_ids = model.generate(
+    # 最近几轮加入问题提示（帮助检索）
+    history_tail = " ".join([u for u, _ in history[-3:]]) if history else ""
+    query = (message or "").strip()
+    full_query = (query + " " + history_tail).strip() or "Ivy summary"
+
+    # 检索
+    retrieved = retrieve(full_query, top_k=4)
+
+    # 构造提示词
+    prompt = build_prompt(query, retrieved)
+
+    # 生成
+    inputs = tokenizer([prompt], return_tensors="pt")
+    output_ids = gen_model.generate(
         **inputs,
-        max_new_tokens=120,
-        do_sample=True,
-        top_p=0.9,
+        max_new_tokens=140,
         temperature=0.7,
+        top_p=0.9,
+        do_sample=True,
+        no_repeat_ngram_size=3,
+        repetition_penalty=1.15,
         pad_token_id=tokenizer.eos_token_id,
     )
-    reply = tokenizer.decode(reply_ids[0], skip_special_tokens=True)
+    text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+    # 抽取 "Ivy:" 之后的回答
+    if "Ivy:" in text:
+        text = text.split("Ivy:", 1)[-1].strip()
 
-    # 由于我们把上下文拼成单段，这里简单抽取最后一段作为回答
-    if "Ivy:" in reply:
-        reply = reply.split("Ivy:")[-1].strip()
-    return reply
+    return text
 
+# ---------------------------
+# 6) UI
+# ---------------------------
 demo = gr.ChatInterface(
     fn=chat_fn,
-    title="Ivy Chatbot (Free Demo)",
-    description="Ask Ivy about education, projects, skills, and experience. (Running on free CPU, first response may take a bit.)",
+    title="Ivy Chatbot (Free RAG Demo)",
+    description=(
+        "Ask me about my education, projects, skills, and experience. "
+        "Runs on free CPU with a small model + retrieval for better accuracy."
+    ),
     chatbot=gr.Chatbot(height=420, label="Chat"),
     textbox=gr.Textbox(placeholder="Ask me anything about Ivy...", autofocus=True),
     theme="soft",