def make_attention_mask(batch_size, query_length, seq_length):
  """Returns an attention mask for question generation.

  Args:
    batch_size: number of windows per batch.
    query_length: length of question including [CLS], [Q] and [SEP].
    seq_length: total number of word pieces in window.

  Returns:
    input_mask: mask used to zero the question token ids at inference time.
    attention_mask: mask used at training and inference time to make the
      representations of question tokens causal.

  The attention mask should be [batch_size, from_seq_length, to_seq_length].
  Every training example will have the same attention mask and it's:
               to:
              [CLS]   [Q]   q1    q2    ...  [SEP]  c1    c2    c3   ...
        [CLS]    1     1     0     0     0     1     1     1     1    1
  from:  [Q]     1     1     0     0     0     1     1     1     1    1
         q1      1     1     1     0     0     1     1     1     1    1
         q2      1     1     1     1     0     1     1     1     1    1
         ...     1     1     1     1     1     1     1     1     1    1
        [SEP]    1     1     0     0     0     1     1     1     1    1
         c1      1     1     0     0     0     1     1     1     1    1
         c2      1     1     0     0     0     1     1     1     1    1
         c3      1     1     0     0     0     1     1     1     1    1
         ...     1     1     0     0     0     1     1     1     1    1
  """
  # TODO(chrisalberti): add test for this.
  input_mask = np.reshape(
      np.array(
          [1, 1] + [0] * (query_length - 2) + [1] * (seq_length - query_length),
          dtype=np.int64), [1, 1, seq_length])
  attention_mask = np.reshape(input_mask, [1, 1, seq_length]) * np.ones(
      [1, seq_length, 1])  # Uses broadcasting.
  for f in range(1, query_length):
    attention_mask[0, f, 1:f + 1] = 1
  attention_mask = np.ones([batch_size, 1, 1]) * attention_mask  # Broadcasting.
  return input_mask, attention_mask
Returns an attention mask for question generation.

  Args:
    batch_size: number of windows per batch.
    query_length: length of question including [CLS], [Q] and [SEP].
    seq_length: total number of word pieces in window.

  Returns:
    input_mask: mask used to zero the question token ids at inference time.
    attention_mask: mask used at training and inference time to make the
      representations of question tokens causal.

  The attention mask should be [batch_size, from_seq_length, to_seq_length].
  Every training example will have the same attention mask and it's:
               to:
              [CLS]   [Q]   q1    q2    ...  [SEP]  c1    c2    c3   ...
        [CLS]    1     1     0     0     0     1     1     1     1    1
  from:  [Q]     1     1     0     0     0     1     1     1     1    1
         q1      1     1     1     0     0     1     1     1     1    1
         q2      1     1     1     1     0     1     1     1     1    1
         ...     1     1     1     1     1     1     1     1     1    1
        [SEP]    1     1     0     0     0     1     1     1     1    1
         c1      1     1     0     0     0     1     1     1     1    1
         c2      1     1     0     0     0     1     1     1     1    1
         c3      1     1     0     0     0     1     1     1     1    1
         ...     1     1     0     0     0     1     1     1     1    1
  
language/capwap/synthetic/generate_questions.py