Embedding Atlas 项目说明

环境准备

项目使用 uv 管理依赖，所有必需依赖已经记录在 pyproject.toml 的 [project.dependencies] 中。建议的设置流程如下：

# 可选：配置清华镜像（也可以写入 pyproject.toml 的 [tool.uv.pip]）
export UV_PIP_INDEX_URL=https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple

# 同步依赖并创建虚拟环境（默认 .venv）
uv lock            # 首次或需要更新锁文件时执行
uv sync            # 生成或更新虚拟环境

# 进入虚拟环境
source .venv/bin/activate

# 或直接使用 uv run
uv run streamlit run app.py

如需固定镜像源，可在 pyproject.toml 中追加：

[tool.uv.pip]
index-url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple"

运行前如果需要离线或国内镜像 Hugging Face，可以设置：

export HF_HUB_OFFLINE=1
export HF_ENDPOINT=https://hf-mirror.com

数据可视化工具

项目包含一个强大的数据可视化工具 comparison.py，可以比较两个 CSV 文件中的数据并在 2D 空间中可视化。

安装和首次使用注意事项

首次运行时，系统需要下载 embedding 模型权重（如 all-MiniLM-L6-v2）。如果下载失败，请设置 Hugging Face 镜像：

export HF_ENDPOINT=https://hf-mirror.com

工具使用原则

在使用该工具时，请遵循以下核心设计原则：

1. 让 make_server() 自动处理数据库配置
2. 使用完整的 props 格式提供前端所需的所有信息
3. 避免手动添加可能与系统自动添加冲突的配置字段

命令行使用方式

python src/visualization/comparison.py file1.csv file2.csv \
  --column1 smiles --column2 smiles \
  --label1 "Dataset A" --label2 "Dataset B" \
  --interactive --port 5055

参数说明：

• file1.csv 和 file2.csv：要比较的两个 CSV 文件
• --column1 和 --column2：分别指定两个文件中用于生成 embedding 的列名
• --label1 和 --label2：在可视化中显示的数据集标签
• --interactive：启动交互式 Web 查看器
• --port：指定 Web 服务器端口
• --model：指定要使用的 embedding 模型（默认：all-MiniLM-L6-v2）
• --batch-size：指定处理数据的批大小（默认：32）
• --output 或 -o：指定输出图像文件路径（默认：comparison_visualization.png）
• --host：指定 Web 服务器主机地址（默认：0.0.0.0）

Python API 调用方式

工具也支持作为 Python 模块直接调用：

from src.visualization.comparison import visualize_csv_comparison

# 基本用法
visualize_csv_comparison(
    "file1.csv", 
    "file2.csv",
    column1="smiles",
    column2="smiles",
    launch_interactive=True,
    port=5055
)

# 高级用法 - 自定义模型和参数
visualize_csv_comparison(
    "file1.csv",
    "file2.csv",
    column1="smiles",
    column2="smiles",
    model="sentence-transformers/all-mpnet-base-v2",  # 使用不同的模型
    batch_size=16,  # 调整批处理大小
    output_path="custom_output.png",  # 自定义输出路径
    launch_interactive=True,
    port=8080,  # 自定义端口
    umap_args={  # 自定义 UMAP 参数
        "n_neighbors": 20,
        "min_dist": 0.2,
        "metric": "cosine"
    }
)

# 自定义 embedding 服务
from src.visualization.comparison import create_embedding_service

create_embedding_service(
    ["text1", "text2", "text3"],  # 第一组文本数据
    ["text4", "text5", "text6"],  # 第二组文本数据
    labels=("Group A", "Group B"),
    model="sentence-transformers/all-mpnet-base-v2",  # 指定模型
    batch_size=16,  # 批处理大小
    port=5055
)

支持的模型

工具支持任何兼容 Sentence Transformers 的模型，包括但不限于：

• all-MiniLM-L6-v2（默认）
• all-mpnet-base-v2
• all-distilroberta-v1
• paraphrase-multilingual-MiniLM-L12-v2

高级功能

1. 自定义 UMAP 参数： 可以通过 umap_args 参数调整降维效果：

   umap_args = {
       "n_neighbors": 15,      # 邻居数量
       "min_dist": 0.1,        # 最小距离
       "metric": "cosine"      # 距离度量
   }
   

2. 批处理优化： 对于大型数据集，可以通过调整 batch_size 参数来平衡内存使用和处理速度。

会话编排服务（FastAPI / MCP）

使用 uv run embedding-backend-api 可以启动一个同时兼容 FastAPI 与 FastMCP 的后端服务。该服务监听 /sessions 路径，负责按需拉起 embedding-atlas 容器并在 10 小时后自动清理。

uv run embedding-backend-api
# 或以 MCP 模式启动（stdio）
uv run embedding-backend-mcp

REST API 用法

curl -X POST http://localhost:9000/sessions \
  -H 'Content-Type: application/json' \
  -d '{
    "data_url": "https://example.com/data.csv",
    "extra_args": ["--text", "smiles"]
  }'

# 关闭会话
curl -X DELETE http://localhost:9000/sessions/<session_id>

# 查看当前会话
curl http://localhost:9000/sessions

请求体支持 session_id、port、auto_remove_seconds、environment 等字段；省略 session_id 时会自动生成并在响应中返回。返回结果包含容器名称与可访问的前端地址，供 FastAPI 或 MCP 客户端转发给前端使用。

如果希望以交互方式测试 API，可在浏览器访问 http://localhost:9000/docs 打开自动生成的 Swagger UI。

请求体字段说明

• session_id：可选，自定义的会话 ID；省略时系统自动生成。
• data_url：必填，指向 CSV/Parquet 的 HTTPS 链接，会被下载到 orchestrator 与 DinD 共享的 /sessions/<session_id>/ 中。
• input_filename：可选，保存到共享卷时使用的文件名；默认根据 URL 推断。
• extra_args：可选，附加给 embedding-atlas CLI 的参数数组，例如 ["--text", "smiles"]。
• environment：可选，注入到容器内的环境变量映射。
• labels：可选，附加到容器上的 Docker labels（在默认的 embedding-backend.* 标签之外），便于自定义监控或追踪。
• image：可选，覆盖默认的 embedding-atlas 镜像名。
• host：可选，传递给 embedding-atlas --host 的值，默认 0.0.0.0。
• port：可选，显式指定宿主机端口；若缺省会自动在配置的区间内分配。
• auto_remove_seconds：可选，覆盖默认的 10 小时存活时间。

示例请求：

curl -X POST http://localhost:9000/sessions \
  -H 'Content-Type: application/json' \
  -d '{
    "session_id": "demo-session",
    "data_url": "https://example.com/data.csv",
    "input_filename": "data.csv",
    "extra_args": ["--text", "smiles"],
    "environment": {"HF_ENDPOINT": "https://hf-mirror.com"},
    "labels": {"team": "chem"},
    "auto_remove_seconds": 7200
  }'

GET /sessions 响应示例

无查询参数，返回当前所有会话的列表：

{
  "sessions": [
    {
      "session_id": "demo-session",
      "container_id": "e556c0f1c35b...",
      "container_name": "embedding-atlas_demo-session",
      "port": 6000,
      "host": "0.0.0.0",
      "started_at": "2025-09-22T14:37:25.206038Z",
      "expires_at": "2025-09-23T00:37:24.416241Z",
      "dataset_path": "/sessions/demo-session/data.csv"
    }
  ]
}

• sessions 数组为空时表示当前没有存活的容器。
• dataset_path 指向共享卷内对应数据集的绝对路径。

MCP 集成

根目录的 fastmcp.json 示例可直接将本项目注册为 MCP 工具：

uv run embedding-backend-mcp

FastMCP 客户端加载该配置后，可用标准 MCP 协议转发同一套 REST 接口，从而与传统后端保持一致的行为。

命令行生成嵌入可视化交互

uv run embedding-atlas data/drugbank_pre_filtered_mordred_qed_id_selfies.csv --text smiles
uv run embedding-atlas data/drugbank_pre_filtered_mordred_qed_id_selfies.csv --export-application data/my_visualization.zip

embedding-atlas 更多用法示例：

# 本地文件
embedding-atlas dataset.parquet
# Hugging Face 数据集
embedding-atlas huggingface_org/dataset_name
# 指定文本列
embedding-atlas dataset.parquet --text text_column
# 使用预计算坐标
embedding-atlas dataset.parquet --x projection_x --y projection_y

CSV文件比较可视化

本项目提供了一个强大的工具用于比较两个CSV文件中的分子数据，并使用Embedding Atlas进行可视化。

Python API使用方法

from script.visualize_csv_comparison import visualize_csv_comparison, create_embedding_service

# 比较两个CSV文件
visualize_csv_comparison(
    "file1.csv",
    "file2.csv",
    column1="smiles",
    column2="smiles",
    output_path="comparison.png",
    label1="Dataset A",
    label2="Dataset B",
    launch_interactive=True,
    port=5055,
    model="all-MiniLM-L6-v2"
)

# 直接从文本列表创建可视化服务
texts1 = ["CCO", "CCN", "CCC"]
texts2 = ["c1ccccc1", "CCCN", "CCCO"]

create_embedding_service(
    texts1,
    texts2,
    labels=("Alcohols", "Others"),
    port=8080,
    model="all-MiniLM-L6-v2",
    umap_args={
        "n_neighbors": 30,
        "min_dist": 0.05,
        "metric": "cosine"
    }
)

命令行使用方法

# 基本用法
python script/visualize_csv_comparison.py file1.csv file2.csv

# 指定不同的列名
python script/visualize_csv_comparison.py file1.csv file2.csv \
    --column1 smiles --column2 SMILES

# 自定义标签和输出文件
python script/visualize_csv_comparison.py file1.csv file2.csv \
    --label1 "Dataset A" --label2 "Dataset B" \
    --output comparison.png

# 启动交互式查看器
python script/visualize_csv_comparison.py file1.csv file2.csv \
    --interactive --port 8080

# 使用自定义模型和参数
python script/visualize_csv_comparison.py file1.csv file2.csv \
    --interactive \
    --model all-MiniLM-L6-v2 \
    --batch-size 64

功能特点

1. 双模式支持：既可以生成静态可视化图像，也可以启动交互式查看器
2. 数据源区分：在可视化中使用不同颜色区分两个数据源，并在图例中标识
3. 自动标注：默认使用文件名作为数据源标签，也支持自定义标签
4. 端口配置：可以自定义交互式查看器的端口
5. 模型选择：支持指定不同的SentenceTransformer模型
6. UMAP参数调优：可以自定义UMAP算法参数以获得更好的可视化效果
7. 灵活输入：支持直接从文本列表创建可视化服务，无需CSV文件

划分 MolGen 第一轮微调数据集

uv run python script/split_drugbank.py \
  --in-csv data/drugbank_pre_filtered_mordred_qed_id_selfies.csv \
  --out-dir splits_v2 --seed 20250922 \
  --train-ratio 0.8 --val-ratio 0.1 --test-ratio 0.1 \
  --n_qed_bins 5 --n_mw_bins 5 --largest-first

产物：split_train.csv / split_val.csv / split_test.csv 其中 split_val 和 split_test 中的分子不会出现在训练集里，且整体 QED/MW 分布接近训练集，便于后续"用未见参考分子做条件生成并观察邻域覆盖"。

合并分割的数据集进行可视化

合并数据集：

uv run python script/merge_splits.py --input-dir splits_v2/ --output data/drugbank_split_merge.csv

可视化：

uv run embedding-atlas data/drugbank_split_merge.csv --text smiles

容器化部署

项目提供 docker/ 目录用于快速启动后端：

1. 先构建可被 orchestrator 复用的 embedding-atlas 镜像：

   docker build -f docker/embedding-atlas.Dockerfile -t embedding-atlas:latest .
   

2. 启动 DIND + orchestrator 组合：

   docker compose -f docker/docker-compose.yml up --build
   

   • engine 服务运行 docker:dind，对外暴露 tcp://localhost:2375 供 orchestrator 通过 socket 管理容器；
   • 两个服务通过 sessions-data 卷共享 /sessions 目录，后端会把下载的数据放到这里；
   • orchestrator 服务运行 FastAPI/FastMCP 后端，默认开放 http://localhost:9000；
   • 会话缓存保存在 sessions-data 卷，可在容器重启时保留下载的数据集。

如需调整默认镜像或端口，可在 docker/docker-compose.yml 中覆盖 EMBEDDING_* 环境变量，或在部署时通过 .env 文件注入。

注意：由于 orchestrator 运行在容器内并通过 DinD 调度新容器，如需在宿主机直接访问 embedding-atlas 的 Web UI，需要确保相应端口从 docker_engine_1 转发到宿主，可按需求使用 podman port 或额外的反向代理。