# ann-gen-data

**Repository Path**: SmartDigit/ann-gen-data

## Basic Information

- **Project Name**: ann-gen-data
- **Description**: No description available
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: main
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2025-04-27
- **Last Updated**: 2025-06-12

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# ANN-GEN-DATA

这是一个用于生成和处理向量数据的工具包，主要用于生成用于向量检索基准测试的数据集。

## 功能特点

- 支持多种向量数据格式的读写（HDF5、IVECS、FVECS、BVECS等）
- 支持向量数据的变换和生成
- 支持多种距离度量方式（欧氏距离、余弦相似度、点积）
- 支持大规模数据集的处理
- 支持数据集的扩展和维度调整

## 安装

1. 克隆仓库：
```bash
git clone https://gitee.com/SmartDigit/ann-gen-data.git
cd ann-gen-data
```

2. 安装依赖：
```bash
pip install -r requirements.txt
```

## 快速开始

### CSV转HDF5快速示例

```bash
# 1. 或者使用自己的CSV文件
python csv_to_hdf5.py --train_csv your_train.csv --test_csv your_test.csv --output result.hdf5
```

## 使用方法

### 1. 生成LAION数据集

```python
from dataset.dataset_processor import get_laion_hdf5

# 生成完整的LAION数据集
get_laion_hdf5(chunk_size=1000000)

# 生成部分LAION数据集
get_laion_hdf5(
    chunk_size=100000, 
    output_file="laion-10M-768.hdf5", 
    max_train_size=10000000,
    max_test_size=1000
)
```

### 2. 生成OpenAI数据集

```python
from dataset.dataset_processor import get_openai_hdf5

# 生成OpenAI数据集
get_openai_hdf5(
    chunk_size=1000000, 
    output_file="openai_1536_500W_100.hdf5", 
    dataset_size=50000
)
```

### 3. 扩展HDF5文件

```python
from data.io import expand_hdf5

# 扩展HDF5文件并调整维度
expand_hdf5(
    input_file="input.hdf5",
    output_file="output.hdf5",
    expand_factor=2,
    randomness_factor=0.2,
    similarity_type='cosine',
    target_dim=1024,
    max_train_size=1000000
)
```

### 4. CSV转HDF5格式

将CSV格式的向量数据转换为HDF5格式，用于向量检索基准测试。

#### 命令行使用

```bash
# 基本用法
python csv_to_hdf5.py --train_csv train.csv --test_csv test.csv --output output.hdf5

# 指定距离度量和参数
python csv_to_hdf5.py \
    --train_csv train.csv \
    --test_csv test.csv \
    --output output.hdf5 \
    --distance angular \
    --k 50 \
    --chunk_size 50000

# 限制数据量
python csv_to_hdf5.py \
    --train_csv train.csv \
    --test_csv test.csv \
    --output output.hdf5 \
    --max_train_rows 100000 \
    --max_test_rows 1000
```

#### Python代码使用

```python
from csv_to_hdf5 import csv_to_hdf5

# 转换CSV到HDF5
csv_to_hdf5(
    train_csv="train.csv",
    test_csv="test.csv", 
    output_hdf5="output.hdf5",
    distance_metric="euclidean",
    k=100,
    chunk_size=100000
)
```

#### CSV文件格式要求

**训练数据文件 (train.csv):**
- **列数**: 2列
- **第一列**: ID（字符串或数字，用于标识向量）
- **第二列**: 向量数据（支持多种格式，见下方说明）
- **编码**: UTF-8
- **分隔符**: 逗号(,)

示例：
```csv
id_0,"[0.0, 16.0, 35.0, 5.0, 32.0]"
id_1,"[14.0, 35.0, 19.0, 20.0, 3.0]"
100,"{1.0,2.0,3.0,4.0}"
train_3,"1.2,2.2,3.2,4.2"
```

**测试数据文件 (test.csv):**
- **列数**: 1列
- **第一列**: 向量数据（格式与训练数据相同）
- **编码**: UTF-8

示例：
```csv
"[1.0, 3.0, 11.0, 110.0, 62.0]"
"{2.0,4.0,6.0,8.0}"
"5.1,7.2,9.3,11.4"
```

#### 支持的向量格式

工具支持以下多种向量数据格式，会自动解析：

1. **方括号格式**：
   - `[1.0,2.0,3.0]` - 无空格
   - `[1.0, 2.0, 3.0]` - 有空格
   - `"[1.0,2.0,3.0]"` - 带外层引号

2. **花括号格式**：
   - `{1.0,2.0,3.0}` - 无空格
   - `{1.0, 2.0, 3.0}` - 有空格
   - `"{1.0,2.0,3.0}"` - 带外层引号

3. **纯数字格式**：
   - `1.0,2.0,3.0` - 逗号分隔
   - `1.0 2.0 3.0` - 空格分隔
   - `"1.0,2.0,3.0"` - 带引号

4. **混合格式**：
   - `'[1.0,2.0,3.0]'` - 单引号包围
   - `"  [1.0, 2.0, 3.0]  "` - 带前后空格

#### 数据类型支持

- **整数**: `[1,2,3,4]`
- **浮点数**: `[1.0,2.5,3.14,4.99]`
- **科学记数法**: `[1e-3,2.5e2,3.14e-1]`
- **混合数值**: `[1,2.5,3,4.99]`

#### 格式要求和限制

1. **维度一致性**: 所有向量必须具有相同的维度
2. **数值有效性**: 向量中的所有元素必须是有效的数值
3. **文件大小**: 
   - 建议单个CSV文件不超过2GB
   - 超大文件可能导致内存不足
4. **行数限制**: 
   - 训练数据：建议不超过1000万行
   - 测试数据：建议不超过10万行
5. **特殊字符**: 
   - 避免在向量数据中使用特殊字符
   - ID列可以包含字母、数字、下划线
6. **空值处理**: 不支持空值或NaN值，请确保数据完整

#### 常见错误和解决方案

❌ **错误格式示例**：
```csv
# 错误1：维度不一致
id_0,"[1.0,2.0,3.0]"
id_1,"[1.0,2.0]"

# 错误2：包含非数值
id_0,"[1.0,nan,3.0]"
id_1,"[1.0,null,3.0]"

# 错误3：格式混乱
id_0,"[1.0,2.0,3.0"
id_1,"1.0;2.0;3.0]"
```

✅ **正确格式示例**：
```csv
id_0,"[1.0,2.0,3.0]"
id_1,"[4.0,5.0,6.0]"
id_2,"{7.0,8.0,9.0}"
id_3,"10.0,11.0,12.0"
```

## 项目结构

```
ann-gen-data/
├── datasets.py         # 数据集生成和处理
├── csv_to_hdf5.py     # CSV转HDF5格式工具
├── hdf5_to_csv.py     # HDF5转CSV格式工具
├── hdf5_to_txt.py     # HDF5转TXT格式工具
├── bruteforce.py      # 暴力搜索算法实现
├── distance.py        # 距离度量函数
├── base.py           # 基础类定义
├── vectordb_bench/   # 向量数据库基准测试
├── requirements.txt  # 项目依赖
└── README.md        # 项目说明
```

## 参数说明

### CSV转HDF5工具参数

- `--train_csv`: 训练数据CSV文件路径（必需）
- `--test_csv`: 测试数据CSV文件路径（必需）
- `--output`: 输出HDF5文件路径（必需）
- `--distance`: 距离度量方式，可选值：`euclidean`（欧氏距离）、`angular`（余弦距离）、`hamming`（汉明距离），默认为`euclidean`
- `--k`: 最近邻数量，默认为100
- `--chunk_size`: 计算最近邻时的分块大小，默认为100000
- `--max_train_rows`: 最大训练数据行数，默认读取全部
- `--max_test_rows`: 最大测试数据行数，默认读取全部

### 距离度量说明

- **euclidean**: 欧氏距离，适用于一般的向量数据
- **angular**: 余弦距离（角度距离），适用于归一化向量或需要考虑向量方向的场景
- **hamming**: 汉明距离，适用于二进制向量数据

## 注意事项

### 基本要求

1. **内存管理**: 处理大规模数据集时，建议使用适当的分块大小（chunk_size）以避免内存溢出
2. **数据格式**: CSV文件中的向量数据支持多种格式，工具会自动解析
3. **维度一致性**: 训练数据和测试数据的向量维度必须一致
4. **存储空间**: 使用前请确保有足够的磁盘空间存储生成的HDF5文件
5. **计算时间**: 最近邻计算可能需要较长时间，特别是对于大规模数据集
6. **数据质量**: 确保CSV文件中的向量数据格式正确，避免解析错误

### CSV格式特别注意事项

1. **编码格式**: CSV文件必须使用UTF-8编码，避免中文乱码问题
2. **文件头**: CSV文件不需要表头，直接从数据行开始
3. **引号处理**: 
   - 如果向量字符串包含逗号，必须用引号包围
   - 推荐所有向量数据都用引号包围以避免解析错误
4. **换行符**: 支持Windows（\r\n）和Unix（\n）换行符
5. **空行**: 工具会自动跳过空行，但建议避免在文件中包含空行
6. **特殊字符**: 向量数据中避免使用制表符、换行符等特殊字符

### 性能优化建议

1. **文件大小**: 
   - 小文件（< 100MB）：chunk_size 设置为 10000-50000
   - 中等文件（100MB-1GB）：chunk_size 设置为 50000-100000  
   - 大文件（> 1GB）：chunk_size 设置为 100000-500000

2. **系统配置**:
   - 推荐至少 8GB 内存用于处理大规模数据集
   - 使用SSD存储可以显著提升I/O性能
   - 多核CPU可以提升最近邻计算速度

3. **数据预处理**:
   - 预先验证CSV文件格式的正确性
   - 对于超大数据集，考虑先转换小样本进行测试
   - 删除重复的向量数据以提升效率

## 许可证

MIT License