fastNLP/tests/io/pipe/test_classification.py

import pytest
import os

from fastNLP.io import DataBundle
from fastNLP.io.pipe.classification import SSTPipe, SST2Pipe, IMDBPipe, YelpFullPipe, YelpPolarityPipe, \
    AGsNewsPipe, DBPediaPipe
from fastNLP.io.pipe.classification import ChnSentiCorpPipe, THUCNewsPipe, WeiboSenti100kPipe


@pytest.mark.skipif('download' not in os.environ, reason="Skip download")
class TestClassificationPipe:
    def test_process_from_file(self):
        for pipe in [YelpPolarityPipe, SST2Pipe, IMDBPipe, YelpFullPipe, SSTPipe]:
            print(pipe)
            data_bundle = pipe(tokenizer='raw').process_from_file()
            print(data_bundle)


class TestRunPipe:
    def test_load(self):
        for pipe in [IMDBPipe]:
            data_bundle = pipe(tokenizer='raw').process_from_file('tests/data_for_tests/io/imdb')
            print(data_bundle)


@pytest.mark.skipif('download' not in os.environ, reason="Skip download")
class TestCNClassificationPipe:
    def test_process_from_file(self):
        for pipe in [ChnSentiCorpPipe]:
            data_bundle = pipe(bigrams=True, trigrams=True).process_from_file()
            print(data_bundle)


@pytest.mark.skipif('download' not in os.environ, reason="Skip download")
class TestRunClassificationPipe:
    def test_process_from_file(self):
        data_set_dict = {
            'yelp.p': ('tests/data_for_tests/io/yelp_review_polarity', YelpPolarityPipe,
                       {'train': 6, 'dev': 6, 'test': 6}, {'words': 1176, 'target': 2},
                       False),
            'yelp.f': ('tests/data_for_tests/io/yelp_review_full', YelpFullPipe,
                       {'train': 6, 'dev': 6, 'test': 6}, {'words': 1166, 'target': 5},
                       False),
            'sst-2': ('tests/data_for_tests/io/SST-2', SST2Pipe,
                      {'train': 5, 'dev': 5, 'test': 5}, {'words': 139, 'target': 2},
                      True),
            'sst': ('tests/data_for_tests/io/SST', SSTPipe,
                    {'train': 354, 'dev': 6, 'test': 6}, {'words': 232, 'target': 5},
                    False),
            'imdb': ('tests/data_for_tests/io/imdb', IMDBPipe,
                     {'train': 6, 'dev': 6, 'test': 6}, {'words': 1670, 'target': 2},
                     False),
            'ag': ('tests/data_for_tests/io/ag', AGsNewsPipe,
                   {'train': 4, 'test': 5}, {'words': 257, 'target': 4},
                   False),
            'dbpedia': ('tests/data_for_tests/io/dbpedia', DBPediaPipe,
                        {'train': 14, 'test': 5}, {'words': 496, 'target': 14},
                        False),
            'ChnSentiCorp': ('tests/data_for_tests/io/ChnSentiCorp', ChnSentiCorpPipe,
                             {'train': 6, 'dev': 6, 'test': 6},
                             {'chars': 529, 'bigrams': 1296, 'trigrams': 1483, 'target': 2},
                             False),
            'Chn-THUCNews': ('tests/data_for_tests/io/THUCNews', THUCNewsPipe,
                             {'train': 9, 'dev': 9, 'test': 9}, {'chars': 1864, 'target': 9},
                             False),
            'Chn-WeiboSenti100k': ('tests/data_for_tests/io/WeiboSenti100k', WeiboSenti100kPipe,
                                   {'train': 6, 'dev': 6, 'test': 7}, {'chars': 452, 'target': 2},
                                   False),
        }
        for k, v in data_set_dict.items():
            path, pipe, data_set, vocab, warns = v
            if 'Chn' not in k:
                if warns:
                    data_bundle = pipe(tokenizer='raw').process_from_file(path)
                else:
                    data_bundle = pipe(tokenizer='raw').process_from_file(path)
            else:
                data_bundle = pipe(bigrams=True, trigrams=True).process_from_file(path)

            assert(isinstance(data_bundle, DataBundle))
            assert(len(data_set) == data_bundle.num_dataset)
            for name, dataset in data_bundle.iter_datasets():
                assert(name in data_set.keys())
                assert(data_set[name] == len(dataset))

            assert(len(vocab) == data_bundle.num_vocab)
            for name, vocabs in data_bundle.iter_vocabs():
                assert(name in vocab.keys())
                assert(vocab[name] == len(vocabs))