/
githubmirror
/
the-algorithm
Обзор
Документация
Войти
/
githubmirror
/
the-algorithm
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
twml/libtwml/src/ops/feature_id.cpp
58 строк
2 KB
twitter-team
Twitter Recommendation Algorithm
01 апр 2023, 01:36
01 апр 2023, 01:36
ef4c5eb
Код
Авторство
О чём код?
#include "tensorflow/core/framework/op.h" #include "tensorflow/core/framework/shape_inference.h" #include "tensorflow/core/framework/op_kernel.h" #include <twml.h> #include "tensorflow_utils.h" using namespace tensorflow; REGISTER_OP("FeatureId") .Attr("feature_names: list(string)") .Output("output: int64") .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { return Status::OK(); }).Doc(R"doc( A tensorflow OP that hashes a list of strings into int64. This is used for feature name hashing. Attr feature_names: a list of string feature names (list(string)). Outputs ouput: hashes corresponding to the string feature names (int64). )doc"); class FeatureId : public OpKernel { private: std::vector<string> input_vector; public: explicit FeatureId(OpKernelConstruction* context) : OpKernel(context) { OP_REQUIRES_OK(context, context->GetAttr("feature_names", &input_vector)); } void Compute(OpKernelContext* context) override { // Get size of the input_vector and create TensorShape shape const int total_size = static_cast<int>(input_vector.size()); TensorShape shape = {total_size}; // Create an output tensor Tensor* output_tensor = nullptr; OP_REQUIRES_OK(context, context->allocate_output(0, shape, &output_tensor)); auto output_flat = output_tensor->flat<int64>(); // Transform the input tensor into a int64 for (int i = 0; i < total_size; i++) { output_flat(i) = twml::featureId(input_vector[i]); } } }; REGISTER_KERNEL_BUILDER( Name("FeatureId") .Device(DEVICE_CPU), FeatureId);