您可以直接通过一个管道将此模型用于掩码语言建模:
from openmind import pipeline, AutoTokenizer, is_torch_npu_available
import argparse
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument(
"--model_name_or_path",
type=str,
help="Path to model",
default="ChongqingAscend/distilbert-base-cased",
)
args = parser.parse_args()
return args
def main():
args = parse_args()
model_path = args.model_name_or_path
if is_torch_npu_available():
device = "npu:0"
else:
device = "cpu"
tokenizer = AutoTokenizer.from_pretrained(model_path)
pipe = pipeline("fill-mask", model=model_path, tokenizer=tokenizer, device=device)
out = pipe("Hello I'm a [MASK] model.")
print(out)
if __name__ == "__main__":
main()