/*
 * Copyright (c) 2023, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once

#include <embedding/common.hpp>

namespace embedding {

class ILookup {
 public:
  virtual ~ILookup() = default;

  virtual void lookup(const core23::Tensor& keys, size_t num_keys,
                      const core23::Tensor& num_keys_per_table_offset, size_t num_table_offset,
                      const core23::Tensor& table_id_list, core23::Tensor& embedding_vec) = 0;

  virtual void ratio_filter(const core23::Tensor& keys, size_t num_keys,
                            const core23::Tensor& id_space_offset, size_t num_id_space_offset,
                            const core23::Tensor& id_space, core23::Tensor& filtered){};
};

}  // namespace embedding
