mirror of
https://github.com/huggingface/transformers.git
synced 2025-07-15 18:48:24 +06:00
908a28894c
* Add new token classification example * Remove txt file * Add test * With actual testing done * Less warmup is better * Update examples/token-classification/run_ner_new.py Co-authored-by: Thomas Wolf <thomwolf@users.noreply.github.com> * Address review comments * Fix test * Make Lysandre happy * Last touches and rename * Rename in tests * Address review comments * More run_ner -> run_ner_old Co-authored-by: Thomas Wolf <thomwolf@users.noreply.github.com>
56 lines
1.8 KiB
Python
56 lines
1.8 KiB
Python
import logging
|
|
import sys
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import run_ner_old as run_ner
|
|
from transformers.testing_utils import slow
|
|
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
|
|
logger = logging.getLogger()
|
|
|
|
|
|
class ExamplesTests(unittest.TestCase):
|
|
@slow
|
|
def test_run_ner(self):
|
|
stream_handler = logging.StreamHandler(sys.stdout)
|
|
logger.addHandler(stream_handler)
|
|
|
|
testargs = """
|
|
--model_name distilbert-base-german-cased
|
|
--output_dir ./tests/fixtures/tests_samples/temp_dir
|
|
--overwrite_output_dir
|
|
--data_dir ./tests/fixtures/tests_samples/GermEval
|
|
--labels ./tests/fixtures/tests_samples/GermEval/labels.txt
|
|
--max_seq_length 128
|
|
--num_train_epochs 6
|
|
--logging_steps 1
|
|
--do_train
|
|
--do_eval
|
|
""".split()
|
|
with patch.object(sys, "argv", ["run.py"] + testargs):
|
|
result = run_ner.main()
|
|
self.assertLess(result["eval_loss"], 1.5)
|
|
|
|
def test_run_ner_pl(self):
|
|
stream_handler = logging.StreamHandler(sys.stdout)
|
|
logger.addHandler(stream_handler)
|
|
|
|
testargs = """
|
|
--model_name distilbert-base-german-cased
|
|
--output_dir ./tests/fixtures/tests_samples/temp_dir
|
|
--overwrite_output_dir
|
|
--data_dir ./tests/fixtures/tests_samples/GermEval
|
|
--labels ./tests/fixtures/tests_samples/GermEval/labels.txt
|
|
--max_seq_length 128
|
|
--num_train_epochs 6
|
|
--logging_steps 1
|
|
--do_train
|
|
--do_eval
|
|
""".split()
|
|
with patch.object(sys, "argv", ["run.py"] + testargs):
|
|
result = run_ner.main()
|
|
self.assertLess(result["eval_loss"], 1.5)
|