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

[testing] add timeout_decorator (#3543)

Browse Source
This commit is contained in:
Sam Shleifer
2020-05-01 09:05:47 -04:00
committed by GitHub
parent b8686174be
commit 18db92dd9a
3 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
setup.cfg
View File

@@ -23,6 +23,7 @@ known_third_party =
tensorboardX tensorboardX
tensorflow tensorflow
tensorflow_datasets tensorflow_datasets
timeout_decorator
torch torch
torchtext torchtext
torchvision torchvision

2
setup.py
View File

@@ -72,7 +72,7 @@ extras["torch"] = ["torch==1.4.0"]
extras["serving"] = ["pydantic", "uvicorn", "fastapi", "starlette"] extras["serving"] = ["pydantic", "uvicorn", "fastapi", "starlette"]
extras["all"] = extras["serving"] + ["tensorflow", "torch"] extras["all"] = extras["serving"] + ["tensorflow", "torch"]
extras["testing"] = ["pytest", "pytest-xdist"] extras["testing"] = ["pytest", "pytest-xdist", "timeout-decorator"]
extras["docs"] = ["recommonmark", "sphinx", "sphinx-markdown-tables", "sphinx-rtd-theme"] extras["docs"] = ["recommonmark", "sphinx", "sphinx-markdown-tables", "sphinx-rtd-theme"]
extras["quality"] = [ extras["quality"] = [
"black", "black",

4
tests/test_modeling_bart.py
View File

@@ -13,10 +13,11 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import tempfile import tempfile
import unittest import unittest
import timeout_decorator # noqa
from transformers import is_torch_available from transformers import is_torch_available
from .test_configuration_common import ConfigTester from .test_configuration_common import ConfigTester
@@ -357,6 +358,7 @@ class BartHeadTests(unittest.TestCase):
loss = outputs[0] loss = outputs[0]
self.assertIsInstance(loss.item(), float) self.assertIsInstance(loss.item(), float)
@timeout_decorator.timeout(1)
def test_lm_forward(self): def test_lm_forward(self):
config, input_ids, batch_size = self._get_config_and_data() config, input_ids, batch_size = self._get_config_and_data()
lm_labels = ids_tensor([batch_size, input_ids.shape[1]], self.vocab_size).to(torch_device) lm_labels = ids_tensor([batch_size, input_ids.shape[1]], self.vocab_size).to(torch_device)