lingyuzeng
60c5ce152b
主要改动： 1. 新增CSV文件比较可视化功能： - 添加了src/visualization/comparison.py模块，支持比较两个CSV文件并使用不同颜色可视化 - 支持命令行和API两种使用方式 - 可生成静态图像或启动交互式查看器 - 支持自定义标签、模型和UMAP参数 2. 更新pixi.toml配置： - 添加linux-64平台支持 - 增加多个依赖项：ipykernel, anywidget, rdkit, selfies, fastapi, fastmcp, docker等 - 完善依赖版本约束 3. 更新README.md文档： - 添加CSV文件比较可视化功能说明和使用示例
273 lines
9.1 KiB
Markdown
273 lines
9.1 KiB
Markdown
# Embedding Atlas 项目说明
|
||
|
||
## 环境准备
|
||
|
||
项目使用 [uv](https://github.com/astral-sh/uv) 管理依赖，所有必需依赖已经记录在 `pyproject.toml` 的 `[project.dependencies]` 中。建议的设置流程如下：
|
||
|
||
```bash
|
||
# 可选：配置清华镜像（也可以写入 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` 中追加：
|
||
|
||
```toml
|
||
[tool.uv.pip]
|
||
index-url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple"
|
||
```
|
||
|
||
运行前如果需要离线或国内镜像 Hugging Face，可以设置：
|
||
|
||
```bash
|
||
export HF_HUB_OFFLINE=1
|
||
export HF_ENDPOINT=https://hf-mirror.com
|
||
```
|
||
|
||
## 会话编排服务（FastAPI / MCP）
|
||
|
||
使用 `uv run embedding-backend-api` 可以启动一个同时兼容 FastAPI 与 FastMCP 的后端服务。该服务监听 `/sessions` 路径，负责按需拉起 `embedding-atlas` 容器并在 10 小时后自动清理。
|
||
|
||
```bash
|
||
uv run embedding-backend-api
|
||
# 或以 MCP 模式启动（stdio）
|
||
uv run embedding-backend-mcp
|
||
```
|
||
|
||
### REST API 用法
|
||
|
||
```bash
|
||
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 小时存活时间。
|
||
|
||
示例请求：
|
||
|
||
```bash
|
||
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 响应示例
|
||
|
||
无查询参数，返回当前所有会话的列表：
|
||
|
||
```json
|
||
{
|
||
"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 工具：
|
||
|
||
```bash
|
||
uv run embedding-backend-mcp
|
||
```
|
||
|
||
FastMCP 客户端加载该配置后，可用标准 MCP 协议转发同一套 REST 接口，从而与传统后端保持一致的行为。
|
||
|
||
## 命令行生成嵌入可视化交互
|
||
|
||
```bash
|
||
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` 更多用法示例：
|
||
|
||
```bash
|
||
# 本地文件
|
||
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使用方法
|
||
|
||
```python
|
||
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"
|
||
}
|
||
)
|
||
```
|
||
|
||
### 命令行使用方法
|
||
|
||
```bash
|
||
# 基本用法
|
||
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 第一轮微调数据集
|
||
|
||
```bash
|
||
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 分布接近训练集，便于后续"用未见参考分子做条件生成并观察邻域覆盖"。
|
||
|
||
## 合并分割的数据集进行可视化
|
||
|
||
合并数据集：
|
||
|
||
```bash
|
||
uv run python script/merge_splits.py --input-dir splits_v2/ --output data/drugbank_split_merge.csv
|
||
```
|
||
|
||
可视化：
|
||
|
||
```bash
|
||
uv run embedding-atlas data/drugbank_split_merge.csv --text smiles
|
||
```
|
||
|
||
## 容器化部署
|
||
|
||
项目提供 `docker/` 目录用于快速启动后端：
|
||
|
||
1. 先构建可被 orchestrator 复用的 `embedding-atlas` 镜像：
|
||
```bash
|
||
docker build -f docker/embedding-atlas.Dockerfile -t embedding-atlas:latest .
|
||
```
|
||
2. 启动 DIND + orchestrator 组合：
|
||
```bash
|
||
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` 或额外的反向代理。 |