/*
|
* Copyright (C) 2018 The Android Open Source Project
|
*
|
* 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.
|
*/
|
|
#include "actions/lua-utils.h"
|
|
namespace libtextclassifier3 {
|
namespace {
|
static constexpr const char* kTextKey = "text";
|
static constexpr const char* kTimeUsecKey = "parsed_time_ms_utc";
|
static constexpr const char* kGranularityKey = "granularity";
|
static constexpr const char* kCollectionKey = "collection";
|
static constexpr const char* kNameKey = "name";
|
static constexpr const char* kScoreKey = "score";
|
static constexpr const char* kPriorityScoreKey = "priority_score";
|
static constexpr const char* kTypeKey = "type";
|
static constexpr const char* kResponseTextKey = "response_text";
|
static constexpr const char* kAnnotationKey = "annotation";
|
static constexpr const char* kSpanKey = "span";
|
static constexpr const char* kMessageKey = "message";
|
static constexpr const char* kBeginKey = "begin";
|
static constexpr const char* kEndKey = "end";
|
static constexpr const char* kClassificationKey = "classification";
|
static constexpr const char* kSerializedEntity = "serialized_entity";
|
static constexpr const char* kEntityKey = "entity";
|
} // namespace
|
|
template <>
|
int AnnotationIterator<ClassificationResult>::Item(
|
const std::vector<ClassificationResult>* annotations, StringPiece key,
|
lua_State* state) const {
|
// Lookup annotation by collection.
|
for (const ClassificationResult& annotation : *annotations) {
|
if (key.Equals(annotation.collection)) {
|
PushAnnotation(annotation, entity_data_schema_, env_);
|
return 1;
|
}
|
}
|
TC3_LOG(ERROR) << "No annotation with collection: " << key.ToString()
|
<< " found.";
|
lua_error(state);
|
return 0;
|
}
|
|
template <>
|
int AnnotationIterator<ActionSuggestionAnnotation>::Item(
|
const std::vector<ActionSuggestionAnnotation>* annotations, StringPiece key,
|
lua_State* state) const {
|
// Lookup annotation by name.
|
for (const ActionSuggestionAnnotation& annotation : *annotations) {
|
if (key.Equals(annotation.name)) {
|
PushAnnotation(annotation, entity_data_schema_, env_);
|
return 1;
|
}
|
}
|
TC3_LOG(ERROR) << "No annotation with name: " << key.ToString() << " found.";
|
lua_error(state);
|
return 0;
|
}
|
|
void PushAnnotation(const ClassificationResult& classification,
|
const reflection::Schema* entity_data_schema,
|
LuaEnvironment* env) {
|
if (entity_data_schema == nullptr ||
|
classification.serialized_entity_data.empty()) {
|
// Empty table.
|
lua_newtable(env->state());
|
} else {
|
env->PushFlatbuffer(entity_data_schema,
|
flatbuffers::GetRoot<flatbuffers::Table>(
|
classification.serialized_entity_data.data()));
|
}
|
lua_pushinteger(env->state(),
|
classification.datetime_parse_result.time_ms_utc);
|
lua_setfield(env->state(), /*idx=*/-2, kTimeUsecKey);
|
lua_pushinteger(env->state(),
|
classification.datetime_parse_result.granularity);
|
lua_setfield(env->state(), /*idx=*/-2, kGranularityKey);
|
env->PushString(classification.collection);
|
lua_setfield(env->state(), /*idx=*/-2, kCollectionKey);
|
lua_pushnumber(env->state(), classification.score);
|
lua_setfield(env->state(), /*idx=*/-2, kScoreKey);
|
env->PushString(classification.serialized_entity_data);
|
lua_setfield(env->state(), /*idx=*/-2, kSerializedEntity);
|
}
|
|
void PushAnnotation(const ClassificationResult& classification,
|
StringPiece text,
|
const reflection::Schema* entity_data_schema,
|
LuaEnvironment* env) {
|
PushAnnotation(classification, entity_data_schema, env);
|
env->PushString(text);
|
lua_setfield(env->state(), /*idx=*/-2, kTextKey);
|
}
|
|
void PushAnnotatedSpan(
|
const AnnotatedSpan& annotated_span,
|
const AnnotationIterator<ClassificationResult>& annotation_iterator,
|
LuaEnvironment* env) {
|
lua_newtable(env->state());
|
{
|
lua_newtable(env->state());
|
lua_pushinteger(env->state(), annotated_span.span.first);
|
lua_setfield(env->state(), /*idx=*/-2, kBeginKey);
|
lua_pushinteger(env->state(), annotated_span.span.second);
|
lua_setfield(env->state(), /*idx=*/-2, kEndKey);
|
}
|
lua_setfield(env->state(), /*idx=*/-2, kSpanKey);
|
annotation_iterator.NewIterator(kClassificationKey,
|
&annotated_span.classification, env->state());
|
lua_setfield(env->state(), /*idx=*/-2, kClassificationKey);
|
}
|
|
MessageTextSpan ReadSpan(LuaEnvironment* env) {
|
MessageTextSpan span;
|
lua_pushnil(env->state());
|
while (lua_next(env->state(), /*idx=*/-2)) {
|
const StringPiece key = env->ReadString(/*index=*/-2);
|
if (key.Equals(kMessageKey)) {
|
span.message_index =
|
static_cast<int>(lua_tonumber(env->state(), /*idx=*/-1));
|
} else if (key.Equals(kBeginKey)) {
|
span.span.first =
|
static_cast<int>(lua_tonumber(env->state(), /*idx=*/-1));
|
} else if (key.Equals(kEndKey)) {
|
span.span.second =
|
static_cast<int>(lua_tonumber(env->state(), /*idx=*/-1));
|
} else if (key.Equals(kTextKey)) {
|
span.text = env->ReadString(/*index=*/-1).ToString();
|
} else {
|
TC3_LOG(INFO) << "Unknown span field: " << key.ToString();
|
}
|
lua_pop(env->state(), 1);
|
}
|
return span;
|
}
|
|
int ReadAnnotations(const reflection::Schema* entity_data_schema,
|
LuaEnvironment* env,
|
std::vector<ActionSuggestionAnnotation>* annotations) {
|
if (lua_type(env->state(), /*idx=*/-1) != LUA_TTABLE) {
|
TC3_LOG(ERROR) << "Expected annotations table, got: "
|
<< lua_type(env->state(), /*idx=*/-1);
|
lua_pop(env->state(), 1);
|
lua_error(env->state());
|
return LUA_ERRRUN;
|
}
|
|
// Read actions.
|
lua_pushnil(env->state());
|
while (lua_next(env->state(), /*idx=*/-2)) {
|
if (lua_type(env->state(), /*idx=*/-1) != LUA_TTABLE) {
|
TC3_LOG(ERROR) << "Expected annotation table, got: "
|
<< lua_type(env->state(), /*idx=*/-1);
|
lua_pop(env->state(), 1);
|
continue;
|
}
|
annotations->push_back(ReadAnnotation(entity_data_schema, env));
|
lua_pop(env->state(), 1);
|
}
|
return LUA_OK;
|
}
|
|
ActionSuggestionAnnotation ReadAnnotation(
|
const reflection::Schema* entity_data_schema, LuaEnvironment* env) {
|
ActionSuggestionAnnotation annotation;
|
lua_pushnil(env->state());
|
while (lua_next(env->state(), /*idx=*/-2)) {
|
const StringPiece key = env->ReadString(/*index=*/-2);
|
if (key.Equals(kNameKey)) {
|
annotation.name = env->ReadString(/*index=*/-1).ToString();
|
} else if (key.Equals(kSpanKey)) {
|
annotation.span = ReadSpan(env);
|
} else if (key.Equals(kEntityKey)) {
|
annotation.entity = ReadClassificationResult(entity_data_schema, env);
|
} else {
|
TC3_LOG(ERROR) << "Unknown annotation field: " << key.ToString();
|
}
|
lua_pop(env->state(), 1);
|
}
|
return annotation;
|
}
|
|
ClassificationResult ReadClassificationResult(
|
const reflection::Schema* entity_data_schema, LuaEnvironment* env) {
|
ClassificationResult classification;
|
lua_pushnil(env->state());
|
while (lua_next(env->state(), /*idx=*/-2)) {
|
const StringPiece key = env->ReadString(/*index=*/-2);
|
if (key.Equals(kCollectionKey)) {
|
classification.collection = env->ReadString(/*index=*/-1).ToString();
|
} else if (key.Equals(kScoreKey)) {
|
classification.score =
|
static_cast<float>(lua_tonumber(env->state(), /*idx=*/-1));
|
} else if (key.Equals(kTimeUsecKey)) {
|
classification.datetime_parse_result.time_ms_utc =
|
static_cast<int64>(lua_tonumber(env->state(), /*idx=*/-1));
|
} else if (key.Equals(kGranularityKey)) {
|
classification.datetime_parse_result.granularity =
|
static_cast<DatetimeGranularity>(
|
lua_tonumber(env->state(), /*idx=*/-1));
|
} else if (key.Equals(kSerializedEntity)) {
|
classification.serialized_entity_data =
|
env->ReadString(/*index=*/-1).ToString();
|
} else if (key.Equals(kEntityKey)) {
|
auto buffer = ReflectiveFlatbufferBuilder(entity_data_schema).NewRoot();
|
env->ReadFlatbuffer(buffer.get());
|
classification.serialized_entity_data = buffer->Serialize();
|
} else {
|
TC3_LOG(INFO) << "Unknown classification result field: "
|
<< key.ToString();
|
}
|
lua_pop(env->state(), 1);
|
}
|
return classification;
|
}
|
|
void PushAnnotation(const ActionSuggestionAnnotation& annotation,
|
const reflection::Schema* entity_data_schema,
|
LuaEnvironment* env) {
|
PushAnnotation(annotation.entity, annotation.span.text, entity_data_schema,
|
env);
|
env->PushString(annotation.name);
|
lua_setfield(env->state(), /*idx=*/-2, kNameKey);
|
{
|
lua_newtable(env->state());
|
lua_pushinteger(env->state(), annotation.span.message_index);
|
lua_setfield(env->state(), /*idx=*/-2, kMessageKey);
|
lua_pushinteger(env->state(), annotation.span.span.first);
|
lua_setfield(env->state(), /*idx=*/-2, kBeginKey);
|
lua_pushinteger(env->state(), annotation.span.span.second);
|
lua_setfield(env->state(), /*idx=*/-2, kEndKey);
|
}
|
lua_setfield(env->state(), /*idx=*/-2, kSpanKey);
|
}
|
|
void PushAction(
|
const ActionSuggestion& action,
|
const reflection::Schema* entity_data_schema,
|
const AnnotationIterator<ActionSuggestionAnnotation>& annotation_iterator,
|
LuaEnvironment* env) {
|
if (entity_data_schema == nullptr || action.serialized_entity_data.empty()) {
|
// Empty table.
|
lua_newtable(env->state());
|
} else {
|
env->PushFlatbuffer(entity_data_schema,
|
flatbuffers::GetRoot<flatbuffers::Table>(
|
action.serialized_entity_data.data()));
|
}
|
env->PushString(action.type);
|
lua_setfield(env->state(), /*idx=*/-2, kTypeKey);
|
env->PushString(action.response_text);
|
lua_setfield(env->state(), /*idx=*/-2, kResponseTextKey);
|
lua_pushnumber(env->state(), action.score);
|
lua_setfield(env->state(), /*idx=*/-2, kScoreKey);
|
lua_pushnumber(env->state(), action.priority_score);
|
lua_setfield(env->state(), /*idx=*/-2, kPriorityScoreKey);
|
annotation_iterator.NewIterator(kAnnotationKey, &action.annotations,
|
env->state());
|
lua_setfield(env->state(), /*idx=*/-2, kAnnotationKey);
|
}
|
|
ActionSuggestion ReadAction(
|
const reflection::Schema* actions_entity_data_schema,
|
const reflection::Schema* annotations_entity_data_schema,
|
LuaEnvironment* env) {
|
ActionSuggestion action;
|
lua_pushnil(env->state());
|
while (lua_next(env->state(), /*idx=*/-2)) {
|
const StringPiece key = env->ReadString(/*index=*/-2);
|
if (key.Equals(kResponseTextKey)) {
|
action.response_text = env->ReadString(/*index=*/-1).ToString();
|
} else if (key.Equals(kTypeKey)) {
|
action.type = env->ReadString(/*index=*/-1).ToString();
|
} else if (key.Equals(kScoreKey)) {
|
action.score = static_cast<float>(lua_tonumber(env->state(), /*idx=*/-1));
|
} else if (key.Equals(kPriorityScoreKey)) {
|
action.priority_score =
|
static_cast<float>(lua_tonumber(env->state(), /*idx=*/-1));
|
} else if (key.Equals(kAnnotationKey)) {
|
ReadAnnotations(actions_entity_data_schema, env, &action.annotations);
|
} else if (key.Equals(kEntityKey)) {
|
auto buffer =
|
ReflectiveFlatbufferBuilder(actions_entity_data_schema).NewRoot();
|
env->ReadFlatbuffer(buffer.get());
|
action.serialized_entity_data = buffer->Serialize();
|
} else {
|
TC3_LOG(INFO) << "Unknown action field: " << key.ToString();
|
}
|
lua_pop(env->state(), 1);
|
}
|
return action;
|
}
|
|
int ReadActions(const reflection::Schema* actions_entity_data_schema,
|
const reflection::Schema* annotations_entity_data_schema,
|
LuaEnvironment* env, std::vector<ActionSuggestion>* actions) {
|
if (lua_type(env->state(), /*idx=*/-1) != LUA_TTABLE) {
|
TC3_LOG(ERROR) << "Expected actions table, got: "
|
<< lua_type(env->state(), /*idx=*/-1);
|
lua_pop(env->state(), 1);
|
lua_error(env->state());
|
return LUA_ERRRUN;
|
}
|
|
// Read actions.
|
lua_pushnil(env->state());
|
while (lua_next(env->state(), /*idx=*/-2)) {
|
if (lua_type(env->state(), /*idx=*/-1) != LUA_TTABLE) {
|
TC3_LOG(ERROR) << "Expected action table, got: "
|
<< lua_type(env->state(), /*idx=*/-1);
|
lua_pop(env->state(), 1);
|
continue;
|
}
|
actions->push_back(ReadAction(actions_entity_data_schema,
|
annotations_entity_data_schema, env));
|
lua_pop(env->state(), /*n=1*/ 1);
|
}
|
lua_pop(env->state(), /*n=*/1);
|
|
return LUA_OK;
|
}
|
|
int ConversationIterator::Item(const std::vector<ConversationMessage>* messages,
|
const int64 pos, lua_State* state) const {
|
const ConversationMessage& message = (*messages)[pos];
|
lua_newtable(state);
|
lua_pushinteger(state, message.user_id);
|
lua_setfield(state, /*idx=*/-2, "user_id");
|
env_->PushString(message.text);
|
lua_setfield(state, /*idx=*/-2, "text");
|
lua_pushinteger(state, message.reference_time_ms_utc);
|
lua_setfield(state, /*idx=*/-2, "time_ms_utc");
|
env_->PushString(message.reference_timezone);
|
lua_setfield(state, /*idx=*/-2, "timezone");
|
annotated_span_iterator_.NewIterator("annotation", &message.annotations,
|
state);
|
lua_setfield(state, /*idx=*/-2, "annotation");
|
return 1;
|
}
|
|
} // namespace libtextclassifier3
|
