def get_features(self, doc_tokens, qry_tokens, ent_tokens, uid):
    """Convert list of tokens to a feature dict."""
    max_tokens_doc = self.max_seq_length - 2
    max_tokens_qry = self.max_qry_length - 2
    max_tokens_ent = self.max_ent_length
    doc_input_ids = self.tokenizer.convert_tokens_to_ids(
        ["[CLS]"] + doc_tokens[:max_tokens_doc] + ["[SEP]"])
    doc_segment_ids = [1] * len(doc_input_ids)
    doc_input_mask = [1] * len(doc_input_ids)
    while len(doc_input_ids) < self.max_seq_length:
      doc_input_ids.append(0)
      doc_input_mask.append(0)
      doc_segment_ids.append(0)
    qry_input_ids = self.tokenizer.convert_tokens_to_ids(
        ["[CLS]"] + qry_tokens[:max_tokens_qry] + ["[SEP]"])
    qry_segment_ids = [0] * len(qry_input_ids)
    qry_input_mask = [1] * len(qry_input_ids)
    while len(qry_input_ids) < self.max_qry_length:
      qry_input_ids.append(0)
      qry_input_mask.append(0)
      qry_segment_ids.append(0)
    ent_input_ids = self.tokenizer.convert_tokens_to_ids(
        ent_tokens[:max_tokens_ent])
    ent_input_mask = [1] * len(ent_input_ids)
    while len(ent_input_ids) < self.max_ent_length:
      ent_input_ids.append(0)
      ent_input_mask.append(0)
    return {
        "unique_ids": uid,
        "doc_input_ids": doc_input_ids,
        "doc_input_mask": doc_input_mask,
        "doc_segment_ids": doc_segment_ids,
        "qry_input_ids": qry_input_ids,
        "qry_input_mask": qry_input_mask,
        "qry_segment_ids": qry_segment_ids,
        "ent_input_ids": ent_input_ids,
        "ent_input_mask": ent_input_mask,
    }
Convert list of tokens to a feature dict.
language/labs/drkit/bert_utils.py