extract_features 句嵌入 BERT

Question

我正在使用此代码来获取我数据集中句子的嵌入（我正在使用我的预训练模型）。

`python extract_features.py \
  --input_file=/tmp/input.txt \
  --output_file=/tmp/output.jsonl \
  --vocab_file=$BERT_BASE_DIR/vocab.txt \
  --bert_config_file=$BERT_BASE_DIR/bert_config.json \
  --init_checkpoint=$BERT_BASE_DIR/bert_model.ckpt \
  --layers=-1,-2,-3,-4 \
  --max_seq_length=128 \
  --batch_size=32`

但是，有一个问题：有没有办法更快地获得嵌入？ 因为 2000 个句子需要 6 个小时。 我的数据集包含 20000 个句子； 60 小时对于 Colab 来说太长了。 谢谢。

Answer 1

我解决了。 我在 input.txt 中写下了所有的句子，然后我使用了这段代码：

import jsonlines
df_emb=pd.DataFrame()
with jsonlines.open('/content/tmp/output.jsonl') as f:
    for line in f.iter():
        s=line['features'][0]['layers'][0]['values']
        df_tmp=pd.DataFrame(s).T
        df_emb=df_emb.append(df_tmp,ignore_index=True)

之后我将 dataframe 保存在 csv 文件中