| 241 | class NegTrainOp : public OpKernel { |
| 242 | public: |
| 243 | explicit NegTrainOp(OpKernelConstruction* ctx) : OpKernel(ctx) { |
| 244 | base_.Init(0, 0); |
| 245 | |
| 246 | OP_REQUIRES_OK(ctx, ctx->GetAttr("num_negative_samples", &num_samples_)); |
| 247 | |
| 248 | std::vector<int32> vocab_count; |
| 249 | OP_REQUIRES_OK(ctx, ctx->GetAttr("vocab_count", &vocab_count)); |
| 250 | |
| 251 | std::vector<float> vocab_weights; |
| 252 | vocab_weights.reserve(vocab_count.size()); |
| 253 | for (const auto& f : vocab_count) { |
| 254 | float r = std::pow(static_cast<float>(f), 0.75f); |
| 255 | vocab_weights.push_back(r); |
| 256 | } |
| 257 | sampler_ = new random::DistributionSampler(vocab_weights); |
| 258 | } |
| 259 | |
| 260 | ~NegTrainOp() override { delete sampler_; } |
| 261 | |