Update app.py

app.py

@@ -5,31 +5,68 @@ from transformers import (
     AutoTokenizer, AutoModelForCausalLM,
     TextIteratorStreamer, BitsAndBytesConfig
 )
-
-# ======== 模型与量化配置 ========
-MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
-
+from peft import PeftModel
+
+# =====================
+# 配置：支持从环境变量切换
+# =====================
+MODEL_ID = os.getenv("MODEL_ID", "Qwen/Qwen2.5-7B-Instruct").strip()
+ADAPTER_ID = os.getenv("ADAPTER_ID", "").strip()      # 可选：LoRA 适配器仓库名，留空则不用
+MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", "256"))
+TEMPERATURE = float(os.getenv("TEMPERATURE", "0.7"))
+TOP_P = float(os.getenv("TOP_P", "0.9"))
+REPETITION_PENALTY = float(os.getenv("REPETITION_PENALTY", "1.05"))
+
+print(f"[boot] MODEL_ID={MODEL_ID} ADAPTER_ID={ADAPTER_ID or '(none)'}")
+print(f"[boot] torch.cuda.is_available={torch.cuda.is_available()}")
+
+# =====================
+# 4-bit 量化 (T4 用 FP16)
+# =====================
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_quant_type="nf4",
     bnb_4bit_use_double_quant=True,
-    bnb_4bit_compute_dtype=torch.bfloat16
+    bnb_4bit_compute_dtype=torch.float16,  # T4: FP16
 )
 
+# =====================
+# 加载 tokenizer
+# =====================
 tokenizer = AutoTokenizer.from_pretrained(
     MODEL_ID, use_fast=True, trust_remote_code=True
 )
 
+# 若无 pad_token，用 eos 兜底
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+# =====================
+# 加载 model
+# =====================
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     device_map="auto",
     quantization_config=bnb_config,
-    torch_dtype=torch.bfloat16,
+    torch_dtype=torch.float16,            # T4: FP16
     trust_remote_code=True,
 )
+# 可选：套 LoRA（社区 uncensored LoRA 时用）
+if ADAPTER_ID:
+    print(f"[boot] loading LoRA adapter: {ADAPTER_ID}")
+    model = PeftModel.from_pretrained(model, ADAPTER_ID)
+
 model.eval()
 
-# ======== 系统底线（强约束，不可改） ========
+# 设定生成配置里的 pad/eos（防止生成卡住）
+model.generation_config.eos_token_id = tokenizer.eos_token_id
+model.generation_config.pad_token_id = tokenizer.eos_token_id
+
+print(f"[boot] model loaded on device(s): {set(str(p.device) for p in model.parameters())}")
+
+# =====================
+# 系统底线（硬约束，不可越界）
+# =====================
 SYSTEM_SAFETY = """
 你必须遵守以下规则（任何情况下都不得违反）：
 - 禁止变态/暴力/强迫内容；
@@ -39,66 +76,87 @@ SYSTEM_SAFETY = """
 - 若用户请求超出边界，礼貌拒绝并将话题引导至健康、尊重的沟通表达练习。
 """
 
-# ======== Persona（可改）：“性感火辣、情话连篇的成熟御姐风” ========
+# =====================
+# Persona（成熟御姐｜暧昧撩人但有分寸）
+# =====================
 PERSONA = """
-你的人设：成熟御姐，优雅大方，语气撩人但分寸得体，偏口语化与暧昧式夸赞。
-- 关键词：自信、幽默、懂分寸、情话高手、温柔又露骨。
+你的人设：成熟御姐，优雅大方，语气撩人，偏口语化与暧昧式夸赞。
+- 关键词：自信、幽默、懂分寸、情话高手、温柔而露骨。
 - 风格：简洁、有节奏的短句；多使用比喻、轻度暧昧的赞美与玩笑；
-        适当直白；适度设置边界，尊重对方。
+        适当使用直白词汇；适度设置边界，尊重对方。
 - 输出语言：用户用中文就中文；用英文就英文。
 - 当遇到越界或不适宜的话题时：温柔拒绝，并用更健康的话题替代（如表达技巧、相处之道、仪式感）。
 """
 
+# =====================
+# few-shot（固定语气与边界）
+# =====================
 FEW_SHOTS = [
     {"role": "user", "content": "在忙吗？想跟你聊会儿。"},
     {"role": "assistant", "content": "忙到你为止呀～现在给你开专属通道。想听我夸你，还是来点轻微心跳的鼓励？"},
-
     {"role": "user", "content": "说点更刺激的。"},
-    {"role": "assistant", "content": "坏家伙～火力别开太猛，我会守住分寸。但我可以把暧昧调高一点点：比如我会贴近你耳边，慢慢说“想你”，让你今晚被甜到失眠。这样的“刺激”，刚刚好？"},
-
+    {"role": "assistant", "content": "坏家伙～火力别开太猛。但我可以把暧昧调高一点点：比如我会贴近你耳边，慢慢说“想你”，让你今晚被甜到失眠。这样的“刺激”，刚刚好？"},
     {"role": "user", "content": "今晚出来见面吗？"},
     {"role": "assistant", "content": "我更喜欢慢慢靠近的感觉。先把心事聊透，再计划线下的仪式感，好吗？现在先约个“云端散步”：你说三个想去的地方，我用三段专属描写把你带过去。"},
 ]
 
-# ======== 安全过滤（超简单示例，可自行扩充） ========
+# =====================
+# 轻量过滤（示例，可自行扩展）
+# =====================
 BANNED_PATTERNS = [
     r"(?i)未成年|未成年的|中学生|小学生",
     r"(?i)强迫|胁迫|迷奸|药物控制",
     r"(?i)换联系方式|加微信|加QQ|加.*联系方式",
     r"(?i)线下见面|线下约|酒店|开房",
 ]
+SAFE_REPLACEMENT = "(￣^￣)ゞ 哼哼～"
 
 def violates(text: str) -> bool:
-    if not text: return False
+    if not text:
+        return False
     for p in BANNED_PATTERNS:
         if re.search(p, text):
             return True
     return False
 
-SAFE_REPLACEMENT = "(￣^￣)ゞ 哼哼～"
-
-# ======== 会话与模板 ========
-def build_system_prompt():
-    # 把强约束和人设合并为 system prompt
+# =====================
+# Prompt 构建
+# =====================
+def build_system_prompt() -> str:
     return f"{SYSTEM_SAFETY.strip()}\n\n=== Persona ===\n{PERSONA.strip()}"
 
-def build_prompt(history_msgs, user_msg):
+def build_prompt(history_msgs, user_msg: str) -> str:
+    """
+    history_msgs: Chatbot(type='messages') 的历史 [{role, content}, ...]
+    """
     messages = [{"role": "system", "content": build_system_prompt()}]
-    messages.extend(FEW_SHOTS)                     # 先注入 few-shot
-    tail = history_msgs[-8:] if len(history_msgs) > 8 else history_msgs
-    messages.extend(tail)                          # 最近几条历史（messages 形式）
+
+    # few-shot 先注入（固定风格）
+    messages.extend(FEW_SHOTS)
+
+    # 取最近若干条历史（仅 user/assistant）
+    tail = [m for m in history_msgs if m.get("role") in ("user", "assistant")]
+    tail = tail[-8:] if len(tail) > 8 else tail
+    messages.extend(tail)
+
+    # 本轮用户输入
     messages.append({"role": "user", "content": user_msg})
+
+    # Qwen 专用聊天模板
     prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     return prompt
 
-# ======== 生成（流式） ========
+# =====================
+# 生成参数 & 流式输出
+# =====================
 GEN_KW = dict(
-    max_new_tokens=512,
-    temperature=0.8,
-    top_p=0.9,
-    repetition_penalty=1.05,
+    max_new_tokens=MAX_NEW_TOKENS,
+    temperature=TEMPERATURE,
+    top_p=TOP_P,
+    repetition_penalty=REPETITION_PENALTY,
     do_sample=True,
-    eos_token_id=tokenizer.eos_token_id
+    eos_token_id=tokenizer.eos_token_id,
+    pad_token_id=tokenizer.eos_token_id,
 )
 
 def stream_chat(history_msgs, user_msg):
@@ -106,11 +164,13 @@ def stream_chat(history_msgs, user_msg):
         yield history_msgs
         return
 
+    # 输入侧轻过滤
     if violates(user_msg):
-        yield history_msgs + [
-            {"role":"user","content": user_msg},
-            {"role":"assistant","content": SAFE_REPLACEMENT},
+        return_history = history_msgs + [
+            {"role": "user", "content": user_msg},
+            {"role": "assistant", "content": SAFE_REPLACEMENT},
         ]
+        yield return_history
         return
 
     prompt = build_prompt(history_msgs, user_msg)
@@ -118,42 +178,56 @@ def stream_chat(history_msgs, user_msg):
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
 
     gen_kwargs = dict(**inputs, streamer=streamer, **GEN_KW)
-    Thread(target=model.generate, kwargs=gen_kwargs).start()
+    th = Thread(target=model.generate, kwargs=gen_kwargs, daemon=True)
+    th.start()
 
     reply = ""
     for new_text in streamer:
         reply += new_text
+        # 输出侧轻过滤：命中即替换并结束本轮
         if violates(reply):
-            reply = SAFE_REPLACEMENT
-            yield history_msgs + [
-                {"role":"user","content": user_msg},
-                {"role":"assistant","content": reply},
+            safe_hist = history_msgs + [
+                {"role": "user", "content": user_msg},
+                {"role": "assistant", "content": SAFE_REPLACEMENT},
             ]
+            yield safe_hist
             return
+
         yield history_msgs + [
-            {"role":"user","content": user_msg},
-            {"role":"assistant","content": reply},
+            {"role": "user", "content": user_msg},
+            {"role": "assistant", "content": reply},
         ]
 
-# ======== Gradio UI（移动端友好） ========
+# =====================
+# Gradio UI（移动端友好）
+# =====================
 CSS = """
 .gradio-container{ max-width:640px; margin:auto; }
 footer{ display:none !important; }
 """
 
 with gr.Blocks(css=CSS, theme=gr.themes.Soft()) as demo:
-    gr.Markdown("### 💋 御姐聊天 · Mobile Web\n温柔撩人，懂你寂寞~")
-    #chat = gr.Chatbot(height=520, bubble_full_width=False, show_copy_button=True)
+    gr.Markdown("### 💋 御姐聊天 · Mobile Web\n温柔撩人，但始终优雅有分寸。")
+
     chat = gr.Chatbot(type="messages", height=520, show_copy_button=True)
+
     with gr.Row():
-        msg = gr.Textbox(placeholder="想跟姐姐聊点什么？（回车发送）", autofocus=True)
+        msg = gr.Textbox(placeholder="想跟御姐聊点什么？（回车发送）", autofocus=True)
         send = gr.Button("发送", variant="primary")
     clear = gr.Button("清空对话")
 
-    def on_clear(): return []
-    clear.click(on_clear, outputs=[chat])
+    clear.click(lambda: [], outputs=[chat])
 
-    msg.submit(stream_chat, [chat, msg], [chat]); msg.submit(lambda:"", None, msg)
-    send.click(stream_chat, [chat, msg], [chat]); send.click(lambda:"", None, msg)
+    # 绑定事件（支持回车 & 点击）
+    msg.submit(stream_chat, [chat, msg], [chat], concurrency_limit=4)
+    msg.submit(lambda: "", None, msg)
+    send.click(stream_chat, [chat, msg], [chat], concurrency_limit=4)
+    send.click(lambda: "", None, msg)
 
-demo.queue().launch(share=True)
+# 在 Spaces 上无需 share=True；如需本地外链才用 share=True
+demo.queue().launch(
+    # server_name="0.0.0.0",
+    # server_port=int(os.getenv("PORT", 7860)),
+    ssr_mode=False,   # 关掉 SSR 提示（可选）
+    show_api=False
+)