tensorflow - 从 3D 张量中提取某个标记的嵌入

Question

每个训练示例包含一个句子（一个标记列表）、一个起始标记索引和一个结束标记索引。开始和结束索引突出显示句子中的一些单词。标签为 True 或 False，表示突出显示的单词是否“有趣”。
例如，“新闻”是此示例中突出显示的单词。

tokens=["This", "news", "is", "sad"], start=1, end=2, label=1 

这是我的模型设置。句子被馈送到预训练的 bert 层，该层返回一个 3D 张量，其中包含每个句子中每个标记的嵌入。

我想从 3D 张量中提取每个示例的高亮词的嵌入并获得它们的平均值，这样我就可以将它提供给密集层，我应该怎么做？假设我已经有了 3D 张量。

我正在使用 tensorflow 2.0 提供的 keras API。我是 tensorflow 的新手，所以非常感谢一个具体的例子。

谢谢！

标签： tensorflowkeras