chromium/services/passage_embeddings/passage_embedder.cc

// Copyright 2024 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "services/passage_embeddings/passage_embedder.h"

#include "base/containers/heap_array.h"
#include "base/files/file.h"
#include "base/metrics/histogram_functions.h"
#include "base/timer/elapsed_timer.h"
#include "base/trace_event/trace_event.h"
#include "base/trace_event/trace_id_helper.h"
#include "base/trace_event/typed_macros.h"
#include "components/history_embeddings/history_embeddings_features.h"
#include "components/optimization_guide/core/tflite_op_resolver.h"
#include "third_party/sentencepiece/src/src/sentencepiece_model.pb.h"

namespace {
// Records duration and trace event for embeddings generation.
void RecordEmbeddingsDurationMetrics(
    bool is_passive,
    base::TimeTicks start_time,
    base::TimeDelta elapsed_time,
    std::optional<base::TimeDelta> elapsed_thread_time) { … }
}  // namespace

namespace passage_embeddings {

PassageEmbedder::PassageEmbedder(
    mojo::PendingReceiver<mojom::PassageEmbedder> receiver)
    : … { … }

PassageEmbedder::~PassageEmbedder() = default;

bool PassageEmbedder::LoadModels(
    base::File* embeddings_model_file,
    base::File* sp_file,
    std::unique_ptr<tflite::task::core::TfLiteEngine> tflite_engine) { … }

void PassageEmbedder::SetEmbeddingsModelInputWindowSize(uint32_t size) { … }

bool PassageEmbedder::LoadSentencePieceModelFile(base::File* sp_file) { … }

bool PassageEmbedder::LoadEmbeddingsModelFile(
    base::File* embeddings_file,
    std::unique_ptr<tflite::task::core::TfLiteEngine> tflite_engine) { … }

bool PassageEmbedder::BuildExecutionTask() { … }

void PassageEmbedder::UnloadModelFiles() { … }

std::optional<OutputType> PassageEmbedder::Execute(InputType input) { … }

void PassageEmbedder::GenerateEmbeddings(
    const std::vector<std::string>& inputs,
    mojom::PassagePriority priority,
    PassageEmbedder::GenerateEmbeddingsCallback callback) { … }

}  // namespace passage_embeddings