Function ExecuteTfLite

tensorflow/lite/models/smartreply/predictor.cc:45–68

Predict with TfLite model.

// Predict with TfLite model.
void ExecuteTfLite(const std::string& sentence,
                   ::tflite::Interpreter* interpreter,
                   std::map<std::string, float>* response_map) {
  {
    TfLiteTensor* input = interpreter->tensor(interpreter->inputs()[0]);
    tflite::DynamicBuffer buf;
    buf.AddString(sentence.data(), sentence.length());
    buf.WriteToTensorAsVector(input);
    interpreter->AllocateTensors();

    interpreter->Invoke();

    TfLiteTensor* messages = interpreter->tensor(interpreter->outputs()[0]);
    TfLiteTensor* confidence = interpreter->tensor(interpreter->outputs()[1]);

    for (int i = 0; i < confidence->dims->data[0]; i++) {
      float weight = confidence->data.f[i];
      auto response_text = tflite::GetString(messages, i);
      if (response_text.len > 0) {
        (*response_map)[string(response_text.str, response_text.len)] += weight;
      }
    }
  }
}

void GetSegmentPredictions(
    const std::vector<std::string>& input,