Repository: nomic-ai/gpt4all
File: gpt4all-chat/src/database.cpp (C++)

Code before:
            (us::byte_t const *)query.data(),           1,      n_embd * sizeof(float),
            kBatch, metric
        );

        for (int i = 0; i < kBatch; ++i) {
            auto offset = batchResults.at(0)[i].offset;
            us::distance_punned_t distance = batchResults.at(0)[i].distance;
            results.append({batchChunkIds[offset], distance});
        }
    }

    // get top-k nearest neighbors of combined results
    nNeighbors = qMin(nNeighbors, results.size());
    std::partial_sort(
        results.begin(), results.begin() + nNeighbors, results.end(),
        [](const SearchResult &a, const SearchResult &b) { return a.dist < b.dist; }
    );

    QList<int> chunkIds;
    chunkIds.reserve(nNeighbors);
    for (int i = 0; i < nNeighbors; i++)
        chunkIds << results[i].chunkId;
    return chunkIds;
}

QList<int> Database::searchEmbeddings(const std::vector<float> &query, const QList<QString> &collections,
    int nNeighbors)
{
    QSqlQuery q(m_db);
    if (!q.exec(GET_COLLECTION_EMBEDDINGS_SQL.arg(collections.join("', '")))) {
        qWarning() << "Database ERROR: Failed to exec embeddings query:" << q.lastError();
        return {};
    }
    return searchEmbeddingsHelper(query, q, nNeighbors);
}

QList<int> Database::scoreChunks(const std::vector<float> &query, const QList<int> &chunks)
{
    QList<QString> chunkStrings;
    for (int id : chunks)
        chunkStrings << QString::number(id);
    QSqlQuery q(m_db);
    if (!q.exec(GET_CHUNK_EMBEDDINGS_SQL.arg(chunkStrings.join("', '")))) {
        qWarning() << "Database ERROR: Failed to exec embeddings query:" << q.lastError();
        return {};
    }
    return searchEmbeddingsHelper(query, q, chunks.size());
}

QList<Database::BM25Query> Database::queriesForFTS5(const QString &input)
{

Review comment (suggestion):
I don't think it's possible for results.size() to exceed k. If we're not certain, we should add an assertion instead of using qMin. If we can assume this, the final loop can become a standard foreach.

Code after:
    QSqlQuery sqlQuery(m_db);
    sqlQuery.prepare(SELECT_CHUNKS_FTS_SQL.arg(k));

    QList<SearchResult> results;
    for (auto &bm25Query : std::as_const(bm25Queries)) {
        sqlQuery.addBindValue(bm25Query.query);

        if (!sqlQuery.exec()) {
            qWarning() << "Database ERROR: Failed to execute BM25 query:" << sqlQuery.lastError();
            return {};
        }

        if (sqlQuery.next()) {
            // Save the query that was used to produce results
            bm25q = bm25Query;
            break;
        }
    }

    do {
        const int chunkId = sqlQuery.value(0).toInt();
        const float score = sqlQuery.value(1).toFloat();
        results.append({chunkId, score});
    } while (sqlQuery.next());

    k = qMin(k, results.size());
    std::partial_sort(
        results.begin(), results.begin() + k, results.end(),
        [](const SearchResult &a, const SearchResult &b) { return a.score < b.score; }
    );

    QList<int> chunkIds;
    chunkIds.reserve(k);
    for (int i = 0; i < k; i++)
        chunkIds << results[i].chunkId;
    return chunkIds;
}

float Database::computeBM25Weight(const Database::BM25Query &bm25q)
{
    float bmWeight = 0.0f;
    if (bm25q.isExact) {
        bmWeight = 0.9f; // the highest we give
    } else {
        // qlength is the length of the phrases in the query by number of distinct words
        // ilength is the length of the natural language query by number of distinct words
        // rlength is the number of stop words removed from the natural language query to form the query

        // calculate the query length weight based on the ratio of query terms to meaningful terms.
        // this formula adjusts the weight with the empirically determined insight that BM25's
        // effectiveness decreases as query length increases.