Skip to main content
Glama

Two-Tower Recsys MCP — 神经检索,通过 MCP 提供服务

一个深度学习双塔推荐模型,在亚马逊自己的2023年评论语料库上训练,作为 MCP(Model Context Protocol)工具服务器提供服务,并带有 Streamlit 聊天前端,让 Gemini 代理代表你调用这些工具。

本 README 从头到尾介绍了整个流水线:模型是什么、如何训练、实际表现如何(实测而非估算)、MCP 服务器如何暴露它,以及如何运行或部署前端。


1. 这是什么

双塔模型是大型工业推荐系统背后的标准架构(这种模式——将用户和物品嵌入到同一向量空间的独立“塔”,训练使得相关对彼此靠近——与 YouTube、Pinterest 和亚马逊自己的检索系统在生产中使用的形状相同)。本项目从头实现了一个,在真实的亚马逊交互数据上训练,并通过 MCP 包装以供代理使用,而不是典型的 REST API。

为什么用 MCP 而不是 REST API? MCP 是 Anthropic 为连接 LLM 代理与工具和数据而引入的协议。将训练好的模型包装为 MCP 工具(而不是,比如说,Flask 端点)意味着任何兼容 MCP 的代理——Claude Desktop、本项目自己的 Streamlit+Gemini 前端,或任何其他 MCP 客户端——都可以直接调用 recommend_for_user、similar_items 等,由 LLM 根据自然语言决定何时以及如何调用它们。

Related MCP server: Modular RAG MCP Server

2. 架构

  • 用户塔:学习到的用户 ID 嵌入(64维)→ 2层MLP → 64维输出。

  • 物品塔:学习到的物品 ID 嵌入(64维)与产品标题的冻结 all-MiniLM-L6-v2 句子嵌入(384维,投影到64维)拼接 → 2层MLP → 64维输出。冻结的文本嵌入赋予模型冷启动能力——即使没有交互历史,仅凭标题也能将物品合理地放置在向量空间中。

  • 两个塔都输出 L2 归一化向量;相似度是点积(等价于余弦相似度)。

  • 训练损失:批内采样 softmax——对于一批 B 个(用户,物品)正样本对,批次中的每个其他物品都充当每个用户的负样本,并在生成的 B×B 相似度矩阵上应用交叉熵。这是在不进行显式负采样的情况下训练检索塔的标准且计算高效的方法。

  • 服务:物品嵌入预先计算一次,并在 FAISS(IndexFlatIP)中建立索引,用于快速最近邻检索。第二个 FAISS 索引基于原始(未训练的)MiniLM 标题嵌入构建,支持独立于训练协同信号的冷启动文本搜索。

        ┌────────────┐                          ┌────────────┐
        │  User ID   │                          │  Item ID   │
        └─────┬──────┘                          └─────┬──────┘
              │ embed(64)                              │ embed(64)
              ▼                                         ▼
        ┌────────────┐                    ┌──────────────────────────┐
        │  MLP (128) │                    │  Item title → MiniLM(384) │
        └─────┬──────┘                    └─────────────┬─────────────┘
              │                                          │ project(64)
              │                                          ▼
              │                                   concat(128) → MLP(128)
              ▼                                          ▼
        user vector (64, L2-norm)          item vector (64, L2-norm)
              └──────────────┬───────────────────────────┘
                              ▼
                    dot product = relevance score

3. 数据集

McAuley-Lab/Amazon-Reviews-2023(加州大学圣迭戈分校 McAuley 实验室),Video_Games 类别——原始评论 + 物品元数据,直接从 HuggingFace 下载。

步骤

数量

原始评论

4,624,615

原始用户 / 物品

2,766,656 / 137,249

5-core 过滤后(用户和物品交互数≥5)

857,505 交互

用户 / 物品(过滤后)

98,906 / 26,354

训练 / 验证 / 测试交互

659,693 / 98,906 / 98,906

分割协议——按时间戳排序,每个用户留出最后两个:每个用户最近的交互 → 测试,次近的 → 验证,其余 → 训练。这是一个时间分割,因此模型评估的是相对于训练数据预测真正未来行为的能力,而不是随机留出的交互(这会向训练泄露未来信息并夸大数字)。

4. 评估(真实、实测数字)

评估使用全目录排名——每个候选都与全部 26,354 个物品进行评分,而不是一小部分负样本子集。采样负样本评估(在较旧的 RecSys 论文中常见,例如仅对 99 个随机负样本进行排名)已知会大幅夸大离线指标,因此这是更困难、更诚实的协议。每个用户已看过的物品会从他们自己的候选排名中排除。

测试集——98,906 个用户,每个用户留出的最终交互:

指标

值

Recall@10

1.40%

NDCG@10

0.70%

HitRate@10

1.40%(在留一法下与 Recall@10 相同:每个用户恰好一个相关物品)

作为背景:在 26,354 个物品的目录中,k=10 的随机概率为 10/26,354 = 0.038%。训练后的模型在全目录排名下比随机好约 37 倍。

训练期间验证 Recall@10 在(第 142/150 轮)达到峰值 2.43%——测试数字较低,因为测试交互是每个用户相对于其训练历史最远的未来交互,这本质上更难预测。这种差距是时间分割的预期行为,而不是错误。测试数字(1.40%)是任何地方都应引用的数字——验证仅用于在训练期间选择最佳检查点,因此将其作为最终结果报告将是一种挑选。

完整训练曲线:models/train_history.csv。原始结果:models/test_results.json。

5. MCP 工具(mcp_server.py)

工具

描述

recommend_for_user(user_id, k)

Top-k 个性化推荐,排除用户已交互过的物品

similar_items(item_id, k)

通过训练后的物品塔嵌入进行物品到物品的相似度

search_items(query_text, k)

对物品标题进行冷启动语义搜索(仅 MiniLM——适用于协同模型信号较弱的物品)

explain_recommendation(user_id, item_id)

相似度分数加上与目标最相似的用户历史物品,用于可解释性

6. 前端(streamlit_app.py)

一个与 weather-mcp-server 风格相同的聊天 UI:它将 MCP 服务器作为子进程通过 stdio 启动,获取其工具模式,将其转换为 Gemini 函数调用声明,并运行一个代理循环——Gemini 根据你的消息决定调用 4 个工具中的哪一个(如果有),工具针对真实训练模型执行,结果被反馈以生成最终的自然语言响应。侧边栏显示每个工具的描述以及一个使用训练目录中真实 ID 的一键示例,还有一个包含模型评估统计的展开器。

7. 本地运行

uv venv --python 3.11 .venv
uv pip install -p .venv/bin/python -r requirements.txt

# one-time: reproduce the trained model from scratch
.venv/bin/python src/data_prep.py                  # downloads + filters the dataset
.venv/bin/python src/precompute_text_embeddings.py
.venv/bin/python src/train.py                       # ~150 epochs, ~40s/epoch on an M2 CPU
.venv/bin/python src/evaluate.py                    # writes models/test_results.json
.venv/bin/python src/build_index.py                 # builds FAISS indices for serving

# run the MCP server standalone (stdio transport)
.venv/bin/python mcp_server.py

# or run the chat frontend (spawns the MCP server itself)
cp .streamlit/secrets.toml.example .streamlit/secrets.toml   # then fill in your key
.venv/bin/streamlit run streamlit_app.py

如果未设置 .streamlit/secrets.toml(或 GEMINI_API_KEY 环境变量),应用会在运行时回退到在侧边栏请求密钥。

macOS 注意事项

在 macOS 上,faiss 和 torch 在 OpenMP 运行时初始化上存在冲突,这会导致 FAISS 搜索调用段错误,除非在 faiss 之前导入 torch/numpy,并设置 KMP_DUPLICATE_LIB_OK=TRUE 和 OMP_NUM_THREADS=1。这两点已在 mcp_server.py 和 src/build_index.py 内部处理。

8. 在 Streamlit Community Cloud 上部署

  1. 将此仓库推送到 GitHub(公开或私有——Community Cloud 都可以为个人账户部署)。

  2. 前往 share.streamlit.io,点击 New app,并将此仓库指向 streamlit_app.py 作为入口点。

  3. 在应用的 设置 → 密钥 中,添加:

    GEMINI_API_KEY = "your_gemini_api_key_here"

    这与本地 .streamlit/secrets.toml 使用的机制相同——密钥仅存在于 Streamlit 的密钥存储中,绝不会出现在仓库或 git 历史中,应用会自动读取它,因此访问者无需输入密钥。

  4. 部署。首次启动会较慢(约 1-2 分钟),因为它会下载 MiniLM 模型并加载 FAISS 索引;后续加载会很快。

关于仓库大小的说明:models/(约 110MB:训练检查点 + FAISS 索引)已提交,因此部署的应用无需在每次冷启动时重新训练。data/raw/(约 2.9GB 的原始 HuggingFace 下载)已被 gitignore,仅在你希望从头复现训练时才需要。

9. 技术栈

Python、PyTorch、FAISS、Sentence-Transformers(MiniLM)、FastMCP、MCP Python SDK、Google Gemini API、Streamlit、pandas、HuggingFace datasets/huggingface_hub。

Related MCP Connectors

Related MCP Servers

  • F
    license
    Not graded
    quality
    C
    maintenance
    Provides RAG-based knowledge retrieval and document management as MCP tools, supporting hybrid search, reranking, and retrieval process visualization.
    -
  • A
    license
    Not graded
    quality
    C
    maintenance
    A modular RAG framework exposing knowledge retrieval tools via MCP, enabling AI assistants to perform hybrid search, reranking, and multimodal document queries with full observability and evaluation.
    MIT
  • A
    license
    Not graded
    quality
    B
    maintenance
    Enables agents to turn documents into vectors, perform approximate nearest neighbor search with HNSW, and filter by metadata via MCP tools for RAG workflows.
    MIT