Gradio + Transformers** 实现带记忆功能的对话系统完整代码示例

以下是一个使用 Gradio + Transformers 实现带记忆功能的对话系统完整代码示例,无需额外数据库依赖:

import gradio as gr
import numpy as np
from transformers import AutoTokenizer, AutoModel
from typing import Dict, List, Tuple

# ----------------- 初始化模型 -----------------
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm3-6b", trust_remote_code=True)
model = AutoModel.from_pretrained("THUDM/chatglm3-6b", trust_remote_code=True).quantize(4).cuda()

# ----------------- 记忆存储结构 -----------------
class MemoryStorage:
    def __init__(self):
        self.sessions: Dict[str, dict] = {
   }  # {session_id: {"history": [], "vectors": []}}
    
    def add_memory(self, session_id: str, text: str):
        # 生成文本向量(使用模型最后一层隐状态)
        inputs = tokenizer(text, return_tensors="pt").to(model.device)
        outputs = model(**inputs)
        vector = outputs.last_hidden_state[:,0,:].cpu().detach().numpy()[0]
        
        # 初始化会话存储
        if session_id not in self.sessions:
            self.sessions[session_id] = {
   "history": [], "vectors": []}
        
        # 保存记忆
        self.sessions[session_id]["history"]<

你可能感兴趣的:(python,python,开发语言)