SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Add bart-base (#5014)

Browse Source
This commit is contained in:
Sam Shleifer
2020-06-15 13:29:26 -04:00
committed by GitHub
parent 7b685f5229
commit a9f1fc6c94
3 changed files with 20 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
tests/test_modeling_bart.py
View File

@@ -40,6 +40,7 @@ if is_torch_available():
BartTokenizer,
MBartTokenizer,
BatchEncoding,
pipeline,
)
from transformers.modeling_bart import (
BART_PRETRAINED_MODEL_ARCHIVE_LIST,
@@ -565,6 +566,22 @@ class BartModelIntegrationTests(unittest.TestCase):
)
self.assertTrue(torch.allclose(output[:, :3, :3], expected_slice, atol=TOLERANCE))
@slow
def test_bart_base_mask_filling(self):
pbase = pipeline(task="fill-mask", model="facebook/bart-base")
src_text = [" I went to the <mask>."]
results = [x["token_str"] for x in pbase(src_text)]
expected_results = ["Ġbathroom", "Ġrestroom", "Ġhospital", "Ġkitchen", "Ġcar"]
self.assertListEqual(results, expected_results)
@slow
def test_bart_large_mask_filling(self):
pbase = pipeline(task="fill-mask", model="facebook/bart-large")
src_text = [" I went to the <mask>."]
results = [x["token_str"] for x in pbase(src_text)]
expected_results = ["Ġbathroom", "Ġgym", "Ġwrong", "Ġmovies", "Ġhospital"]
self.assertListEqual(results, expected_results)
@slow
def test_mnli_inference(self):