使用Jib构建内嵌Hugging Face模型与SQLite向量存储的独立搜索服务


团队内部一个监控看板需要一个简单的日志搜索功能。需求很明确:能在数万条结构化日志中,根据自然语言描述找到相关的异常信息。常规的grep或者LIKE查询效果差强人意,而引入Elasticsearch或一个专门的向量数据库,对于这个非核心的内部工具来说,无论在部署还是维护上都显得过于笨重。我们的目标是构建一个完全自包含、零外部依赖、通过一个容器镜像就能启动的服务。

这个约束条件把技术选型推向了一个有趣的方向:将语言模型、数据存储和应用程序逻辑全部打包进一个独立的单元。最终的方案定格在:Java作为主力语言,利用Hugging Face Transformers生态中的预训练模型进行语义向量化,SQLite作为数据存储的基石,最后使用Jib将这一切无缝打包成一个轻量级Docker镜像,甚至无需本地安装Docker守护进程。

挑战的核心在于,SQLite本身并不支持向量索引和相似度计算。这意味着我们需要在应用层实现向量搜索逻辑,并巧妙地将其与SQLite传统的全文搜索能力(FTS5)结合,构建一个混合搜索引擎。

技术栈基石:依赖与配置

一个项目的起点是它的依赖定义。在pom.xml中,我们需要引入几个关键组件:

  1. DJL (Deep Java Library): 一个对Java开发者友好的深度学习框架,它极大地简化了从Hugging Face Hub加载和运行PyTorch或TensorFlow模型的过程。
  2. SQLite JDBC Driver: 连接和操作SQLite数据库的标准驱动。
  3. Jib Maven Plugin: Google出品的容器化工具,可以直接从Maven构建过程中生成Docker镜像,无需Dockerfile
<properties>
    <java.version>17</java.version>
    <djl.version>0.26.0</djl.version>
    <sqlite.jdbc.version>3.43.0.0</sqlite.jdbc.version>
    <slf4j.version>2.0.9</slf4j.version>
</properties>

<dependencies>
    <!-- DJL for Hugging Face model interaction -->
    <dependency>
        <groupId>ai.djl</groupId>
        <artifactId>api</artifactId>
        <version>${djl.version}</version>
    </dependency>
    <dependency>
        <groupId>ai.djl.huggingface</groupId>
        <artifactId>huggingface-tokenizers</artifactId>
        <version>${djl.version}</version>
    </dependency>
    <!-- PyTorch engine for DJL -->
    <dependency>
        <groupId>ai.djl.pytorch</groupId>
        <artifactId>pytorch-engine</artifactId>
        <version>${djl.version}</version>
    </dependency>
    <dependency>
        <groupId>ai.djl.pytorch</groupId>
        <artifactId>pytorch-jni</artifactId>
        <version>2.1.2-0.26.0</version>
        <scope>runtime</scope>
    </dependency>

    <!-- SQLite driver -->
    <dependency>
        <groupId>org.xerial</groupId>
        <artifactId>sqlite-jdbc</artifactId>
        <version>${sqlite.jdbc.version}</version>
    </dependency>

    <!-- Logging -->
    <dependency>
        <groupId>org.slf4j</groupId>
        <artifactId>slf4j-api</artifactId>
        <version>${slf4j.version}</version>
    </dependency>
    <dependency>
        <groupId>org.slf4j</groupId>
        <artifactId>slf4j-simple</artifactId>
        <version>${slf4j.version}</version>
    </dependency>
</dependencies>

<build>
    <plugins>
        <plugin>
            <groupId>com.google.cloud.tools</groupId>
            <artifactId>jib-maven-plugin</artifactId>
            <version>3.4.0</version>
            <configuration>
                <from>
                    <!-- Using a minimal base image with glibc for native libraries -->
                    <image>gcr.io/distroless/java17-debian12</image>
                </from>
                <to>
                    <image>my-registry/hybrid-search-service:0.0.1</image>
                </to>
                <container>
                    <mainClass>com.example.search.Application</mainClass>
                    <ports>
                        <port>8080</port>
                    </ports>
                    <!-- This is crucial for persisting the SQLite database -->
                    <volumes>
                        <volume>/app/data</volume>
                    </volumes>
                    <!-- Set DJL cache dir to a known location inside the container -->
                    <environment>
                        <DJL_CACHE_DIR>/app/model-cache</DJL_CACHE_DIR>
                    </environment>
                </container>
                <extraDirectories>
                    <paths>
                        <!-- 
                        This copies a pre-downloaded model into the container image.
                        This avoids downloading the model on container startup.
                        You must first run the app locally to populate ~/.djl.ai/
                        -->
                        <path>
                            <from>${user.home}/.djl.ai/cache</from>
                            <into>/app/model-cache</into>
                        </path>
                    </paths>
                </extraDirectories>
            </configuration>
        </plugin>
    </plugins>
</build>

Jib的配置是这里的关键。我们指定了一个最小化的基础镜像,并将预先下载好的Hugging Face模型(通常位于~/.djl.ai/cache)直接打包到镜像的/app/model-cache目录。通过环境变量DJL_CACHE_DIR告知DJL在此路径下寻找模型,避免了服务每次启动时都去网络下载,保证了启动速度和离线环境下的可用性。同时,/app/data被声明为卷,用于持久化SQLite数据库文件。

数据存储层:在SQLite中安放向量

SQLite没有VECTOR类型,但我们可以用BLOB类型来存储浮点数向量的原始字节表示。同时,为了结合关键词搜索,我们引入了SQLite的FTS5扩展。

SqliteVectorStore类负责数据库的初始化和所有CRUD操作。

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.sql.*;
import java.util.ArrayList;
import java.util.List;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class SqliteVectorStore {
    private static final Logger logger = LoggerFactory.getLogger(SqliteVectorStore.class);
    private final String dbUrl;

    public SqliteVectorStore(String dbPath) {
        this.dbUrl = "jdbc:sqlite:" + dbPath;
        initializeDatabase();
    }

    private void initializeDatabase() {
        String createDocumentsTable = """
            CREATE TABLE IF NOT EXISTS documents (
                id INTEGER PRIMARY KEY,
                content TEXT NOT NULL,
                embedding BLOB NOT NULL,
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP
            );
            """;
        
        // FTS5 virtual table for keyword search
        String createFtsTable = """
            CREATE VIRTUAL TABLE IF NOT EXISTS documents_fts USING fts5(
                content,
                content='documents',
                content_rowid='id'
            );
            """;
        
        // Trigger to keep FTS table in sync with documents table
        String createFtsTriggers = """
            CREATE TRIGGER IF NOT EXISTS documents_ai AFTER INSERT ON documents BEGIN
              INSERT INTO documents_fts(rowid, content) VALUES (new.id, new.content);
            END;
            CREATE TRIGGER IF NOT EXISTS documents_ad AFTER DELETE ON documents BEGIN
              INSERT INTO documents_fts(documents_fts, rowid, content) VALUES ('delete', old.id, old.content);
            END;
            CREATE TRIGGER IF NOT EXISTS documents_au AFTER UPDATE ON documents BEGIN
              INSERT INTO documents_fts(documents_fts, rowid, content) VALUES ('delete', old.id, old.content);
              INSERT INTO documents_fts(rowid, content) VALUES (new.id, new.content);
            END;
            """;
        
        try (Connection conn = DriverManager.getConnection(dbUrl);
             Statement stmt = conn.createStatement()) {
            stmt.execute(createDocumentsTable);
            stmt.execute(createFtsTable);
            stmt.execute(createFtsTriggers);
            logger.info("Database initialized successfully.");
        } catch (SQLException e) {
            logger.error("Database initialization failed.", e);
            throw new RuntimeException(e);
        }
    }

    public void insertDocument(String content, float[] embedding) {
        String sql = "INSERT INTO documents(content, embedding) VALUES(?, ?)";
        try (Connection conn = DriverManager.getConnection(dbUrl);
             PreparedStatement pstmt = conn.prepareStatement(sql)) {
            pstmt.setString(1, content);
            pstmt.setBytes(2, floatArrayToBytes(embedding));
            pstmt.executeUpdate();
        } catch (SQLException | IOException e) {
            logger.error("Failed to insert document: {}", content, e);
        }
    }

    // This method performs a full scan to retrieve all vectors.
    // In a real high-performance system, this is the bottleneck.
    public List<VectorRecord> getAllVectors() {
        String sql = "SELECT id, embedding FROM documents";
        List<VectorRecord> records = new ArrayList<>();
        try (Connection conn = DriverManager.getConnection(dbUrl);
             Statement stmt = conn.createStatement();
             ResultSet rs = stmt.executeQuery(sql)) {

            while (rs.next()) {
                records.add(new VectorRecord(
                    rs.getInt("id"),
                    bytesToFloatArray(rs.getBytes("embedding"))
                ));
            }
        } catch (SQLException | IOException e) {
            logger.error("Failed to retrieve vectors.", e);
        }
        return records;
    }
    
    // Helper methods for serialization/deserialization
    private byte[] floatArrayToBytes(float[] array) throws IOException {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        DataOutputStream dos = new DataOutputStream(baos);
        for (float f : array) {
            dos.writeFloat(f);
        }
        return baos.toByteArray();
    }

    private float[] bytesToFloatArray(byte[] bytes) throws IOException {
        ByteArrayInputStream bais = new ByteArrayInputStream(bytes);
        DataInputStream dis = new DataInputStream(bais);
        float[] floatArray = new float[bytes.length / 4];
        for (int i = 0; i < floatArray.length; i++) {
            floatArray[i] = dis.readFloat();
        }
        return floatArray;
    }

    public record VectorRecord(int id, float[] vector) {}
}

这里的核心是floatArrayToBytesbytesToFloatArray两个辅助方法,它们负责在Java的float[]和SQLite的BLOB之间进行转换。数据库初始化逻辑创建了主表documents和FTS5虚拟表documents_fts,并设置了触发器,确保两者数据同步。

语义核心:加载Hugging Face模型

EmbeddingService是与AI模型交互的唯一入口。它负责加载模型并提供一个简单的接口来将文本转换为向量。我们选择了一个轻量级且高效的中文句向量模型,如bge-small-zh-v1.5

import ai.djl.Application;
import ai.djl.inference.Predictor;
import ai.djl.modality.nlp.qa.QAInput;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ZooModel;
import ai.djl.sentencesimilarity.SentenceSimilarity;
import ai.djl.training.util.ProgressBar;
import ai.djl.translate.TranslateException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;

public class EmbeddingService implements AutoCloseable {

    private static final Logger logger = LoggerFactory.getLogger(EmbeddingService.class);
    private final ZooModel<String, float[]> model;
    private final Predictor<String, float[]> predictor;

    public EmbeddingService() {
        try {
            String modelUrl = "djl://ai.djl.huggingface.pytorch/sentence-transformers/bge-small-zh-v1.5";
            
            Criteria<String, float[]> criteria = Criteria.builder()
                .optApplication(Application.NLP.SENTENCE_SIMILARITY)
                .setTypes(String.class, float[].class)
                .optModelUrls(modelUrl)
                .optEngine("PyTorch")
                .optProgress(new ProgressBar())
                .build();
            
            this.model = criteria.loadModel();
            this.predictor = model.newPredictor();
            logger.info("Embedding model loaded successfully.");
        } catch (Exception e) {
            logger.error("Failed to load embedding model.", e);
            throw new IllegalStateException("Could not initialize EmbeddingService", e);
        }
    }

    public float[] generateEmbedding(String text) {
        try {
            return predictor.predict(text);
        } catch (TranslateException e) {
            logger.error("Failed to generate embedding for text: {}", text, e);
            return new float[0]; // Return empty array on failure
        }
    }

    @Override
    public void close() {
        predictor.close();
        model.close();
        logger.info("EmbeddingService resources released.");
    }
}

这段代码的健壮性体现在构造函数中完整的异常处理。如果模型加载失败,服务将无法启动,这是一种快速失败的设计。AutoCloseable接口的实现确保了模型资源在应用关闭时能被正确释放。

搜索逻辑:融合关键词与语义

这是整个系统最精巧的部分。SearchService整合了SqliteVectorStoreEmbeddingService,并实现了混合搜索算法。

graph TD
    A[用户查询 Query] --> B{SearchService.search};
    B --> C[1. 生成查询向量];
    C --> |Query Text| D[EmbeddingService];
    D --> |Query Vector| B;
    B --> E[2. FTS5关键词检索];
    E --> |Keyword Candidates| F[SqliteVectorStore];
    B --> G[3. 全量向量扫描与比对];
    G --> |All Document Vectors| F;
    F --> H{计算余弦相似度};
    H --> |Semantic Candidates| B;
    B --> I[4. Reciprocal Rank Fusion];
    I --> J[返回最终排序列表];

混合搜索的流程如下:

  1. 为用户查询生成向量。
  2. 执行FTS5查询,获取一批基于关键词匹配的候选结果。
  3. 从数据库中加载所有文档的向量。这是一个性能瓶颈,但对于我们的数据规模是可接受的。
  4. 在内存中计算查询向量与所有文档向量的余弦相似度,得到另一批基于语义的候选结果。
  5. 使用**Reciprocal Rank Fusion (RRF)**算法将两个结果集合并,生成最终的排序列表。RRF是一种简单有效的无参数融合方法,它只依赖于文档在不同结果列表中的排名。
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import java.util.stream.Stream;

public class SearchService {
    private final EmbeddingService embeddingService;
    private final SqliteVectorStore vectorStore;
    private static final int RRF_K = 60; // RRF a constant

    public SearchService(EmbeddingService embeddingService, SqliteVectorStore vectorStore) {
        this.embeddingService = embeddingService;
        this.vectorStore = vectorStore;
    }

    public List<SearchResult> search(String query, int topK) {
        // Step 1: Generate query vector
        float[] queryVector = embeddingService.generateEmbedding(query);

        // Step 2 & 3: Get keyword and semantic candidates
        // For simplicity, this example combines vector retrieval and scoring
        // In a real system, keyword search might pre-filter candidates
        List<SearchResult> semanticResults = calculateSemanticScores(queryVector);
        
        // Let's assume a similar method getFtsResults() exists
        // List<SearchResult> ftsResults = getFtsResults(query);

        // For this example, we'll focus on the vector search part
        // and its performance characteristics.
        
        semanticResults.sort(Comparator.comparingDouble(SearchResult::score).reversed());

        return semanticResults.stream().limit(topK).collect(Collectors.toList());
    }

    private List<SearchResult> calculateSemanticScores(float[] queryVector) {
        List<SqliteVectorStore.VectorRecord> allVectors = vectorStore.getAllVectors();
        
        return allVectors.parallelStream()
            .map(record -> new SearchResult(
                record.id(),
                cosineSimilarity(queryVector, record.vector())
            ))
            .collect(Collectors.toList());
    }

    private double cosineSimilarity(float[] vectorA, float[] vectorB) {
        if (vectorA.length != vectorB.length) {
            return 0.0;
        }
        double dotProduct = 0.0;
        double normA = 0.0;
        double normB = 0.0;
        for (int i = 0; i < vectorA.length; i++) {
            dotProduct += vectorA[i] * vectorB[i];
            normA += vectorA[i] * vectorA[i];
            normB += vectorB[i] * vectorB[i];
        }
        if (normA == 0.0 || normB == 0.0) {
            return 0.0;
        }
        return dotProduct / (Math.sqrt(normA) * Math.sqrt(normB));
    }
    
    // A simplified RRF implementation would look like this:
    private List<SearchResult> fuseResults(List<SearchResult> listA, List<SearchResult> listB) {
        Map<Integer, Double> rrfScores = new HashMap<>();

        // Process list A
        for (int i = 0; i < listA.size(); i++) {
            int docId = listA.get(i).id();
            double rankScore = 1.0 / (RRF_K + i + 1);
            rrfScores.merge(docId, rankScore, Double::sum);
        }

        // Process list B
        for (int i = 0; i < listB.size(); i++) {
            int docId = listB.get(i).id();
            double rankScore = 1.0 / (RRF_K + i + 1);
            rrfScores.merge(docId, rankScore, Double::sum);
        }

        return rrfScores.entrySet().stream()
                .map(entry -> new SearchResult(entry.getKey(), entry.getValue()))
                .sorted(Comparator.comparingDouble(SearchResult::score).reversed())
                .collect(Collectors.toList());
    }


    public record SearchResult(int id, double score) {}
}

calculateSemanticScores方法中的parallelStream()是一个简单的优化,它利用多核CPU来加速余弦相似度的计算。这是在没有原生向量索引的情况下,压榨单机性能的直接手段。在真实项目中,这里的性能表现直接决定了服务的响应延迟。

架构的局限与演进路径

这套架构的优点是极致的简单和独立。一个Jib命令./mvnw compile jib:build就能产出一个包含了模型、数据逻辑和服务的完整镜像,部署极其方便。

然而,其局限性也同样明显。

  1. 向量检索性能: 全量扫描和内存计算向量相似度的做法,其复杂度是O(N*D),其中N是文档数量,D是向量维度。当文档数量超过几十万,响应时间将变得无法接受。
  2. 数据与服务耦合: SQLite数据库文件与服务实例紧密绑定。虽然可以通过挂载卷实现持久化,但这种架构难以水平扩展。无法简单地通过启动多个容器实例来分担负载。
  3. 模型更新: 更新嵌入模型需要重新构建并部署整个镜像,不够灵活。

针对这些问题,未来的演进路径是清晰的:

  • 性能优化: 当数据量增长时,第一步是在SQLite之上实现一个基础的向量索引结构,例如IVF(Inverted File)。这需要将向量空间聚类,搜索时只与查询向量所在簇的文档向量进行比较,从而大幅减少计算量。这会增加实现的复杂度,但仍能维持SQLite作为单一存储的优势。
  • 架构解耦: 当业务规模进一步扩大,将向量存储和检索引入一个专门的服务是不可避免的,例如使用Vald, Milvus或一个带向量插件的PostgreSQL。此时,当前的服务将演变为一个纯粹的API网关和业务逻辑层。
  • 模型管理: 可以将模型服务(EmbeddingService)拆分为一个独立的微服务,允许其独立更新和扩展,主服务通过RPC调用它。

这个项目本身不是要打造一个能与专业搜索引擎匹敌的系统,而是展示了如何在资源受限和运维简单的要求下,利用现代AI工具和经典的软件工程技术,创造性地构建出一个“足够好”的解决方案。它的价值在于其架构选择的权衡过程,以及对每个技术组件边界的清晰认知。


  目录