Package: pie-datasets
pie-datasetsDataset building scripts and utilities forPyTorch-IE. We parse all datasets into a common format that can be
loaded directly from the Huggingface Hub. Taking advantage ofHuggingface datasets, the documents are cached in an arrow table and
serialized / deserialized on the fly. Any changes or preprocessing applied to the documents will be cached as well.Setuppipinstallpie-datasetsTo install the latest version from GitHub:pipinstallgit+https://git@github.com/ArneBinder/pie-datasets.gitAvailable datasetsSeeherefor a list of available datasets. Note, that you can easily add your own
datasets by following theinstructions below.UsageGeneralfrompie_datasetsimportload_dataset# load the dataset from https://huggingface.co/datasets/pie/conll2003dataset=load_dataset("pie/conll2003")print(dataset["train"][0])# >>> CoNLL2003Document(text='EU rejects German call to boycott British lamb .', id='0', metadata={})dataset["train"][0].entities# >>> AnnotationLayer([LabeledSpan(start=0, end=2, label='ORG', score=1.0), LabeledSpan(start=11, end=17, label='MISC', score=1.0), LabeledSpan(start=34, end=41, label='MISC', score=1.0)])entity=dataset["train"][0].entities[1]print(f"[{entity.start},{entity.end}]{entity}")# >>> [11, 17] German{name:len(split)forname,splitindataset.items()}# >>> {'train': 14041, 'validation': 3250, 'test': 3453}Adjusting splitsSimilar toHuggingface datasets, you can adjust the splits of a dataset in
various ways. Here are some examples:frompie_datasetsimportload_datasetdataset=load_dataset("pie/conll2003")# re-create a validation split from train split concatenated with the original validation splitdataset_with_new_val=dataset.concat_splits(["train","validation"],target="train").add_test_split(source_split="train",target_split="my_validation",test_size=0.2,seed=42){name:len(split)forname,splitindataset_with_new_val.items()}# >>> {'test': 3453, 'train': 13832, 'my_validation': 3459}# drop the test splitdataset_without_test=dataset_with_new_val.drop_splits(["test"]){name:len(split)forname,splitindataset_without_test.items()}# >>> {'train': 13832, 'my_validation': 3459}Adjusting dataset entriesCallingmapon the dataset will apply the given function to all its documents. Internally, that relies
ondatasets.Dataset.map.
Thus, the function can be any function that takes a document as input and returns a document as output. If the
function returns a different document type, you need to provide it asresult_document_typeargument tomap. Note, thatthe result is cached for each split, so that re-running the same function on the
same dataset will be a no-op.Example where the function returns the same document type:frompie_datasetsimportload_datasetdefduplicate_entities(document):new_document=document.copy()forentityindocument.entities:# we need to copy the entity because each annotation can only be part of one documentnew_document.entities.append(entity.copy())returnnew_documentdataset=load_dataset("pie/conll2003")len(dataset["train"][0].entities)# >>> 3converted_dataset=dataset.map(duplicate_entities)# Map: 100%|██████████| 14041/14041 [00:02<00:00, 4697.18 examples/s]# Map: 100%|██████████| 3250/3250 [00:00<00:00, 4583.95 examples/s]# Map: 100%|██████████| 3453/3453 [00:00<00:00, 4614.67 examples/s]len(converted_dataset["train"][0].entities)# >>> 6Example where the function returns a different document type:fromdataclassesimportdataclassfrompytorch_ie.coreimportAnnotationLayer,annotation_fieldfrompytorch_ie.documentsimportTextBasedDocumentfrompytorch_ie.annotationsimportLabeledSpan,Spanfrompie_datasetsimportload_dataset@dataclassclassCoNLL2003DocumentWithWords(TextBasedDocument):entities:AnnotationLayer[LabeledSpan]=annotation_field(target="text")words:AnnotationLayer[Span]=annotation_field(target="text")defadd_words(document)->CoNLL2003DocumentWithWords:new_document=CoNLL2003DocumentWithWords(text=document.text,id=document.id)forentityindocument.entities:new_document.entities.append(entity.copy())start=0forwordindocument.text.split():word_start=document.text.index(word,start)word_annotation=Span(start=word_start,end=word_start+len(word))new_document.words.append(word_annotation)returnnew_documentdataset=load_dataset("pie/conll2003")dataset.document_type# >>> <class 'datasets_modules.datasets.pie--conll2003.821bfce48d2ebc3533db067c4d8e89396155c65cd311d2341a82acf81f561885.conll2003.CoNLL2003Document'>converted_dataset=dataset.map(add_words,result_document_type=CoNLL2003DocumentWithWords)# Map: 100%|██████████| 14041/14041 [00:03<00:00, 3902.00 examples/s]# Map: 100%|██████████| 3250/3250 [00:00<00:00, 3929.52 examples/s]# Map: 100%|██████████| 3453/3453 [00:00<00:00, 3947.49 examples/s]converted_dataset.document_type# >>> <class '__main__.CoNLL2003DocumentWithWords'>converted_dataset["train"][0].words# >>> AnnotationLayer([Span(start=0, end=2), Span(start=3, end=10), Span(start=11, end=17), Span(start=18, end=22), Span(start=23, end=25), Span(start=26, end=33), Span(start=34, end=41), Span(start=42, end=46), Span(start=47, end=48)])[str(word)forwordinconverted_dataset["train"][0].words]# >>> ['EU', 'rejects', 'German', 'call', 'to', 'boycott', 'British', 'lamb', '.']We can alsoregister a document converterfor a specific document type. This will be used when callingto_document_typewith the respective document type. The following code will produce the same result
as the previous one:dataset=load_dataset("pie/conll2003")# Register add_words as a converter function for the target document type CoNLL2003DocumentWithWords.# Since add_words specifies the return type, we can omit the document type here.dataset.register_document_converter(add_words)# Determine the matching converter entry for the target document type and apply it with dataset.map.converted_dataset=dataset.to_document_type(CoNLL2003DocumentWithWords)Note, that some of the PIE datasets come with default document converters. For instance, thePIE conll2003 datasetcomes with one that converts
the dataset topytorch_ie.documents.TextDocumentWithLabeledSpans. These documents work with the
PIE taskmodules fortoken classificationandspan classificationout-of-the-box. The following code will load the dataset and convert it to the required document type:frompie_datasetsimportload_datasetfrompytorch_ie.taskmodulesimportTransformerTokenClassificationTaskModuletaskmodule=TransformerTokenClassificationTaskModule(tokenizer_name_or_path="bert-base-cased")# the taskmodule expects TextDocumentWithLabeledSpans as input and the conll2003 dataset comes with a# default converter for that document type. Thus, we can directly load the dataset and convert it.dataset=load_dataset("pie/conll2003").to_document_type(taskmodule.document_type)...How to create your own PIE datasetPIE datasets are built on top of Huggingface datasets. For instance, considerconll2003 at the Huggingface Huband especially its respectivedataset loading script. To create a PIE
dataset from that, you have to implement:A Document class. This will be the type of individual dataset examples.fromdataclassesimportdataclassfrompytorch_ie.annotationsimportLabeledSpanfrompytorch_ie.coreimportAnnotationLayer,annotation_fieldfrompytorch_ie.documentsimportTextBasedDocument@dataclassclassCoNLL2003Document(TextBasedDocument):entities:AnnotationLayer[LabeledSpan]=annotation_field(target="text")Here we derive fromTextBasedDocumentthat has a simpletextstring as base annotation target. TheCoNLL2003Documentadds one single annotation layer calledentitiesthat consists ofLabeledSpans which
reference thetextfield of the document. You can add further annotation types by addingAnnotationLayerfields that may also reference (i.e.target) other annotations as you like. The packagepytorch_ie.annotationscontains some predefined annotation types and the packagepytorch_ie.documentsdefines
some document types that you can use as base classes.A dataset config. This is similar tocreating a Huggingface dataset config.importdatasetsclassCoNLL2003Config(datasets.BuilderConfig):"""BuilderConfig for CoNLL2003"""def__init__(self,**kwargs):"""BuilderConfig for CoNLL2003.Args:**kwargs: keyword arguments forwarded to super."""super().__init__(**kwargs)A dataset builder class. This should inherit frompie_datasets.GeneratorBasedBuilderwhich is a wrapper around theHuggingface dataset builder classwith some utility functionality to work with PyTorch-IEDocuments. The key elements to implement are:DOCUMENT_TYPE,BASE_DATASET_PATH, and_generate_document.frompytorch_ie.documentsimportTextDocumentWithLabeledSpansfrompytorch_ie.utils.spanimporttokens_and_tags_to_text_and_labeled_spansfrompie_datasetsimportGeneratorBasedBuilderclassConll2003(GeneratorBasedBuilder):# Specify the document type. This will be the class of individual dataset examples.DOCUMENT_TYPE=CoNLL2003Document# The Huggingface identifier that points to the base dataset. This may be any string that works# as path with Huggingface `datasets.load_dataset`.BASE_DATASET_PATH="conll2003"# It is strongly recommended to also specify the revision (tag name, or branch name, or commit hash)# of the base dataset. This ensures that the dataset will not change unexpectedly when the base dataset# is updated.BASE_DATASET_REVISION="01ad4ad271976c5258b9ed9b910469a806ff3288"# The builder configs, see https://huggingface.co/docs/datasets/dataset_script for further information.BUILDER_CONFIGS=[CoNLL2003Config(name="conll2003",version=datasets.Version("1.0.0"),description="CoNLL2003 dataset"),]# [Optional] Define additional keyword arguments which will be passed to `_generate_document` below.def_generate_document_kwargs(self,dataset):return{"int_to_str":dataset.features["ner_tags"].feature.int2str}# Define how a Pytorch-IE Document will be created from a Huggingface dataset example.def_generate_document(self,example,int_to_str):doc_id=example["id"]tokens=example["tokens"]ner_tags=[int_to_str(tag)fortaginexample["ner_tags"]]text,ner_spans=tokens_and_tags_to_text_and_labeled_spans(tokens=tokens,tags=ner_tags)document=CoNLL2003Document(text=text,id=doc_id)forspaninsorted(ner_spans,key=lambdaspan:span.start):document.entities.append(span)returndocument# [OPTIONAL] Define how the dataset will be converted to a different document type. Here, we add a# converter for the generic document type `TextDocumentWithLabeledSpans` that is used by the PIE# taskmodules for token and span classification. This allows to directly call# `pie_datasets.load_dataset("pie/conll2003").to_document_type(TextDocumentWithLabeledSpans)`.DOCUMENT_CONVERTERS={TextDocumentWithLabeledSpans:{# if the converter is a simple dictionary, just rename the layer according that"entities":"labeled_spans",}}The full script can be found here:dataset_builders/pie/conll2003/conll2003.py. Note, that to
load the dataset withpie_datasets.load_dataset, the script has to be located in a directory with the same name
(as it is the case for standard Huggingface dataset loading scripts).DevelopmentSetupThis project is build withPoetry. See here forinstallation instructions.Get the code and switch into the project directory:gitclonehttps://github.com/ArneBinder/pie-datasetscdpie-datasetsCreate a virtual environment and install the dependencies:poetryinstallFinally, to run any of the below commands, you need to activate the virtual environment:poetryshellNote: You can also run commands in the virtual environment without activating it first:poetry run <command>.Code Formatting, Linting and Static Type Checkingpre-commitrun-aTestingrun all tests with coverage:pytest--cov--cov-reportterm-missingReleasingCreate the release branch:git switch --create release mainIncrease the project version:poetry version <PATCH|MINOR|MAJOR>,
e.g.poetry version patchfor a patch release. If the release contains new features, or breaking changes,
bump the minor version (this project has no main release yet). If the release contains only bugfixes, bump
the patch version. SeeSemantic Versioningfor more information.Increase the maximum version in all datasetrequirements.txtfiles, if not a patch release.Commit the changes:git commit --message="release <NEW VERSION>" -a,
e.g.git commit --message="release 0.13.0" -aPush the changes to GitHub:git push origin releaseCreate a PR for thatreleasebranch on GitHub.Wait until checks passed successfully.Merge the PR into the main branch. This triggers the GitHub Action that creates all relevant release
artefacts and also uploads them to PyPI.Cleanup: Delete thereleasebranch. This is important, because otherwise the next release will fail.Update the scripts athuggingface.co/piewith the new versions.