Tensorflow 将两个列表合并到PCollection

我正在使用apachebeam。在写入tfRecord时，我需要包括项的ID及其文本和嵌入。 本教程只处理一个文本列表，但我也有一个ID列表来匹配文本列表，因此我想知道如何将ID传递给以下函数：

  def to_tf_example(entries):
  examples = []

  text_list, embedding_list = entries
  for i in range(len(text_list)):
    text = text_list[i]
    embedding = embedding_list[i]

    features = {
        # need to pass in ID here like so:
        'id': tf.train.Feature(
            bytes_list=tf.train.BytesList(value=[ids.encode('utf-8')])),
        'text': tf.train.Feature(
            bytes_list=tf.train.BytesList(value=[text.encode('utf-8')])),
        'embedding': tf.train.Feature(
            float_list=tf.train.FloatList(value=embedding.tolist()))
    }
  
    example = tf.train.Example(
        features=tf.train.Features(
            feature=features)).SerializeToString(deterministic=True)
  
    examples.append(example)
  
  return examples

我的第一个想法是将ID包含在数据库的文本列中，然后通过切片或正则表达式或其他方式提取它们，但我想知道是否有更好的方法，我假设转换为PCollection，但不知道从何处开始。以下是管道：

    with beam.Pipeline(args.runner, options=options) as pipeline:
        query_data = pipeline | 'Read data from BigQuery' >> 
        beam.io.Read(beam.io.BigQuerySource(project='my-project', query=get_data(args.limit), use_standard_sql=True))
        # list of texts
        text = query_data | 'get list of text' >> beam.Map(lambda x: x['text'])
        # list of ids
        ids = query_data | 'get list of ids' >> beam.Map(lambda x: x['id'])
    
        ( text
            | 'Batch elements' >> util.BatchElements(
            min_batch_size=args.batch_size, max_batch_size=args.batch_size)
            | 'Generate embeddings' >> beam.Map(
            generate_embeddings, args.module_url, args.random_projection_matrix)
            | 'Encode to tf example' >> beam.FlatMap(to_tf_example)
            | 'Write to TFRecords files' >> beam.io.WriteToTFRecord(
            file_path_prefix='{0}'.format(args.output_dir),
            file_name_suffix='.tfrecords')
        )

        query_data | 'Convert to entity and write to datastore' >> beam.Map(
                lambda input_features: create_entity(
                    input_features, args.kind))

---

这里我假设

generate_embeddings

具有签名

List[str]，…->（列表[str]，列表[List[float]]）

您要做的是避免将文本和ID拆分为单独的PCollection。所以你可能想写一些像

def generate_embeddings_for_batch(
    batch,
    module_url,
    random_projection_matrix) -> Tuple[int, str, List[float]]:
  embeddings = generate_embeddings(
      [x['text'] for x in batch], module_url, random_projection_matrix)
  text_to_embedding = dict(embeddings)
  for id, text in batch:
    yield x['id'], x['text'], text_to_embedding[x['text']]

从那里你应该可以写一个例子

---

看看如何使用可能是有意义的。

我改变了generate_嵌入以返回List[int]、List[string]、List[List[float]]，然后使用以下函数将id和文本列表传入：

 def generate_embeddings_for_batch(batch, module_url, random_projection_matrix):
  embeddings = generate_embeddings([x['id'] for x in batch], [x['text'] for x in batch], module_url, random_projection_matrix)
  return embeddings

---

最后，我将generate_嵌入更改为返回List[int]、List[string]、List[List[float]]。然后通过列表理解传递ID。谢谢你的建议。