Package: pytorch-ie
PyTorch-IE: State-of-the-art Information Extraction in PyTorch🤯 What's this about?This is an experimental framework that aims to combine the lessons learned from five years of information extraction research.Focus on the core task:The main goal is to develop information extraction methods not dataset loading and evaluation logic. We use external well-maintained libraries for non-core functionality. PyTorch-Lightning for training and logging, Huggingface datasets for dataset reading, and Huggingface evaluate for evaluation (coming soon).Sharing is caring:Being able to quickly and easily share models is key to promote your work and facilitate further research. All models developed in PyTorch-IE can be easily shared via the Huggingface model hub. This further allows to quickly build demos based on Huggingface spaces, gradio or streamlit.Unified document format:A unified document format allows for quick experimentation on any dataset or task.Beyond sentence level:Most information extraction frameworks assume text inputs at a sentence granularity. We do not make any assumption on the granularity but generally aim for document-level information extraction.Beyond unstructured text:Unstructured text is only one possible area for information extraction. We developed the framework to also support information extraction from semi-structured text (e.g. HTML), two-dimensional text (e.g. OCR'd images), and images.Character-level annotation and evaluation:Many information extraction frameworks annotate and evaluate on a token level. We believe that annotation and evaluation should be done on a character level as this also considers the suitability of the tokenizer for the task.Make no assumptions on the structure of models:The last years have seen many different and creative approaches to information extraction and a framework that imposes a structure on those will most certainly be to limiting. With PyTorch-iE you have full control over how a document is prepared for a model and how the model is structured. The logic is self-contained and thus can be easily shared and inspected by others. The only assumption we make is that the input is a document and the output are targets (training) or annotations (inference).🔭 DemosTaskLinkNamed Entity Recognition (Span-based)Joint Named Entity Recognition and Relation Classification🚀️ Quickstart$pipinstallpytorch-ieFor even faster prototyping with pre-defined, but fully configurable training pipelines and much more useful tooling, have a look into thePyTorch-IE-Hydra-Template.🥧 Concepts & ArchitecturePyTorch-IE builds on three core concepts: the📃 Document, the🔤 ⇔ 🔢 Taskmodule, and the🧮 Model. In a
nutshell, the Document says how your data is structured, the Model defines your trainable logic and the Taskmodule
converts from one end to the other. All three concepts are represented as abstract classes that should be used to
derive use-case specific versions. In the following, they are explained in detail.📃 DocumentTheDocumentclass is a specialdataclassthat defines the document model. Derivations can contain several
elements:Data fieldslike strings to represent one or multiple texts or arrays for image data. These elements can be
arbitrary python objects.Annotation fieldslike labeled spans for entities or labeled tuples of spans for relations. These elements have
to be of a certain container typeAnnotationListthat is dynamically typed with the actual annotation type, e.g.entities: AnnotationList[LabeledSpan]. Furthermore, annotation elements define one or multiple annotationtargets.
An annotation target is either a data element or another annotation container. Internally, targets are used to construct the
annotation graph, i.e. data elements and annotation containers are the nodes and targets define the edges. The
annotation graph defines the (de-)serialization order and what is accessible from within an annotation. To
facilitate the setup of annotation containers, there is theannotation_field()method.Other fieldsto save metadata, ids, etc. They are not constrained in any way, but can not be accessed from within
annotations.Example Document ModelfromtypingimportOptionalfrompytorch_ie.coreimportDocument,AnnotationLayer,annotation_fieldfrompytorch_ie.annotationsimportLabeledSpan,BinaryRelation,LabelclassMyDocument(Document):# data fields (any field that is targeted by an annotation fields)text:str# annotation fieldsentities:AnnotationLayer[LabeledSpan]=annotation_field(target="text")relations:AnnotationLayer[BinaryRelation]=annotation_field(target="entities")label:AnnotationLayer[Label]=annotation_field()# other fieldsdoc_id:Optional[str]=NoneNote that thelabelis a special annotation field that does not define a target because it belongs to the whole document.
You can also have more complex constructs, like annotation fields that target multiple other fields by usingannotation_field(targets)orannotation_field(named_targets). The latter is useful if you want to access the
targets by name from within the annotation, see below for an example.AnnotationsThere are several predefinedannotation typesinpytorch_ie.annotations, however, feel free to define your own.
Annotations have to be dataclasses that subclasspytorch_ie.core.Annotation. They also need to be hashable and
immutable. The following is a simple example:@dataclass(eq=True,frozen=True)classSimpleLabeledSpan(Annotation):start:intend:intlabel:strAccessing Target ContentWe can expand the above example a little to have a nice string representation:@dataclass(eq=True,frozen=True)classLabeledSpan(Annotation):start:intend:intlabel:strdef__str__(self)->str:ifself.targetsisNone:return""returnstr(self.target[self.start:self.end])The content ofself.targetis lazily assigned as soon as the annotation is added to a document.Note that this now expects a singlecollections.abc.Sequenceastarget, e.g.:my_spans:AnnotationLayer[Span]=annotation_field(target="<NAME_OF_THE_SEQUENCE_FIELD>")If we have multiple targets, we need to define target names to access them. For this, we need to set the special
fieldTARGET_NAMES:@dataclass(eq=True,frozen=True)classAlignment(Annotation):TARGET_NAMES=("text1","text2")start1:intend1:intstart2:intend2:intdef__str__(self)->str:ifself.targetsisNone:return""# we can access the `named_targets` which has the keys defined in `TARGET_NAMES`span1=self.named_targets["text1"][self.start1:self.end1]span2=self.named_targets["text2"][self.start2:self.end2]returnf'span1="{span1}" is aligned with span2="{span2}"'This requires to define the annotation container as follows:classMyDocumentWithAlignment(Document):text_a:strtext_b:str# `named_targets` defines the mapping from `TARGET_NAMES` to data fieldsmy_alignments:AnnotationLayer[Alignment]=annotation_field(named_targets={"text1":"text_a","text2":"text_b"})Note thattext1andtext2can also target the same field.(De-)Serialization of AnnotationsAs usual for dataclasses, annotations can be converted to json like objects with.asdict(). However, they can be
also created withMyAnnotation.fromdict(dct, annotation_store). Both methods are required because documents and
their annotations are created on the fly when working with PIE datasets (see below).🔤 ⇔ 🔢 TaskmoduleThe taskmodule is responsible for converting documents to model inputs and back. For that purpose, it requires the
user to implement the following methods:encode_input: Taking one document, create one or multipleTaskEncodings. ATaskEncodingrepresents an
example that will be passed to the model later on. It is a container holdinginputs, optionaltargets, the
originaldocument, andmetadata. Note thatencode_inputshould not assign a value totargets.encode_target: This gets a singleTaskEncodingand should produce a target encoding that will be assigned
totargetslater on. As such, it is called only during training / evaluation, but not for inference. Note that,
this is allowed to return None. In this case, the respectiveTaskEncodingwill not be passed to the model at all.collate: Taking a batch ofTaskEncodings, this should produce a batch input for the model. Note that this has to
work with available targets (training and evaluation) and without them (inference).unbatch_output: This gets a batch output from the model and should rearrange that into a sequence ofTaskOutputs.
In that means it can be understood as the opposite tocollate. The number ofTaskOutputs should match the
number ofTaskEncodings that got into the batch because we align them later on for easy creation of new annotations.create_annotations_from_output: This gets a singleTaskEncodingwith its correspondingTaskOutputand
should yield tuples each consisting of an annotation field name and an annotation. The annotations will be added
as predictions to the annotation field with the respective name.prepare(OPTIONAL): This will get the train dataset, i.e. a Sequence or Iterable of Documents, and can be used
to calculate additional parameters like the list of all available labels, etc.You can find some predefined taskmodules fortext-andtoken classification,text classification based relation
extraction,joint entity and relation classificationand other use cases in the packagepytorch_ie.taskmodules. Especially, have a look at theSimpleTransformerTextClassificationTaskModulethat is well documented and should provide a good starting point to implement your own one.🧮 ModelPyTorch-IE models are meant to do the heavy lifting training and inference. They arePytorch-Lightning modules,
enhanced with some functionality to ease persisting them, seeReusability and Sharing.You can find some predefined models for transformer basedtext-andtoken classification,sequence generation,
and other use cases in the packagepytorch_ie.models.Reusability and SharingTaskmodules and Models provide some functionality to ease reusability and reproducibility. Especially, they provide
the methodssave_pretrained()andfrom_pretrained()that can be used to save their specification, i.e. theirconfig, and available model wights to disc and exactly re-create them again from that data.Huggingface Hub and Extended ConfigsThese methods come along
with integration to theHuggingface Hub. By passingpush_to_hub=Truetosave_pretrained(), the taskmodule / model is directly pushed to the Hub and can be loaded again with the respective
identifier (see theExamplesfor how to do so). However, to work properly, each taskmodule / model has to
correctly implement the_config()getter method. Per default, it returns all parameters passed to the__init__method if this callssave_hyperparameters()which is very recommended. But you may have created some further
parameters that should be persisted, for instance a label-to-id mapping. In this case,_config()should be
overwritten to take this into account:def_config(self)->Dict[str,Any]:# add the label-to-id mapping to the configconfig=super()._config()config["label_to_id"]=self.label_to_idreturnconfigFurthermore, you can use the propertyis_from_pretrainedto know if the taskmodule / model is just loaded or created
from scratch. This may be useful, for instance, to avoid downloading a model from Huggingface Transformers when you
in fact want to load your own trained model from disc viafrom_pretrained:fromtransformersimportAutoConfig,AutoModelhf_config=AutoConfig.from_pretrained(model_name_or_path)# If this is already trained, just create an empty transformer model. The weights are loaded afterwards# via the pytorch_ie.Model.from_pretrained() logic.ifself.is_from_pretrained:self.model=AutoModel.from_config(config=hf_config)# Otherwise, download the whole model from the Huggingface Hub.else:self.model=AutoModel.from_pretrained(model_name_or_path,config=hf_config)In short, each taskmodule / model implementation should:callsave_hyperparameters()in__init__to save all constructor arguments,pass remaining__init__kwargs (keyword arguments) to its super to not break some other helpful functionality
(e.g.is_from_pretrained), andoverwrite_config()if additional parameters are calculated, e.g. from the dataset.⚡️ Examples: PredictionThe following examples work out of the box. No further setup like manually downloading a model is needed!Note:Settingnum_workers=0in the pipeline is only necessary when running an example in an
interactive python session. The reason is that multiprocessing doesn't play well with the interactive python
interpreter, seeherefor details.Span-classification-based Named Entity Recognitionfromdataclassesimportdataclassfrompytorch_ie.annotationsimportLabeledSpanfrompytorch_ie.autoimportAutoPipelinefrompytorch_ie.coreimportAnnotationLayer,annotation_fieldfrompytorch_ie.documentsimportTextDocument@dataclassclassExampleDocument(TextDocument):entities:AnnotationLayer[LabeledSpan]=annotation_field(target="text")document=ExampleDocument("“Making a super tasty alt-chicken wing is only half of it,” said Po Bronson, general partner at SOSV and managing director of IndieBio.")# see below for the long versionner_pipeline=AutoPipeline.from_pretrained("pie/example-ner-spanclf-conll03",device=-1,num_workers=0)ner_pipeline(document)forentityindocument.entities.predictions:print(f"{entity}->{entity.label}")# Result:# IndieBio -> ORG# Po Bronson -> PER# SOSV -> ORGTo create the same pipeline as above without `AutoPipeline`frompytorch_ie.autoimportAutoTaskModule,AutoModelfrompytorch_ie.pipelineimportPipelinemodel_name_or_path="pie/example-ner-spanclf-conll03"ner_taskmodule=AutoTaskModule.from_pretrained(model_name_or_path)ner_model=AutoModel.from_pretrained(model_name_or_path)ner_pipeline=Pipeline(model=ner_model,taskmodule=ner_taskmodule,device=-1,num_workers=0)Or, without `Auto` classes at allfrompytorch_ie.pipelineimportPipelinefrompytorch_ie.modelsimportTransformerSpanClassificationModelfrompytorch_ie.taskmodulesimportTransformerSpanClassificationTaskModulemodel_name_or_path="pie/example-ner-spanclf-conll03"ner_taskmodule=TransformerSpanClassificationTaskModule.from_pretrained(model_name_or_path)ner_model=TransformerSpanClassificationModel.from_pretrained(model_name_or_path)ner_pipeline=Pipeline(model=ner_model,taskmodule=ner_taskmodule,device=-1,num_workers=0)Text-classification-based Relation Extractionfromdataclassesimportdataclassfrompytorch_ie.annotationsimportBinaryRelation,LabeledSpanfrompytorch_ie.autoimportAutoPipelinefrompytorch_ie.coreimportAnnotationLayer,annotation_fieldfrompytorch_ie.documentsimportTextDocument@dataclassclassExampleDocument(TextDocument):entities:AnnotationLayer[LabeledSpan]=annotation_field(target="text")relations:AnnotationLayer[BinaryRelation]=annotation_field(target="entities")document=ExampleDocument("“Making a super tasty alt-chicken wing is only half of it,” said Po Bronson, general partner at SOSV and managing director of IndieBio.")re_pipeline=AutoPipeline.from_pretrained("pie/example-re-textclf-tacred",device=-1,num_workers=0)forstart,end,labelin[(65,75,"PER"),(96,100,"ORG"),(126,134,"ORG")]:document.entities.append(LabeledSpan(start=start,end=end,label=label))re_pipeline(document,batch_size=2)forrelationindocument.relations.predictions:print(f"({relation.head}->{relation.tail}) ->{relation.label}")# Result:# (Po Bronson -> SOSV) -> per:employee_of# (Po Bronson -> IndieBio) -> per:employee_of# (SOSV -> Po Bronson) -> org:top_members/employees# (IndieBio -> Po Bronson) -> org:top_members/employees⚡️ Examples: TrainingSpan-classification-based Named Entity Recognitionimportpytorch_lightningasplfrompytorch_lightning.callbacksimportModelCheckpointfromtorch.utils.dataimportDataLoaderimportdatasetsfrompytorch_ie.models.transformer_span_classificationimportTransformerSpanClassificationModelfrompytorch_ie.taskmodules.transformer_span_classificationimport(TransformerSpanClassificationTaskModule,)pl.seed_everything(42)model_output_path="./model_output/"model_name="bert-base-cased"num_epochs=10batch_size=32# Get the PIE dataset consisting of PIE Documents that will be used for training (and evaluation).# IMPORTANT: This requires pie-datasets >=0.3.0 to be installed! See here for further information:# https://github.com/ArneBinder/pie-datasetsdataset=datasets.load_dataset(path="pie/conll2003",)train_docs,val_docs=dataset["train"],dataset["validation"]print("train docs: ",len(train_docs))print("val docs: ",len(val_docs))# Create a PIE taskmodule.task_module=TransformerSpanClassificationTaskModule(tokenizer_name_or_path=model_name,entity_annotation="entities",max_length=128,)# Prepare the taskmodule with the training data. This may collect available labels etc.# The result of this should affect the state of the taskmodule config which will be# persisted (and can be loaded) later on.task_module.prepare(train_docs)# Persist the taskmodule. This writes the taskmodule config as a json file into the# model_output_path directory. The config contains all constructor parameters to# re-create the taskmodule at this state (via AutoTaskmodule.from_pretrained(model_output_path)).task_module.save_pretrained(model_output_path)# Use the taskmodule to encode the train and dev sets. This may use the text and# available annotations of the documents.train_dataset=task_module.encode(train_docs,encode_target=True,as_dataset=True)val_dataset=task_module.encode(val_docs,encode_target=True,as_dataset=True)# Create the dataloaders. Note that the taskmodule provides the collate function!train_dataloader=DataLoader(train_dataset,batch_size=batch_size,shuffle=True,collate_fn=task_module.collate,)val_dataloader=DataLoader(val_dataset,batch_size=batch_size,shuffle=False,collate_fn=task_module.collate,)# Create the PIE model. Note that we use the number of entries in the previously# collected label_to_id mapping to set the number of classes to predict.model=TransformerSpanClassificationModel(model_name_or_path=model_name,num_classes=len(task_module.label_to_id),t_total=len(train_dataloader)*num_epochs,learning_rate=1e-4,)# Optionally, set up a model checkpoint callback. See here for further information:# https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.callbacks.ModelCheckpoint.html# checkpoint_callback = ModelCheckpoint(#     monitor="val/f1",#     dirpath=model_output_path,#     filename="zs-ner-{epoch:02d}-val_f1-{val/f1:.2f}",#     save_top_k=1,#     mode="max",#     auto_insert_metric_name=False,#     save_weights_only=True,# )# Create the pytorch-lightning trainer. See here for further information:# https://pytorch-lightning.readthedocs.io/en/latest/api/pytorch_lightning.trainer.trainer.Trainer.htmltrainer=pl.Trainer(fast_dev_run=False,max_epochs=num_epochs,gpus=0,enable_checkpointing=False,# callbacks=[checkpoint_callback],precision=32,)# Start the training.trainer.fit(model,train_dataloader,val_dataloader)# Persist the trained model. This will save the model weights and the model config that allows# to re-create the model at this state (via AutoModel.from_pretrained(model_output_path)).# model.save_pretrained(model_output_path)📚 DatasetsPyTorch-IE works quite well together with Huggingface datasets. Have a look atpie-datasetsfor helpful tooling and a collection of datasets
that are already converted to the PIE format.🔧 Project DevelopmentSetupThis project is build withPoetry. If not yet done, follow thePoetry installation guideto install it. Then, install all
dependencies (including for development) for PyTorch-IE by calling the following command from the root of the
repository:poetryinstall--withdevNOTE: If the installation gets stuck, try if disabling experimental parallel installer helps
(source):poetry config experimental.new-installer falseTesting and code quality checksWe useNoxto execute any tests and code quality tooling in a reproducible way.To get a list of available toolchains, call:poetryrunnox-lTo run a specific command from that list, call:poetryrunnox-s<command>Note: To run the nox commands in the same, reproducible setup that is specified by the lock file, we call them viapoetry run <nox-command>.For instance, to run static type checking withmypy, call:poetryrunnox-smypy-3.9To run all commands that also run on GitHub CI, call:poetryrunnoxYou can also start a shell with Poetry's virtual environment activated by calling:poetryshellThis allows you to run above commands without thepoetry runprefix.Updating DependenciesCall this to update individual packages:poetryupdate<package>Then, commit the modified lock file to persist the state.ReleasingSince this project is based on theCookiecutter template, we can followtheir release steps:Create the release branch:git switch --create release mainIncrease the version:poetry version <PATCH|MINOR|MAJOR>,
e.g.poetry version patchfor a patch releaseCommit the changes:git commit --message="release <NEW VERSION>" pyproject.toml,
e.g.git commit --message="release 0.13.0" pyproject.tomlPush the changes to GitHub:git push origin releaseCreate a PR for thatreleasebranch on GitHub.Wait until checks passed successfully.Integrate 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.🏅 AcknowledgementsThis package is based on thesourcery-ai/python-best-practices-cookiecutterandcjolowicz/cookiecutter-hypermodern-pythonproject templates.📃 CitationIf you find the framework useful please consider citing it:@misc{alt2022pytorchie,author={Christoph Alt, Arne Binder},title={PyTorch-IE: State-of-the-art Information Extraction in PyTorch},year={2022},publisher={GitHub},journal={GitHub repository},howpublished={\url{https://github.com/ChristophAlt/pytorch-ie}}}