chromium/components/omnibox/browser/on_device_tail_model_service.cc

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

#include "components/omnibox/browser/on_device_tail_model_service.h"

#include <optional>
#include <utility>

#include "base/command_line.h"
#include "base/containers/flat_set.h"
#include "base/files/file_path.h"
#include "base/functional/bind.h"
#include "base/functional/callback.h"
#include "base/logging.h"
#include "base/memory/memory_pressure_monitor.h"
#include "base/metrics/field_trial_params.h"
#include "base/metrics/histogram_macros.h"
#include "base/task/task_traits.h"
#include "base/task/thread_pool.h"
#include "components/omnibox/browser/on_device_tail_model_executor.h"
#include "components/optimization_guide/core/optimization_guide_model_provider.h"
#include "components/optimization_guide/core/optimization_guide_util.h"
#include "components/optimization_guide/proto/models.pb.h"
#include "components/optimization_guide/proto/on_device_tail_suggest_model_metadata.pb.h"

namespace {
// Constants for TFlite model validation.
constexpr std::string kTestPrefix = …;
constexpr std::string_view kModelValidationSwitchName = …;

void InitializeTailModelExecutor(
    OnDeviceTailModelExecutor* executor,
    const base::FilePath& model_file,
    const base::flat_set<base::FilePath>& additional_files,
    const optimization_guide::proto::OnDeviceTailSuggestModelMetadata&
        metadata) { … }

std::vector<OnDeviceTailModelExecutor::Prediction> RunTailModelExecutor(
    OnDeviceTailModelExecutor* executor,
    const OnDeviceTailModelExecutor::ModelInput& input) { … }

void MaybeUnloadModelExecutor(OnDeviceTailModelExecutor* executor) { … }

}  // namespace

OnDeviceTailModelService::OnDeviceTailModelService(
    optimization_guide::OptimizationGuideModelProvider* model_provider)
    : … { … }

OnDeviceTailModelService::OnDeviceTailModelService()
    : … { … }

OnDeviceTailModelService::~OnDeviceTailModelService() { … }

void OnDeviceTailModelService::Shutdown() { … }

void OnDeviceTailModelService::OnModelUpdated(
    optimization_guide::proto::OptimizationTarget optimization_target,
    base::optional_ref<const optimization_guide::ModelInfo> model_info) { … }

void OnDeviceTailModelService::OnMemoryPressure(
    base::MemoryPressureListener::MemoryPressureLevel level) { … }

void OnDeviceTailModelService::GetPredictionsForInput(
    const OnDeviceTailModelExecutor::ModelInput& input,
    ResultCallback result_callback) { … }