向量資料庫落地實踐

京东云开发者發表於2024-04-03

本文基於京東內部向量資料庫vearch進行實踐。Vearch 是對大規模深度學習向量進行高效能相似搜尋的彈性分散式系統。詳見: https://github.com/vearch/zh_docs/blob/v3.3.X/docs/source/overview.rst

二、探索

初次認識向量資料庫,一臉懵逼?

image

向量是什麼?如何將文字轉換為向量?如何確定維度?如何定義表結構?如何選擇索引方式,建表引數如何配置?檢索引數如何配置?分片數副本數如何選擇等等

隨著對文件的逐漸熟悉以及和vearch相關同事的溝通,以上問題迎刃而解,具體的不再贅述。主要記住以下幾點:

1、 文字轉向量:採用大模型閘道器介面 domain/embeddings 傳入對應的模型如:text-embedding-ada-002-2和待轉換的文字即可;

2、 向量維度:這個和向量轉換所採用的模型有關,細節不用關注;

3、 建表引數的選擇以及表結構:主要在於retrieval_type 檢索模型的選擇,具體的可以參考文件。經過綜合考慮,決定採用 HNSW:

欄位標識 欄位含義 型別 是否必填 備註
metric_type 計算方式 string L2或者InnerProduct
nlinks 節點鄰居數量 int 預設32
efConstruction 構圖時尋找節點鄰居過程中在圖中遍歷的深度 int 預設40
"retrieval_type": "HNSW",
"retrieval_param": {
    "metric_type": "InnerProduct",
    "nlinks": 32,
    "efConstruction": 40
}

注意: 1、向量儲存只支援MemoryOnly
      2、建立索引不需要訓練,index_size 值大於0均可

具體的建表示例見後文。

4、 分片數和副本數結合實際資料量評估,如果無法評估,按照最少資源申請即可,後續可擴充套件。

三、實踐

1、 建表(space)

為了簡化操作,實行db(庫)-space(表)一對一的方案,弱化庫的概念。經過一系列探索之後定義出了通用的space結構:

{
    "name": "demphah",
    "partition_num": 3,
    "replica_num": 3,
    "engine": {
        "name": "gamma",
        "index_size": 1,
        "id_type": "String",
        "retrieval_type": "HNSW",
        "retrieval_param": {
            "metric_type": "InnerProduct",
            "nlinks": 32,
            "efConstruction": 100,
            "efSearch": 64
        }
    },
    "properties": {
        "vectorVal": {
            "type": "vector",
            "dimension": 1536
        },
        "contentVal": {
            "type": "string"
        },
        "chunkFlagId": {
            "type": "string",
            "index": true
        },
        "chunkIndexId": {
            "type": "integer",
            "index": true
        }
    }
}

欄位說明:

engine、partition_num等都是固定的引數,properties中所列欄位皆為通用欄位,如果有擴充套件欄位如:skuId,storeId追加即可

欄位名含義型別說明
vectorVal 文字向量 vector 維度與選用模型有關
contentVal 源文字 string
chunkFlagId 檔案唯一id string 檔案的標識id,用於串聯分塊後的片段
chunkIndexId 檔案分段位置 integer 從0開始,遞增
skuId ... 擴充套件欄位見上

這裡file的概念可以理解為一個單元,可能是一個檔案,也可能是一個url,總之就是一個資料整體。

2、 分段寫入

這裡針對通用檔案描述,比如提供一個pdf檔案如何匯入向量庫:

a. 首先上傳檔案到oss,然後根據對應的fileKey獲取到檔案資料流

b. 再根據各種拆分場景(按行、位元組數、正則拆分等)分成片段

c. 分段寫入向量庫:

    /**
     * 將字串轉換為向量並插入資料庫
     * <p>
     * 目前所有的知識庫管理端寫入全走這個方法
     *
     * @param dbName       資料庫名稱
     * @param spaceName    空間名稱
     * @param str          字串
     * @param flagId       標誌ID
     * @param chunkIndexId 塊索引ID
     * @param properties   屬性
     */
    private void embeddingsAndInsert(String dbName, String spaceName, String str, String flagId, Integer chunkIndexId, Map<String, Object> properties) {
        // 先向資料庫寫入一條記錄,記錄當前文件的寫入操作
        int success = knbaseDocRecordService.writeDocRecord(spaceName, flagId, chunkIndexId.longValue(), 0, str);
        if (success <= 0) {
            log.error("writeDocRecord失敗 {},{},{}", spaceName, flagId, chunkIndexId);
        }

        // 分塊轉向量並寫入
        List<Float> embeddings = GatewayUtil.baseEmbeddings(str);
        if (CollectionUtils.isEmpty(embeddings)) {
            return;
        }

        KnBaseVecDto knBaseVecDto = buildKnBaseVecDto(new FeaVector(embeddings),flagId,chunkIndexId,str);

        Map<String, Object> newPros = JsonUtil.obj2Map(knBaseVecDto);
        if (MapUtils.isNotEmpty(properties)) {
            newPros.putAll(properties);
        }
        // {"_index":"kn_base_file_db","_type":"kn_base_file_space","_id":"-8182839813441244911","status":200}
        String insert = VearchUtil.insert(dbName, spaceName, null, newPros);
        if (StringUtils.isBlank(insert) || !insert.contains("_index")) {
            log.error("寫入失敗的塊:{},{}", chunkIndexId, insert);
        }
    }

3、 資料記錄

上文寫知識庫的過程有個 knbaseDocRecordService.writeDocRecord 的邏輯,用於記錄寫入的片段。下文詳細介紹其中用到的mysql表:

1、 表1 space記錄表

注:主要用於記錄建立的space,以及查詢管控,如禁用某個space等

CREATE TABLE `xxx_vearch_spaces` (
  `id` bigint(20) unsigned NOT NULL AUTO_INCREMENT COMMENT '自增主鍵',
  `type` tinyint(3) NOT NULL COMMENT '型別',
  `status` tinyint(3) NOT NULL COMMENT '狀態',
  `space` varchar(127) NOT NULL COMMENT '空間標識',
  `db` varchar(127) NOT NULL COMMENT '庫標識',
  `desc` varchar(127) NOT NULL COMMENT '空間描述',
  `ext` varchar(4095) NOT NULL DEFAULT '' COMMENT '擴充套件欄位',
  `creator` varchar(127) NOT NULL DEFAULT '' COMMENT '建立人',
  `created` timestamp NOT NULL COMMENT '建立時間',
  `modifier` varchar(127) NOT NULL DEFAULT '' COMMENT '修改人',
  `modified` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '修改時間',
  `deleted` tinyint(3) NOT NULL DEFAULT '0' COMMENT '已刪除(0:否;1:是)',
  PRIMARY KEY (`id`) USING BTREE,
  UNIQUE KEY `uniq_space_db` (`space`,`db`) USING BTREE
) ENGINE=InnoDB AUTO_INCREMENT=0 DEFAULT CHARSET=utf8mb4 COMMENT='xxx向量空間'

2、 表2 file記錄表

注:主要用於記錄space下的file,以及查詢管控,如禁用某個file,以及關聯查詢對應的全部片段。

CREATE TABLE `xxx_spaces_knbase` (
  `id` bigint(20) unsigned NOT NULL AUTO_INCREMENT COMMENT '自增主鍵',
  `status` tinyint(3) NOT NULL COMMENT '狀態',
  `space` varchar(127) NOT NULL COMMENT '空間標識',
  `file_name` varchar(255) NOT NULL COMMENT '檔名',
  `file_desc` varchar(511) NOT NULL COMMENT '檔案描述',
  `byte_num` bigint(20) unsigned NOT NULL COMMENT '字元數',
  `hit_count` int(10) unsigned NOT NULL COMMENT '命中次數',
  `ext` varchar(4095) NOT NULL DEFAULT '' COMMENT '擴充套件欄位',
  `creator` varchar(127) NOT NULL DEFAULT '' COMMENT '建立人',
  `created` timestamp NOT NULL COMMENT '建立時間',
  `modifier` varchar(127) NOT NULL DEFAULT '' COMMENT '修改人',
  `modified` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '修改時間',
  `deleted` tinyint(3) NOT NULL DEFAULT '0' COMMENT '已刪除(0:否;1:是)',
  `file_flag_id` varchar(255) DEFAULT NULL COMMENT '檔案唯一標識',
  PRIMARY KEY (`id`) USING BTREE,
  KEY `idx_space` (`space`) USING BTREE
) ENGINE=InnoDB AUTO_INCREMENT=0 DEFAULT CHARSET=utf8mb4 COMMENT='xxx空間知識庫'

3、 表3 paragraph記錄表

注:主要用於記錄file拆分的片段,包括當前位置,查詢命中數等。

CREATE TABLE `xxx_knbase_doc_record` (
  `id` bigint(20) unsigned NOT NULL AUTO_INCREMENT COMMENT '自增主鍵',
  `space` varchar(127) NOT NULL COMMENT '空間標識',
  `file_flag_id` varchar(255) NOT NULL COMMENT '檔案標識',
  `d_index` bigint(20) unsigned NOT NULL COMMENT '文件位置',
  `hit_count` int(10) unsigned NOT NULL COMMENT '命中次數',
  `ext` varchar(4095) NOT NULL DEFAULT '' COMMENT '擴充套件欄位',
  `creator` varchar(127) NOT NULL DEFAULT '' COMMENT '建立人',
  `created` timestamp NOT NULL COMMENT '建立時間',
  `modifier` varchar(127) NOT NULL DEFAULT '' COMMENT '修改人',
  `modified` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '修改時間',
  `deleted` tinyint(3) NOT NULL DEFAULT '0' COMMENT '已刪除(0:否;1:是)',
  PRIMARY KEY (`id`) USING BTREE,
  UNIQUE KEY `uniq_space_file_idx` (`space`,`file_flag_id`,`d_index`) USING BTREE
) ENGINE=InnoDB AUTO_INCREMENT=0 DEFAULT CHARSET=utf8mb4 COMMENT='xxx知識文件記錄'

四、總結

向量資料庫對於大模型應用落地來說至關重要,有些不可外露的內部資料可以儲存在向量庫中,用於內部檢索。隨著向量庫中資料的豐富,大模型推理回答的能力也將更加精準。

上文的設計比如space中的chunkFlagId可以關聯出原始的整個檔案;chunkIndexId可以控制資料的查詢範圍,另一方面可以透過此欄位實現分頁(vearch目前不支援分頁查詢)以及全文匯出。xxx_knbase_doc_record表中記錄了片段的記錄,可用於計算片段的chunkIndexId,一方面避免重複,另一方面保證屬性的遞增,可用於擴充套件很多能力。

目前向量資料庫的檢索只支援基本的向量檢索和關鍵字檢索,後續會逐步最佳化混合檢索等方案以提高檢索準確率等。

相關文章