[bart-tiny-random] Put a 5MB model on S3 to allow faster exampl… (#3488)

This commit is contained in:
Sam Shleifer
2020-03-30 12:28:27 -04:00
committed by GitHub
parent 1f72865726
commit 8deff3acf2
3 changed files with 23 additions and 6 deletions

View File

@@ -27,7 +27,9 @@ from .utils import CACHE_DIR, require_torch, slow, torch_device
if is_torch_available():
import torch
from transformers import (
AutoModel,
AutoModelForSequenceClassification,
AutoTokenizer,
BartModel,
BartForConditionalGeneration,
BartForSequenceClassification,
@@ -183,6 +185,15 @@ class BARTModelTest(ModelTesterMixin, unittest.TestCase):
def test_inputs_embeds(self):
pass
def test_tiny_model(self):
model_name = "sshleifer/bart-tiny-random"
tiny = AutoModel.from_pretrained(model_name) # same vocab size
tok = AutoTokenizer.from_pretrained(model_name) # same tokenizer
inputs_dict = tok.batch_encode_plus(["Hello my friends"], return_tensors="pt")
with torch.no_grad():
tiny(**inputs_dict)
@require_torch
class BartHeadTests(unittest.TestCase):