BigBird 是一种基于稀疏注意力机制的 Transformer,它对 BERT 等基于 Transformer 的模型进行了扩展,使其能够处理更长的序列。此外,BigBird 还从理论层面阐释了稀疏模型所能具备的完整 Transformer 的能力。
BigBird 在这篇论文中被首次提出,并在该代码库中首次发布。
BigBird 采用块稀疏注意力机制替代常规注意力机制(如 BERT 的注意力机制),能够以远低于 BERT 的计算成本处理长度达 4096 的序列。它在涉及超长序列的各类任务中(如长文档摘要、长上下文问答)均取得了最先进的性能。
以下是在 PyTorch 中使用该模型获取给定文本特征的方法:
import argparse
from openmind import AutoModel, AutoTokenizer
from openmind import is_torch_npu_available
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--model_name_or_path",type=str,help="Path to model",default=None,)
args = parser.parse_args()
return args
if __name__ == '__main__':
if is_torch_npu_available():
device = "npu:0"
else:
device = "cpu"
args = parse_args()
model_path = args.model_name_or_path
# Note: CodeSage requires adding eos token at the end of
# each tokenized sequence to ensure good performance
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, add_eos_token=True)
model = AutoModel.from_pretrained(model_path, trust_remote_code=True).to(device)
inputs = tokenizer.encode("def print_hello_world():\tprint('Hello World!')", return_tensors="pt").to(device)
embedding = model(inputs)[0]
print(f'Dimension of the embedding: {embedding[0].size()}')
print(embedding)