412 lines
15 KiB
Python
412 lines
15 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Script to visualize two CSV files using embedding-atlas with different colors for each file.
|
||
|
||
This script supports both command-line usage and API usage.
|
||
"""
|
||
|
||
import argparse
|
||
import pandas as pd
|
||
import matplotlib.pyplot as plt
|
||
import os
|
||
from typing import Optional, List, Dict, Any
|
||
import numpy as np
|
||
|
||
def launch_interactive_viewer(df: pd.DataFrame, text_column: str, port: int = 5055, host: str = "0.0.0.0"):
|
||
"""使用Python API启动交互式服务器
|
||
|
||
Args:
|
||
df: 原始的pandas DataFrame，包含投影坐标和邻居信息（作为Python对象，而非JSON字符串）
|
||
text_column: 用作文本内容的列名
|
||
port: 启动服务器的端口
|
||
host: 服务器绑定的主机地址
|
||
"""
|
||
try:
|
||
from embedding_atlas.server import make_server
|
||
from embedding_atlas.data_source import DataSource
|
||
from embedding_atlas.utils import Hasher
|
||
import pathlib
|
||
|
||
# 创建metadata - 使用与CLI一致的props字段结构
|
||
metadata = {
|
||
"props": {
|
||
"data": {
|
||
"id": "_row_index",
|
||
"text": text_column,
|
||
"projection": {"x": "projection_x", "y": "projection_y"},
|
||
"neighbors": "__neighbors"
|
||
},
|
||
"initialState": {"version": "0.0.0"}
|
||
}
|
||
}
|
||
|
||
# 生成数据集标识符
|
||
hasher = Hasher()
|
||
hasher.update(["combined_dataset"])
|
||
hasher.update(metadata)
|
||
identifier = hasher.hexdigest()
|
||
|
||
# 创建DataSource - 直接使用传入的原始DataFrame
|
||
# 确保df['__neighbors'] 是 Python dict + numpy arrays 的结构，而不是 JSON 字符串
|
||
dataset = DataSource(identifier, df, metadata)
|
||
|
||
# 获取静态文件路径
|
||
import embedding_atlas
|
||
static_path = str(
|
||
(pathlib.Path(embedding_atlas.__file__).parent / "static").resolve()
|
||
)
|
||
|
||
# 创建并启动服务器 - 使用默认的duckdb_uri参数(即"wasm")
|
||
app = make_server(dataset, static_path=static_path)
|
||
|
||
import uvicorn
|
||
print(f"Starting interactive viewer on http://{host}:{port}")
|
||
uvicorn.run(app, host=host, port=port)
|
||
|
||
except ImportError:
|
||
print("Error: embedding-atlas not found. Please install it with: pip install embedding-atlas")
|
||
except Exception as e:
|
||
print(f"Error launching interactive viewer: {e}")
|
||
|
||
def create_embedding_service(
|
||
texts1: List[str],
|
||
texts2: List[str],
|
||
labels: tuple = ("Group1", "Group2"),
|
||
port: int = 5055,
|
||
host: str = "0.0.0.0",
|
||
text_column: str = "text",
|
||
model: str = "all-MiniLM-L6-v2",
|
||
batch_size: int = 32,
|
||
umap_args: Optional[Dict[str, Any]] = None
|
||
):
|
||
"""从两组文本数据创建embedding可视化服务"""
|
||
import pathlib
|
||
import embedding_atlas
|
||
from embedding_atlas.projection import compute_text_projection
|
||
from embedding_atlas.server import make_server
|
||
from embedding_atlas.data_source import DataSource
|
||
from embedding_atlas.utils import Hasher
|
||
|
||
# 默认UMAP参数
|
||
if umap_args is None:
|
||
umap_args = {
|
||
"n_neighbors": 15,
|
||
"min_dist": 0.1,
|
||
"metric": "cosine"
|
||
}
|
||
|
||
# 1. 创建DataFrame
|
||
df1 = pd.DataFrame({text_column: texts1, 'source': labels[0]})
|
||
df2 = pd.DataFrame({text_column: texts2, 'source': labels[1]})
|
||
combined_df = pd.concat([df1, df2], ignore_index=True)
|
||
|
||
# 2. 添加必需的ID列
|
||
combined_df['_row_index'] = range(len(combined_df))
|
||
|
||
# 3. 计算embeddings和投影
|
||
try:
|
||
compute_text_projection(
|
||
combined_df,
|
||
text=text_column,
|
||
x="projection_x",
|
||
y="projection_y",
|
||
neighbors="__neighbors", # 改为双下划线
|
||
model=model,
|
||
batch_size=batch_size,
|
||
umap_args=umap_args
|
||
)
|
||
except ImportError:
|
||
print("Error: embedding-atlas not found. Please install it with: pip install embedding-atlas")
|
||
return
|
||
except Exception as e:
|
||
print(f"Error computing projections: {e}")
|
||
return
|
||
|
||
# 4. 创建metadata (使用与CLI一致的props字段结构)
|
||
metadata = {
|
||
"props": {
|
||
"data": {
|
||
"id": "_row_index",
|
||
"text": text_column,
|
||
"projection": {
|
||
"x": "projection_x",
|
||
"y": "projection_y"
|
||
},
|
||
"neighbors": "__neighbors"
|
||
},
|
||
"initialState": {"version": "0.0.0"}
|
||
}
|
||
}
|
||
|
||
# 5. 保存原始DataFrame用于DataSource
|
||
df_for_datasource = combined_df.copy()
|
||
|
||
# 6. 转换neighbors列为JSON字符串(仅用于保存parquet)
|
||
if '__neighbors' in combined_df.columns:
|
||
import json
|
||
combined_df['__neighbors'] = combined_df['__neighbors'].apply(
|
||
lambda x: json.dumps({
|
||
'ids': x['ids'].tolist() if hasattr(x['ids'], 'tolist') else list(x['ids']),
|
||
'distances': x['distances'].tolist() if hasattr(x['distances'], 'tolist') else list(x['distances'])
|
||
}) if x is not None else None
|
||
)
|
||
|
||
# 7. 保存parquet文件
|
||
parquet_path = "combined_dataset.parquet"
|
||
combined_df.to_parquet(parquet_path, index=False)
|
||
print(f"Dataset saved to {parquet_path}")
|
||
|
||
# 8. 生成数据集标识符
|
||
hasher = Hasher()
|
||
hasher.update(["custom_dataset"])
|
||
hasher.update(metadata)
|
||
identifier = hasher.hexdigest()
|
||
|
||
# 9. 使用原始DataFrame创建DataSource
|
||
dataset = DataSource(identifier, df_for_datasource, metadata)
|
||
|
||
# 10. 获取静态文件路径
|
||
static_path = str(
|
||
(pathlib.Path(embedding_atlas.__file__).parent / "static").resolve()
|
||
)
|
||
|
||
# 11. 创建并启动服务器 - 使用默认的duckdb_uri参数(即"wasm")
|
||
app = make_server(dataset, static_path=static_path)
|
||
import uvicorn
|
||
print(f"Starting interactive viewer on http://{host}:{port}")
|
||
uvicorn.run(app, host=host, port=port)
|
||
|
||
def visualize_csv_comparison(
|
||
csv1_path: str,
|
||
csv2_path: str,
|
||
column1: str = "smiles",
|
||
column2: str = "smiles",
|
||
output_path: str = "comparison_visualization.png",
|
||
label1: Optional[str] = None,
|
||
label2: Optional[str] = None,
|
||
launch_interactive: bool = False,
|
||
port: int = 5055,
|
||
host: str = "0.0.0.0",
|
||
model: str = "all-MiniLM-L6-v2",
|
||
batch_size: int = 32,
|
||
umap_args: Optional[Dict[str, Any]] = None
|
||
):
|
||
"""
|
||
Visualize two CSV files with embedding-atlas and create a combined plot with different colors.
|
||
|
||
Args:
|
||
csv1_path: Path to the first CSV file
|
||
csv2_path: Path to the second CSV file
|
||
column1: Column name for the first CSV file (default: "smiles")
|
||
column2: Column name for the second CSV file (default: "smiles")
|
||
output_path: Output visualization file path
|
||
label1: Label for the first dataset (default: filename)
|
||
label2: Label for the second dataset (default: filename)
|
||
launch_interactive: Whether to launch interactive viewer (default: False)
|
||
port: Port for interactive viewer (default: 5055)
|
||
host: Host for interactive viewer (default: "0.0.0.0")
|
||
model: Embedding model to use (default: "all-MiniLM-L6-v2")
|
||
batch_size: Batch size for embedding computation (default: 32)
|
||
umap_args: UMAP arguments as dictionary (default: None)
|
||
"""
|
||
|
||
# Generate default labels from filenames if not provided
|
||
if label1 is None:
|
||
label1 = os.path.splitext(os.path.basename(csv1_path))[0]
|
||
if label2 is None:
|
||
label2 = os.path.splitext(os.path.basename(csv2_path))[0]
|
||
|
||
# Read CSV files with error handling
|
||
try:
|
||
df1 = pd.read_csv(csv1_path)
|
||
df2 = pd.read_csv(csv2_path)
|
||
except FileNotFoundError as e:
|
||
print(f"Error: CSV file not found - {e}")
|
||
return
|
||
except pd.errors.EmptyDataError:
|
||
print("Error: One of the CSV files is empty")
|
||
return
|
||
except Exception as e:
|
||
print(f"Error reading CSV files: {e}")
|
||
return
|
||
|
||
# Validate columns
|
||
if column1 not in df1.columns:
|
||
print(f"Error: Column '{column1}' not found in {csv1_path}")
|
||
print(f"Available columns: {list(df1.columns)}")
|
||
return
|
||
if column2 not in df2.columns:
|
||
print(f"Error: Column '{column2}' not found in {csv2_path}")
|
||
print(f"Available columns: {list(df2.columns)}")
|
||
return
|
||
|
||
# Add source column to identify origin of each data point
|
||
df1['source'] = label1
|
||
df2['source'] = label2
|
||
|
||
# Add column to identify which file the data came from
|
||
df1['file_origin'] = 0 # First file
|
||
df2['file_origin'] = 1 # Second file
|
||
|
||
# Combine dataframes
|
||
combined_df = pd.concat([df1, df2], ignore_index=True)
|
||
|
||
# Add _row_index column for embedding-atlas DataSource (moved to earlier)
|
||
combined_df['_row_index'] = range(len(combined_df))
|
||
|
||
# Determine which column to use for text projection
|
||
text_column = column1 if column1 in combined_df.columns else column2
|
||
|
||
# Check if projection columns already exist
|
||
has_projection = 'projection_x' in combined_df.columns and 'projection_y' in combined_df.columns
|
||
|
||
if not has_projection:
|
||
print("Computing embeddings and projections...")
|
||
try:
|
||
from embedding_atlas.projection import compute_text_projection
|
||
|
||
# Use default UMAP args if not provided
|
||
if umap_args is None:
|
||
umap_args = {
|
||
"n_neighbors": 15,
|
||
"min_dist": 0.1,
|
||
"metric": "cosine"
|
||
}
|
||
|
||
# Compute projections
|
||
compute_text_projection(
|
||
combined_df,
|
||
text=text_column,
|
||
x="projection_x",
|
||
y="projection_y",
|
||
neighbors="__neighbors", # 改为双下划线
|
||
model=model,
|
||
batch_size=batch_size,
|
||
umap_args=umap_args
|
||
)
|
||
except ImportError:
|
||
print("Error: embedding-atlas not found. Please install it with: pip install embedding-atlas")
|
||
return
|
||
except Exception as e:
|
||
print(f"Error computing projections: {e}")
|
||
return
|
||
|
||
# Create visualization
|
||
print("Creating visualization...")
|
||
plt.figure(figsize=(12, 8))
|
||
|
||
# Plot data points with different colors based on file origin
|
||
colors = ['blue', 'red']
|
||
for i, (label, color) in enumerate(zip([label1, label2], colors)):
|
||
subset = combined_df[combined_df['file_origin'] == i]
|
||
if not subset.empty:
|
||
plt.scatter(
|
||
subset['projection_x'],
|
||
subset['projection_y'],
|
||
c=color,
|
||
label=label,
|
||
alpha=0.6,
|
||
s=30
|
||
)
|
||
|
||
plt.xlabel('Projection X')
|
||
plt.ylabel('Projection Y')
|
||
plt.title('Embedding Atlas Visualization Comparison')
|
||
plt.legend(loc='upper right')
|
||
plt.grid(True, alpha=0.3)
|
||
|
||
# Save the plot
|
||
plt.tight_layout()
|
||
plt.savefig(output_path, dpi=300, bbox_inches='tight')
|
||
plt.close()
|
||
|
||
print(f"Visualization saved to {output_path}")
|
||
|
||
# Keep a copy of the original DataFrame for the interactive viewer
|
||
# The viewer expects Python objects, not JSON strings
|
||
df_for_viewer = combined_df.copy()
|
||
|
||
# Convert __neighbors column to JSON strings before saving
|
||
# This is needed because DuckDB-WASM expects JSON strings for complex data types
|
||
if '__neighbors' in combined_df.columns:
|
||
import json
|
||
# 为保存到 parquet 文件的数据转换为 JSON 字符串
|
||
combined_df_for_save = combined_df.copy()
|
||
combined_df_for_save['__neighbors'] = combined_df['__neighbors'].apply(
|
||
lambda x: json.dumps({
|
||
'ids': x['ids'].tolist() if hasattr(x['ids'], 'tolist') else list(x['ids']),
|
||
'distances': x['distances'].tolist() if hasattr(x['distances'], 'tolist') else list(x['distances'])
|
||
}) if x is not None else None
|
||
)
|
||
|
||
# Save dataset for interactive viewing
|
||
parquet_path = "combined_dataset.parquet"
|
||
# 保存转换为 JSON 字符串的 DataFrame
|
||
if '__neighbors' in combined_df.columns:
|
||
combined_df_for_save.to_parquet(parquet_path, index=False)
|
||
else:
|
||
combined_df.to_parquet(parquet_path, index=False)
|
||
print(f"Dataset saved to {parquet_path}")
|
||
|
||
# Launch interactive viewer with original DataFrame (not converted to JSON)
|
||
if launch_interactive:
|
||
launch_interactive_viewer(df_for_viewer, text_column, port, host)
|
||
print(f"Interactive viewer available at http://{host}:{port}")
|
||
|
||
|
||
def main():
|
||
"""Command-line interface for the visualization script."""
|
||
parser = argparse.ArgumentParser(
|
||
description="Visualize two CSV files using embedding-atlas with different colors",
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
epilog="""
|
||
Examples:
|
||
python -m visualization.comparison file1.csv file2.csv
|
||
python -m visualization.comparison file1.csv file2.csv --column1 smiles --column2 SMILES
|
||
python -m visualization.comparison file1.csv file2.csv --label1 "Dataset A" --label2 "Dataset B"
|
||
python -m visualization.comparison file1.csv file2.csv --output comparison.png
|
||
python -m visualization.comparison file1.csv file2.csv --interactive --port 8080
|
||
"""
|
||
)
|
||
|
||
parser.add_argument("csv1", help="Path to the first CSV file")
|
||
parser.add_argument("csv2", help="Path to the second CSV file")
|
||
parser.add_argument("--column1", default="smiles", help="Column name for the first CSV file (default: smiles)")
|
||
parser.add_argument("--column2", default="smiles", help="Column name for the second CSV file (default: smiles)")
|
||
parser.add_argument("--output", "-o", default="comparison_visualization.png", help="Output visualization file path")
|
||
parser.add_argument("--label1", help="Label for the first dataset (default: filename)")
|
||
parser.add_argument("--label2", help="Label for the second dataset (default: filename)")
|
||
parser.add_argument("--interactive", "-i", action="store_true", help="Launch interactive viewer")
|
||
parser.add_argument("--port", "-p", type=int, default=5055, help="Port for interactive viewer (default: 5055)")
|
||
parser.add_argument("--host", default="0.0.0.0", help="Host for interactive viewer (default: 0.0.0.0)")
|
||
parser.add_argument("--model", default="all-MiniLM-L6-v2", help="Embedding model to use (default: all-MiniLM-L6-v2)")
|
||
parser.add_argument("--batch-size", type=int, default=32, help="Batch size for embedding computation (default: 32)")
|
||
|
||
args = parser.parse_args()
|
||
|
||
# Parse UMAP arguments
|
||
umap_args = {
|
||
"n_neighbors": 15,
|
||
"min_dist": 0.1,
|
||
"metric": "cosine"
|
||
}
|
||
|
||
visualize_csv_comparison(
|
||
args.csv1,
|
||
args.csv2,
|
||
args.column1,
|
||
args.column2,
|
||
args.output,
|
||
args.label1,
|
||
args.label2,
|
||
args.interactive,
|
||
args.port,
|
||
args.host,
|
||
args.model,
|
||
args.batch_size,
|
||
umap_args
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |