使用向量数据库进行电影相似性搜索

简介

随着 LLM 的普及,我们经常看到它们被用于与文本生成不直接相关的任务。例如,使用 LLM 来构建推荐系统。在这篇文章中,我们将展示如何使用 Quarkus LangChain4j 构建这样一个系统,但无需使用 LLM。更具体地说,我们将使用向量数据库创建一个简单的电影相似性搜索系统。在这个过程中,Quarkus LangChain4j 的作用是通过 `EmbeddingStore` 接口来抽象底层的向量数据库。

最近,一个新的相关示例已添加到 Quarkus LangChain4j 示例中。

嵌入

嵌入是一种将非结构化数据(文本、图像等)结构化的方式。这通过将数据映射到向量来实现。由于我们可以对向量执行数学运算,例如计算它们之间的距离,因此我们可以使用嵌入来计算两块数据的接近程度(或相似度)。在我们的案例中,比较电影概述的嵌入可以让我们衡量两部电影的相似程度。这就是本文的前提。

如何创建嵌入?

Quarkus LangChain4j 提供了以下几种创建嵌入的方法:

在这篇文章中,我们将使用前者。我有没有提到我们不会使用 LLM?相反,我们将使用 pgvector,这是一个 PostgreSQL 扩展,提供向量操作和索引。

我们的项目需要以下依赖项,用于 pgvector进程内嵌入

    <dependency>
        <groupId>dev.langchain4j</groupId>
        <artifactId>langchain4j-embeddings-bge-small-en-q</artifactId>
    </dependency>
    <dependency>
        <groupId>io.quarkiverse.langchain4j</groupId>
        <artifactId>quarkus-langchain4j-pgvector</artifactId>
    </dependency>

为了能够使用这些依赖项而无需指定版本,可以将 BOM 导入项目的 `dependencyManagement` 中。

<dependency>
    <groupId>io.quarkiverse.langchain4j</groupId>
    <artifactId>quarkus-langchain4j-bom</artifactId>
    <version>0.25.0</version>
    <type>pom</type>
    <scope>import</scope>
</dependency>

为了正确使用进程内嵌入模型,我们需要在 `application.properties` 文件中对其进行配置。我们还需要配置 pgvector 的维度,并确保它与嵌入模型的维度保持一致。在我们的案例中,它是 384(Quarkus LangChain4j 文档提供了每个模型使用的尺寸)。

因此,`application.properties` 文件应如下所示:

quarkus.langchain4j.pgvector.dimension=384
quarkus.langchain4j.embedding-model.provider=dev.langchain4j.model.embedding.onnx.bgesmallenq.BgeSmallEnQuantizedEmbeddingModel

注意:我们可以使用 Quarkus LangChain4j 支持的任何其他文档存储,事实上,这是使用它的关键优势之一,即对嵌入存储的抽象。

存储嵌入

为了存储嵌入,我们需要一个 `EmbeddingStoreIngestor`。Ingestor 是使用 `EmbeddingModel` 和 `EmbeddingStore` 创建的,两者都作为 bean 提供,并且可以轻松注入。或者,用户可以指定一个文档拆分器,用于将大型文档拆分成较小的块,但在这篇文章中我们不需要它,因为电影概述相对较小。

ingestion

在将非结构化文本传递给 ingestor 之前,我们需要将其包装在 `Document` 对象中。`Document` 还包含一个 `Metadata` 对象,其中保存键值对。`Metadata` 非常有用,因为我们可以添加需要将 `Document` 与其他数据相关联的信息。在我们的案例中,我们将使用 `Metadata` 来存储电影的 ID。该 ID 将帮助我们关联电影概述和实际电影。

下面的简化代码展示了实际的 ingestor 是如何创建的,以及嵌入是如何存储的。

@Inject
EmbeddingModel embeddingModel;

@Inject
EmbeddingStore embeddingStore;

public EmbeddingStoreIngestor createIngestor()  {
    return EmbeddingStoreIngestor.builder()
    .embeddingModel(embeddingModel)
    .embeddingStore(embeddingStore)
    .build();
}

public void ingest(Long movieId, String overview) {
      Metadata metadata = Metadata.from(Map.of("id", id));
      Document document = Document.from(overview, metadata);
      createIngestor().ingest(document);
}

那么,我们到底如何使用电影 ID 呢?这实际上取决于我们如何存储其余的电影数据。在我们的案例中,我们将电影数据存储在 PostgreSQL 数据库中。这意味着电影 ID 对应于数据库中电影的 ID。

查询嵌入

为了查询嵌入,我们将使用 `EmbeddingStore` 和 `EmbeddingModel` 来执行 `EmbeddingSearchRequest`。代码非常直接。我们使用电影概述来创建搜索请求。`EmbeddingSearchRequest` 构建器还允许我们指定最大结果数和最小相似度阈值。后者允许我们过滤掉与查询嵌入不够相似的嵌入。换句话说,它告诉存储,如果没有足够相似的结果,请避免在响应中添加不相关的结果。

package io.quarkiverse.langchain4j.sample;

import java.util.List;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingStore;
import jakarta.enterprise.context.ApplicationScoped;
import jakarta.inject.Inject;
import jakarta.transaction.Transactional;

@ApplicationScoped
public class MovieRecommendationService {

  @Inject
  EmbeddingStore<TextSegment> embeddingStore;

  @Inject
  EmbeddingModel embeddingModel;

  @Transactional
  public List<Movie> searchSimilarMovies(String overview) {

    Embedding embedding = embeddingModel.embed(overview).content();
    EmbeddingSearchRequest request = EmbeddingSearchRequest.builder()
    .queryEmbedding(embedding)
    .minScore(0.5)
    .maxResults(10)
    .build();

    return embeddingStore.search(request).matches().stream().map(m -> {
      Long id = m.embedded().metadata().getLong("id");
      Movie movie = Movie.findById(id);
      return movie;
    }).toList();

  }
}

加载电影

为了将电影填充到数据库中,我们将使用一个 CSV 文件,其中包含 IMDB 排名前 1000 的电影。对我们而言,重要的列是:

  • `title`:电影标题

  • `overview`:电影概述

  • `link`:海报图片的链接

为了方便地将 CSV 条目映射到 `Movie` 对象,我们将使用 `Jackson`。具体来说,我们将使用 `@JsonProperty` 注释将 CSV 列映射到 `Movie` 字段。此外,我们将使用 `@JsonIgnoreProperties(ignoreUnknown = true)` 来忽略未知字段。

因此,我们的 Movie 实体对象的简化版本如下所示:

package org.acme;

import io.quarkus.hibernate.orm.panache.PanacheEntity;
import jakarta.persistence.Column;
import jakarta.persistence.Entity;
import java.util.List;

import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.fasterxml.jackson.annotation.JsonProperty;


@Entity
@JsonIgnoreProperties(ignoreUnknown = true)
public class Movie extends PanacheEntity {

    @JsonProperty("Poster_Link")
    public String link;
    @JsonProperty("Series_Title")
    public String title;
    @JsonProperty("Overview")
    @Column(length = 1000)
    public String overview;

    public static List<Movie> searchByTitleLike(String title) {
      return find("title like ?1", "%" + title + "%").list();
    }
}

现在,我们准备好将电影从 CSV 加载到我们的关系数据库和向量数据库中了。

要点

  • 观察 `StartupEvent` 允许我们在应用程序启动时加载电影。

  • CSVMapper 用于将 CSV 条目映射到 `Movie` 对象。

  • 我们使用 `@Transaction` 方法保存每部电影,因为我们需要通常由数据库生成的 `id`。

  • 我们分批摄取文档。这可以产生巨大的差异,尤其是在我们从进程内嵌入模型迁移到远程嵌入模型时。

package org.acme;

import io.quarkus.logging.Log;
import io.quarkus.runtime.StartupEvent;
import jakarta.enterprise.context.ApplicationScoped;
import jakarta.enterprise.event.Observes;
import jakarta.transaction.Transactional;

import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import org.eclipse.microprofile.config.inject.ConfigProperty;

import com.fasterxml.jackson.databind.MappingIterator;
import com.fasterxml.jackson.dataformat.csv.CsvMapper;
import com.fasterxml.jackson.dataformat.csv.CsvSchema;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;

@ApplicationScoped
public class MovieLoader {


  public void load(@Observes StartupEvent event, @ConfigProperty(name = "movies.file") Path moviesFile,
                   EmbeddingStore embeddingStore, EmbeddingModel embeddingModel) throws Exception {
    if (!Files.exists(moviesFile)) {
      throw new IllegalStateException("Missing movies file: " + moviesFile);
    }

    embeddingStore.removeAll();

    EmbeddingStoreIngestor ingester = EmbeddingStoreIngestor.builder()
    .embeddingModel(embeddingModel)
    .embeddingStore(embeddingStore)
    .build();

    List<Document> docs = new ArrayList<>();
    try (MappingIterator<Movie> it = new CsvMapper().readerFor(Movie.class).with(CsvSchema.emptySchema().withHeader()).readValues(moviesFile.toFile())) {
      for (Movie movie : it.readAll()) {
        Long id = save(movie).id;
        Metadata metadata = Metadata.from(Map.of("id", id, "title", movie.title));
      Document document = Document.from(movie.overview, metadata);
      docs.add(document);

      }
    }

    Log.info("Ingesting movies...");
    ingester.ingest(docs);
    Log.info("Application initalized!");
  }

  @Transactional
  public Movie save(Movie m) {
    m.persist();
    return m;
  }

}

要使用 CSV 映射器,我们需要 Jackson 的 CSV 数据格式依赖项。

        <dependency>
            <groupId>com.fasterxml.jackson.dataformat</groupId>
            <artifactId>jackson-dataformat-csv</artifactId>
        </dependency>

电影文件的路径在 `application.properties` 文件中指定。

movies.file=src/main/resources/movies.csv

整合

剩下的就是创建一个 REST 端点,允许我们搜索相似的电影。我们也可以使用一个简单的 UI。让我们先从 REST 端点开始。这很简单。我们需要两个方法,一个用于搜索电影,一个用于搜索相似电影。前者我们只使用 `Movie` 实体,后者我们注入并使用我们之前创建的 `MovieRecommendationService`。

package io.quarkiverse.langchain4j.sample;

import jakarta.inject.Inject;
import jakarta.ws.rs.GET;
import jakarta.ws.rs.Path;
import jakarta.ws.rs.Produces;
import jakarta.ws.rs.core.MediaType;
import java.util.List;

@Path("/movies")
public class MovieResource {

    @Inject
    MovieRecommendationService recommendationService;

    @GET
    @Produces(MediaType.APPLICATION_JSON)
    @Path("/by-title/{title}")
    public List<Movie> searchByTitle(String title) {
      return Movie.searchByTitleLike(title);
    }

    @GET
    @Produces(MediaType.APPLICATION_JSON)
    @Path("/similar/{id}")
    public List<Movie> searchSimilar(Long id) {
      Movie m = Movie.findById(id);
      return recommendationService.searchSimilarMovies(m.overview);
    }
}

对于 UI,我们只需要使用一个简单的 HTML 页面,该页面使用 REST 端点来搜索相似的电影。

该页面的关键元素是:

  • movie-box:用于输入电影标题的文本字段

  • search-results:用于显示搜索结果的无序列表

  • movie-overview:用于显示所选电影概述的 div

  • movie-poster:用于显示电影海报的图像

  • similar-results:用于显示相似电影的额外无序列表

重要的是要记住,`Movie` 实体使用 Jackson 将 CSV 列映射到实体字段。这意味着当 `Movie` 序列化为 JSON 时,它将使用 CSV 列名作为字段名,而不是实体字段名。下面的 HTML 代码需要考虑这一点。

<!DOCTYPE html>
<html lang="en">
	<head>
    <title>{{_title_}}</title>
		<meta charset="UTF-8"/>
		<meta name="viewport" content="width=device-width, initial-scale=1"/>
	</head>
	<body>
    <h2>Movie Similarity Search</h2>
    <input type="text" id="movie-box" placeholder="Enter a movie title">
    <h3 hidden="true" id="movie-results-heading">Click on of the movies below</h3>
    <ul id="search-results"></ul>
    <img id="movie-poster"><img>
    <div id="movie-overview"></div>
    <h3 hidden="true" id="similar-heading">Similar movies</h3>
    <ul id="similar-results"></ul>

    <script>
    document.getElementById("movie-box").addEventListener("input", async function() {
        const query = this.value.trim();
        if (query.length === 0) {
          document.getElementById("search-results").innerHTML = "";
          return;
        }
      const response = await fetch(`/movies/by-title/${encodeURIComponent(query)}`);
        const movies = await response.json();

        if (movies.length > 0) {
          document.getElementById("movie-results-heading").hidden = false;
        }
        movies.forEach(movie => {
          const li = document.createElement("li");
          li.textContent = movie.Series_Title;
          li.addEventListener("click", () => displayMovie(movie));
          document.getElementById("search-results").appendChild(li);
        });
    });

    async function displayMovie(movie) {
      console.log('Displaying movie:', movie);
      document.getElementById("search-results").innerHTML = "";
      document.getElementById("movie-poster").src = movie.Poster_Link;
      document.getElementById("movie-poster").style.display = "block";
      document.getElementById("movie-overview").textContent = movie.Overview;
      document.getElementById("similar-heading").hidden = false;
      document.getElementById("movie-results-heading").hidden = true;
      document.getElementById("similar-results").innerHTML = "";

      const response = await fetch(`/movies/similar/${encodeURIComponent(movie.id)}`);
      const similarMovies = await response.json();
        similarMovies.forEach(similarMovie => {
           const li = document.createElement("li");
           li.textContent = similarMovie.Series_Title;
           li.addEventListener("click", () => displayMovie(similarMovie));
           document.getElementById("similar-results").appendChild(li);
        });
    }
    </script>
	</body>
</html>

我不会详细介绍 HTML 代码,因为它超出了本文的范围。最终结果应如下所示:

movie similarity search ui

结论

推荐的质量取决于电影概述的准确性以及嵌入的质量。这意味着更好的嵌入模型可以带来更好的推荐。所使用的向量数据库也可能影响推荐的质量,但这将是另一篇文章的主题。