embedding.h

Go to the documentation of this file.

//
//  embedding.cpp
//
//  Created by MNN on 2023/09/25.
//  ZhaodeWang
//
 
#ifndef _NNDEPLOY_LLM_EMBEDDING_H_
#define _NNDEPLOY_LLM_EMBEDDING_H_
 
#include "nndeploy/base/any.h"
#include "nndeploy/base/common.h"
#include "nndeploy/base/glic_stl_include.h"
#include "nndeploy/base/log.h"
#include "nndeploy/base/macro.h"
#include "nndeploy/base/object.h"
#include "nndeploy/base/opencv_include.h"
#include "nndeploy/base/param.h"
#include "nndeploy/base/status.h"
#include "nndeploy/base/string.h"
#include "nndeploy/dag/composite_node.h"
#include "nndeploy/dag/edge.h"
#include "nndeploy/dag/graph.h"
#include "nndeploy/dag/loop.h"
#include "nndeploy/dag/node.h"
#include "nndeploy/device/buffer.h"
#include "nndeploy/device/device.h"
#include "nndeploy/device/memory_pool.h"
#include "nndeploy/device/tensor.h"
#include "nndeploy/llm/embedding/diskembedding.hpp"
#include "nndeploy/tokenizer/tokenizer.h"
 
namespace nndeploy {
namespace llm {
 
class NNDEPLOY_CC_API EmbeddingParam : public base::Param {
 public:
  EmbeddingParam() = default;
  virtual ~EmbeddingParam() = default;
 
  PARAM_COPY(EmbeddingParam)
  PARAM_COPY_TO(EmbeddingParam)
 
  // 隐藏层维度
  int hidden_size_ = 4096;
  // 嵌入权重文件路径
  std::string embedding_weight_path_ = "";
  // 量化相关参数
  bool use_quantization_ = false;
  //
  int weight_offset_ = 0;
  //
  int a_offset_ = 0;
  int alpha_size_ = 0;
  // 量化比特数
  int quant_bit_ = 8;
  // 量化块大小
  int quant_block_ = 0;
 
  // other
  base::DataType data_type_ = base::dataTypeOf<float>();
  base::DataFormat data_format_ = base::DataFormat::kDataFormatNCHW;
  std::string share_disk_embedding_key_ = "disk_embedding";
 
  std::string getShareKey() {
    std::string key = "";
    key += embedding_weight_path_;
    key += std::to_string(hidden_size_);
    key += base::dataTypeToString(data_type_);
    key += base::dataFormatToString(data_format_);
    key += std::to_string(use_quantization_);
    key += std::to_string(weight_offset_);
    key += std::to_string(a_offset_);
    key += std::to_string(alpha_size_);
    key += std::to_string(quant_bit_);
    return key;
  }
  using base::Param::serialize;
  virtual base::Status serialize(
      rapidjson::Value& json,
      rapidjson::Document::AllocatorType& allocator) override;
  using base::Param::deserialize;
  virtual base::Status deserialize(rapidjson::Value& json) override;
};
 
class NNDEPLOY_CC_API Embedding : public dag::Node {
 public:
  Embedding(const std::string& name, std::vector<dag::Edge*> inputs,
            std::vector<dag::Edge*> outputs);
  virtual ~Embedding();
 
  virtual base::Status init();
  virtual base::Status deinit();
 
  virtual base::Status run();
 
 private:
  std::shared_ptr<MNN::Transformer::DiskEmbedding> disk_embedding_ = nullptr;
};
 
}  // namespace llm
}  // namespace nndeploy
 
#endif