Milvus实战:构建Q&A系统及推荐系统

2023-12-20 03:46:26

Milvus简介

全民AI的时代已经在趋势之中,各类应用层出不穷,而想要构建一个完善的AI应用/系统,底层存储是不可缺少的一个组件。
与传统数据库或大数据存储不同的是,这种场景下则需要选择向量数据库,是专门用来存储和查询向量的数据库,其存储的向量来自于对文本、语音、图像、视频等的向量化数据,向量数据库不仅能够完成基本的CRUD(添加、读取查询、更新、删除)等操作,还能够对向量数据进行更快速的相似性搜索。

Milvus是众多向量库中的之一,适用于多个场景,如Questions & Answering系统、推荐系统等,单节点 Milvus 可以在秒内完成十亿级的向量搜索,分布式架构亦能满足用户的水平扩展需求。

参考文档:
Milvus官网
为AI而生的数据库:Milvus详解及实战

实践之问答系统&推荐系统Mix-in

元数据定义

不论是问答还是推荐,它们对上层暴露的接口仅仅是predict(...)/search(...)/query(...),模式是相同的,因此可以共用一个基本的Schema,固定基本的字段即可。

public class MilvusMeta {
    @Getter
    private final RecommenderSchema defaultMetricsSchema;
    @Getter
    private final QASchema defaultQASchema;

    public interface Schema {
        String getCollectionName();

        CreateCollectionParam getCreateCollectionParam();

        CreateIndexParam getCreateIndexParam();
    }


    @AllArgsConstructor
    public abstract static class BasicSchema implements Schema {
        public static final String SEARCH_PARAM = "{\"nprobe\":10}";    // Params
        public static final String INDEX_PARAM = "{\"nlist\":1024}";     // ExtraParam
        public static final IndexType INDEX_TYPE_DEFAULT = IndexType.IVF_FLAT;
        public static final MetricType METRIC_TYPE_DEFAULT = MetricType.L2;      // metric type
        public static final String INDEX_NAME_DEFAULT = "ivf_flat";
        public static final String PRIMARY_KEY_FIELD_NAME_DEFAULT = "id";
        public static final String PARTITION_KEY_FIELD_NAME_DEFAULT = "public";
        public static final String FIELD_NAME_DEFAULT = "embeddings";

        @Getter
        protected final CreateCollectionParam createCollectionParam;

        @Getter
        protected final CreateIndexParam createIndexParam;

        @Override
        public String getCollectionName() {
            return createCollectionParam.getCollectionName();
        }
    }

    /**
     * This schema is designed for storing Zen metrics.
     * TODO: Add more fields/features to describe a metric.
     */
    public static class RecommenderSchema extends BasicSchema {

        private RecommenderSchema(CreateCollectionParam collectionParam, CreateIndexParam indexParam) {
            super(collectionParam, indexParam);
        }

        public static RecommenderSchema create(ZenAiConfig.Storages.MilvusConf conf) {
            ZenAiConfig.Storages.Collection collection = conf.getActiveRecommenderCollection();
            return new RecommenderSchema(
                defaultCollectionParam(collection, collection.getEmbeddingsDimension()),
                MilvusUtil.createIndexParam(collection));
        }

        private static CreateCollectionParam defaultCollectionParam(ZenAiConfig.Storages.Collection collection,
            int dimension) {
            FieldType pkType = FieldType.newBuilder()
                .withName(collection.getPrimaryKey())
                .withDataType(DataType.VarChar)
                .withPrimaryKey(true)
                .withMaxLength(100)
                .withAutoID(false)
                .build();
            // 被embedding的字段
            FieldType embeddedFieldType = FieldType.newBuilder()
                .withName(collection.getEmbeddedFieldName())
                .withDataType(DataType.VarChar)
                .withMaxLength(255)
                .build();
            // embedding vector字段
            FieldType embeddingFieldType = FieldType.newBuilder()
                .withName(collection.getFieldName())
                .withDataType(DataType.FloatVector)
                .withDimension(dimension)
                .build();
            // 指定分区键字段,每一个Collection都需要指定一个分区键,除了能够Hive/Spark那样切分数据外,还能够加速相似查询。
            // 虽然Milvus支持多种方案以切分数据,但从管理复杂度、查询效率上来看,一个Collection对应多个数据分区,是最佳的方案。
            FieldType partitionKeyType = FieldType.newBuilder()
                .withName(collection.getPartitionKey())
                .withPartitionKey(true)
                .withDataType(DataType.VarChar)
                .withMaxLength(100)
                .build();
            return CreateCollectionParam.newBuilder()
                .withCollectionName(collection.getName())
                .withDescription(collection.getDescription())
                // .withShardsNum(2)
                .addFieldType(pkType)
                .addFieldType(embeddedFieldType)
                .addFieldType(embeddingFieldType)
                .addFieldType(partitionKeyType)
                // 开启动态字段添加功能
                .withEnableDynamicField(true)
                .build();
        }

    }

    public static class QASchema extends BasicSchema {
        public static final String ANSWER_FIELD_NAME = "answer";
        public static final String SCORE_FIELD_NAME = "score";
        public static final float SCORE_MAX_DEFAULT = 5.0f;
        public static final float SCORE_MIN_DEFAULT = 0.0f;

        public static final String INTENTION_FIELD_NAME = "intention";
        public static final String QUESTION_OCCURRENCE = "occurrence";

        public QASchema(CreateCollectionParam createCollectionParam, CreateIndexParam createIndexParam) {
            super(createCollectionParam, createIndexParam);
        }

        public static QASchema create(ZenAiConfig.Storages.MilvusConf conf) {
            ZenAiConfig.Storages.Collection collection = conf.getActiveQACollection();
            return new QASchema(
                defaultCollectionParam(collection, collection.getEmbeddingsDimension()),
                MilvusUtil.createIndexParam(collection));
        }

        private static CreateCollectionParam defaultCollectionParam(ZenAiConfig.Storages.Collection collection,
            int dimension) {
            FieldType pkType = FieldType.newBuilder()
                .withName(collection.getPrimaryKey())
                .withDataType(DataType.VarChar)
                .withPrimaryKey(true)
                .withMaxLength(100)
                .withAutoID(false)
                .build();
            FieldType embeddedFieldType = FieldType.newBuilder()
                .withName(collection.getEmbeddedFieldName())
                .withDataType(DataType.VarChar)
                .withMaxLength(65535)
                .build();
            FieldType embeddingFieldType = FieldType.newBuilder()
                .withName(collection.getFieldName())
                .withDataType(DataType.FloatVector)
                .withDimension(dimension)
                .build();
            FieldType partitionKeyType = FieldType.newBuilder()
                .withName(collection.getPartitionKey())
                .withPartitionKey(true)
                .withDataType(DataType.VarChar)
                .withMaxLength(100)
                .build();
            return CreateCollectionParam.newBuilder()
                .withCollectionName(collection.getName())
                .withDescription(collection.getDescription())
                // .withShardsNum(2)
                .addFieldType(pkType)
                .addFieldType(embeddedFieldType)
                .addFieldType(embeddingFieldType)
                .addFieldType(partitionKeyType)
                .withEnableDynamicField(true) // enable to insert new fields without modifying the code
                .build();
        }
    }

Milvus可行的操作接口定义

public interface IMilvusOperations {

    ZenAiConfig.Storages.Collection getCollection();

    MilvusConnection.MultiStatus delete(Filter filter);

    MilvusConnection.MultiStatus create(MilvusMeta.Index index);

    MilvusConnection.MultiStatus drop(String index);

    MilvusConnection.MultiStatus insert(MilvusData.Dataset dataset);

    MilvusConnection.MultiStatus insertAndFlush(MilvusData.Dataset dataset);
        /**
     * Query records by filter on the specified partition, which works like a normal SQL engine.
     *
     * @param partition which partition to query
     * @param filter boolean expression obeys the rules of Milvus
     * @param outputFields if empty, the result will contain all the fields, including the dynamic;
     *                     otherwise the result only contains the specified fields.
     * @return a nonnull instance, size of which is 0 if no matched records, otherwise is positive.
     */
    MilvusData.BasicPredictData queryByPartition(String partition, Filter filter, List<String> outputFields);

    List<MilvusData.BasicPredictData> search(List<List<Float>> vectors, Filter filter, int topK,
        List<String> outputFields);
 }

抽象系统接口定义

/**
 * 每个系统可能有不同的embedding的实现,因此需要定义一个接口。
 */
public interface IEmbedding {
    ImmutableList<List<Float>> getEmbeddings(List<String> messages);
}

/**
 * 通用接口定义,供应用层使用,可以基于sentence返回Milvus相似性结果集。
 */
public interface INlpSystem extends IMilvusOperations, IDataset, IEmbedding {

    default MilvusData.BasicPredictData predict(String sentence) {
        return predict(sentence, Filter.TRUE);
    }

    default MilvusData.BasicPredictData predict(String sentence, Filter filter) {
        return predict(sentence, filter, getCollection().getOutputFields());
    }

    default MilvusData.BasicPredictData predict(String sentence, Filter filter,
        List<String> outputFields) {
        ImmutableList<List<Float>> vectors = getEmbeddings(Lists.newArrayList(sentence));
        if (vectors.isEmpty()) {
            return MilvusData.BasicPredictData.EMPTY;
        }
        List<String> mergedOutputFields = Sets.union(
                ImmutableSet.copyOf(outputFields),
                ImmutableSet.copyOf(getCollection().getOutputFields()))
            .immutableCopy().asList();
        List<MilvusData.BasicPredictData> res = search(vectors, filter, getCollection().getTopk(), mergedOutputFields);
        return res.isEmpty() ? MilvusData.BasicPredictData.EMPTY : res.get(0);
    }

    default MilvusData.BasicPredictData predictByPartition(String partition, String sentence) {
        return predictByPartition(partition, sentence, Filter.TRUE, getCollection().getOutputFields());
    }

    default MilvusData.BasicPredictData predictByPartition(String partition, String sentence,
        Filter filter, List<String> outputFields) {
        ImmutableList<List<Float>> vectors = getEmbeddings(Lists.newArrayList(sentence));
        if (vectors.isEmpty()) {
            return MilvusData.BasicPredictData.EMPTY;
        }
        List<String> mergedOutputFields = Sets.union(
                ImmutableSet.copyOf(outputFields),
                ImmutableSet.copyOf(getCollection().getOutputFields()))
            .immutableCopy().asList();
        return searchByPartition(partition, vectors.get(0), filter, mergedOutputFields);
    }

}

/**
 * Q & A系统接口。
 */
public interface IQuestionAnswering extends INlpSystem {
}

/**
 * 推荐系统接口。
 */
public interface IRecommender extends INlpSystem, ISyncer {
}

插入数据集定义

列式格式构建插入Milvus的数据集,需要注意的是,Milvus JAVA SDK 2.3.1版本并不支持列式导致dynamic fields,因此我对源码进行了改造,以支持列式插入动态字段。
这个问题,已经反馈给了社区,并且已经在v2.3.2版本中支持。

public interface MilvusData {
    interface BasicData {
    	/**
    	 * Return a list view of the splitted data, to avoid copy.
    	 */
        BasicData[] split(int splitSize);

        /**
         * Return a view of the range [start, end) data, to avoid copy.
         */
        BasicData subData(int groupId, int start, int end);

        int size();
    }

    interface EmbeddingsProducer extends Function<List<String>, ImmutableList<List<Float>>> {
    }

    @Getter
    @Setter
    @AllArgsConstructor
    @NoArgsConstructor
    class Dataset {
        private List<BasicInsertData> inserts;
    }

    abstract class GroupedBasicData implements BasicData {
        @Getter
        private final int groupId;
        @Getter
        @Setter
        @Accessors(chain = true)
        private GroupedBasicData parent;

        protected GroupedBasicData(int groupId) {
            this.groupId = groupId;
        }

        /**
         * Split the data into more more sub-dataset.
         *
         * @param groups the number of expected groups
         * @return an array of sub-dataset views from the original dataset
         */
        public abstract BasicData[] grouped(int groups);
		
		/**
		 *  每一个切分或是extract的子数据集,都应该拥有一个可以唯一标识它的ID
		 */
        public String fullGroupId() {
            if (parent == null) {
                return String.valueOf(groupId);
            }
            return parent.fullGroupId() + "-" + groupId;
        }
    }

    @Getter
    abstract class PartitionedBasicData<T> extends GroupedBasicData {
        // 每一个系统都需要指定一个分区键,因此为了能够最小化存储,这里使用一个变量
        // 保存整个数据集应该插入
        private final T partition;

        protected PartitionedBasicData(T partition, int groupId) {
            super(groupId);
            this.partition = partition;
        }

        public abstract List<T> getPartitions();
    }
    /**
     * 以列式的形式构建插入数据集,并完成数据导入到Milvus。
     * Milvus
     */
    class BasicInsertData extends PartitionedBasicData<String> {
        private final String collection;
        private final ImmutableList<String> ids;
        private final ImmutableList<String> embeddingsInput;
        private final Supplier<ImmutableList<List<Float>>> vectorsSupplier;
        private final EmbeddingsProducer embeddingsProducer;
        private ImmutableMap<String, List<?>> dynamicFields;

        private final AtomicBoolean vectorsInitialized = new AtomicBoolean(false);
}

结果集定义

public interface MilvusData {
    /**
     * 一个通用的数据集,可以保存search/query的结果,行式数据结构。
     */
    @Getter
    class BasicPredictData extends GroupedBasicData {

        @Getter
        @Builder
        @AllArgsConstructor
        public static class Row {
            @JsonProperty
            private String id;
            private String embeddingsInput;
            @JsonProperty
            private Map<String, Object> extensions;
            @JsonProperty
            private float distance;
            @JsonIgnore
            private List<Float> vector;

            public <T> T getAs(String key, Class<T> clazz) {
                return getAs(key, clazz, null);
            }

            public <T> T getAs(String key, Class<T> clazz, T defaultValue) {
                return clazz.cast(extensions.getOrDefault(key, defaultValue));
            }
        }
    }
}

数据插入实例:列式插入

这个代码示例展示了如何构建列式数据集,并将其插入Milvus的流程。

注意到这里特别演示了使用了多线程并行 插入的功能,其原因有二:

  1. 一个批次的数据集过大,Milvus无法一次快速且稳定地完成插入动作,因此需要将原始数据集进行分组,例如这里分成3个组;
  2. 通常LLM(Large language Model)的一次API调用,只能支持生成16个向量数组,因此这里又对每一个分组后的子数据集进行横向切分,产生多个Batch,每个Batch包含一条记录。
@Test
void testSyncCollections() throws ExecutionException, InterruptedException {
        ExecutorService executorService = Executors.newFixedThreadPool(2);
        // 一组唯一值,用于区别每一条数据记录
        ImmutableList<String> ids = ImmutableList.of("1", "2", "3", "4", "5");
        // 一组指标名,这些指标就是待检索的合法指标集。
        ImmutableList<String> metrics = ImmutableList.of("m1", "m2", "m3", "m4", "m5");
        // 生成一组包含5个向量的列表,对应于每一个指标名
        ImmutableList<List<Float>> vectors = generateVectors(5, TEST_COLLECTION_DIMENSION);
        // 构建插入数据集
        MilvusData.BasicInsertData data = new MilvusData.BasicInsertData(
            config.getStorages().getMilvusConf().getActiveRecommenderCollection().getName(), ids, metrics,
            //这里使用Java中的Provider接口,提供执行插入数据任务时,对指标名列表向量化,由于这里事先生成了向量数组,因此直接从索引构建数据
            messages -> messages.stream().map(metrics::indexOf).map(vectors::get).collect(toImmutableList()));

        ImmutableList<Double> randoms = ImmutableList.of(1.0d, 2.0d, 3.0d, 4.0d, 5.0d);
        // 添加动态字段及相应的数据
        data.updateDynamicFields(ImmutableMap.of("random", randoms));
        // 构建并行插入数据任务
        // 3 groups:  ([1, 2]), ([3, 4]), ([5])
        // 1 batche: ([1],[2]),([3],[4]),([5])
        CompletableFuture<Integer>[] futures = milvusService.syncMetrics(data, 3, 1, executorService);
        assertEquals(3, futures.length);
        CompletableFuture.allOf(futures).join();
        assertEquals(2, futures[0].get());
        assertEquals(2, futures[1].get());
        assertEquals(1, futures[2].get());
}

相似性检索实例:指标推荐

用户输入一个指标(Metric)名,或是包含指标名的语句,可以通过Milvus的Search接口,找到最相近的TOP K指标,前提是需要对输入指标名进行向量化,然后以此向量来r从Milvus库中既存的指标集中计算找到最相似的。

    @Test
    void testLookingForMetric() {
        int topK = config..getMilvusConf().getActiveRecommenderCollection().getTopk();
        Optional<MilvusData.BasicPredictData> metrics = milvusService.getActiveRecommenderSys().map(system -> system.predict("销售总额"));
        assertTrue(metrics.isPresent());
        assertEquals(metrics.get().getRows().size(), topK);

        metrics = milvusService.getActiveRecommenderSys().map(system -> system.predict("销售总额", lt("random", 0f)));
        assertFalse(metrics.isPresent());
    }

相似性检索实例:问答

用户输入一段描述,可以通过Milvus提供的Search接口,找到历史相关的问题,并返回与此问题相关的上下文,并辅助回答AI模型回答用户的当前问题。

    @Test
    void testSearchWithSimilarityOfMultiVectors() {
        ImmutableList<String> testQuestions = ImmutableList.of("用柱形图展示2019年12月的总销售额", "用拆线图展示2020年12月的总净利润");
        ImmutableList<String> testIds = testQuestions.stream()
            .map(DefaultQuestionAnswering::encodeQuestion)
            .collect(ImmutableList.toImmutableList());
        IEmbedding embeddingSvc = aiService.getMilvusService().get().getActiveQuestionAnswering().get();
        DefaultQuestionAnswering.QAInsertData insertData = system.getInsertDataBuilder()
            .ids(testIds)
            .questions(testQuestions)
            .answers(ImmutableList.of("很好", "不错"))
            // 用户对于此问题返回结果的评价
            .scores(ImmutableList.of(1.0f, 1.0f))
            // 定义embeddings生成器,在插入时才会计算embeddings
            .embeddingsProducer(
                questions -> {
                    ImmutableList<List<Float>> qvectors = embeddingSvc.getEmbeddings(questions);
                    ImmutableList<List<Float>> mvectors = embeddingSvc.getEmbeddings(ImmutableList.of("总销售额", "总净利润"));
                    return ImmutableList.of(merge(qvectors.get(0), mvectors.get(0)), merge(qvectors.get(1), mvectors.get(1)));
                })
            .build();
        system.insert(new MilvusData.Dataset(ImmutableList.of(insertData)));

        // Case 1:
        // 用柱形图展示2019年12月的总销售额: 52.0181
        // 用拆线图展示2020年12月的总净利润: 147.0664
        verifySearch(system, "用拆线图展示2020年12月的总销售额", "总销售额", 0, "销售额", this::merge);

        // Case 2:
        // 用柱形图展示2019年12月的总销售额: 313.70105
        // 用拆线图展示2020年12月的总净利润: 181.3783
        verifySearch(system, "今年5月的净利润详情", "净利润", 0, "利润", this::merge);

        // Case 3:
        // 用柱形图展示2019年12月的总销售额: 160.30568
        // 用拆线图展示2020年12月的总净利润: 357.7008
        verifySearch(system, "今年5月的销售额详情", "销售额", 0, "销售额", this::merge);
    }

总结

Milvus对上层提供了与传统数据库相似的接口,以管理Milvus数据,同时提供了带有过滤功能的数据检索接口,使得上层应用能够很方便地利用传统数据库思维,来设计 和实现自己的系统。
但在使用中也感受到一些局限性或可能提升的点:

  1. 库中的一行记录只能对应一个embedding vector:只能使用相同模型生成的vector才能更好地检索向量,如果想一处持编码的文本对应多个vectors是不可能的,用户不得不创建新的Collection存储相同文本的不同向量。
  2. 用户显示Flush/Load Collection:每一次更新数据集,客户端必须要显示地load collection的操作,才能将新的数据加载到Server结点的内存中,同时第一次加载Collection必须是全量。
  3. 粗糙的表达式字符串:对于API接口的使用,缺少便利的表达式类定义,只能传递字符串,很容易出错,只能在运行时才知道哪些出错了。
  4. 缓存特性的支持:通常Milvus被用作Cache角色被引入系统中,但Milvus缺少一些缓存特性,如过期自动清理、partial dataset的load/unload功能等。

文章来源:https://blog.csdn.net/u014445499/article/details/135079187
本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。