# -*- coding: utf-8 -*-
from setuptools import setup

package_dir = \
{'': 'src'}

packages = \
['pytorch_ie',
 'pytorch_ie.core',
 'pytorch_ie.data',
 'pytorch_ie.data.datamodules',
 'pytorch_ie.data.datasets',
 'pytorch_ie.data.datasets.hf_datasets',
 'pytorch_ie.models',
 'pytorch_ie.models.genre',
 'pytorch_ie.models.modules',
 'pytorch_ie.taskmodules',
 'pytorch_ie.utils']

package_data = \
{'': ['*']}

install_requires = \
['datasets>=2.4.0,<3.0.0',
 'huggingface-hub>=0.5.1,<0.6.0',
 'pytorch-lightning>=1.6.1,<2.0.0',
 'torchmetrics>=0.8.0,<0.9.0',
 'transformers>=4.18.0,<5.0.0']

setup_kwargs = {
    'name': 'pytorch-ie',
    'version': '0.11.0',
    'description': 'State-of-the-art Information Extraction in PyTorch',
    'long_description': '# PyTorch-IE: State-of-the-art Information Extraction in PyTorch\n\n[![PyPI](https://img.shields.io/pypi/v/pytorch-ie.svg)][pypi status]\n[![Status](https://img.shields.io/pypi/status/pytorch-ie.svg)][pypi status]\n[![Python Version](https://img.shields.io/pypi/pyversions/pytorch-ie)][pypi status]\n[![License](https://img.shields.io/pypi/l/pytorch-ie)][license]\n\n[![Read the documentation at https://pytorch-ie.readthedocs.io/](https://img.shields.io/readthedocs/pytorch-ie/latest.svg?label=Read%20the%20Docs)][read the docs]\n[![Tests](https://github.com/christophalt/pytorch-ie/workflows/Tests/badge.svg)][tests]\n[![Codecov](https://codecov.io/gh/christophalt/pytorch-ie/branch/main/graph/badge.svg)][codecov]\n\n[![pre-commit](https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit&logoColor=white)][pre-commit]\n[![Black](https://img.shields.io/badge/code%20style-black-000000.svg)][black]\n\n[pypi status]: https://pypi.org/project/pytorch-ie/\n[read the docs]: https://pytorch-ie.readthedocs.io/\n[tests]: https://github.com/christophalt/pytorch-ie/actions?workflow=Tests\n[codecov]: https://app.codecov.io/gh/christophalt/pytorch-ie\n[pre-commit]: https://github.com/pre-commit/pre-commit\n[black]: https://github.com/psf/black\n\n## 🤯 What\'s this about?\n\nThis is an experimental framework that aims to combine the lessons learned from five years of information extraction research.\n\n-   **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).\n-   **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.\n-   **Unified document format:** A unified document format allows for quick experimentation on any dataset or task.\n-   **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.\n-   **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.\n-   **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.\n-   **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).\n\n## 🔭 Demos\n\n| Task                                                       | Link                                                                                                                                                                  |\n| ---------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Named Entity Recognition (Span-based)                      | [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/pie/NER)                               |\n| Joint Named Entity Recognition and Relation Classification | [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/pie/Joint-NER-and-Relation-Extraction) |\n\n## 🚀️ Quickstart\n\n```console\n$ pip install pytorch-ie\n```\n\n## ⚡️ Examples: Prediction\n\n**The following examples work out of the box. No further setup like manually downloading a model is needed!**\n\n**Note:** Setting `num_workers=0` in the pipeline is only necessary when running an example in an\ninteractive python session. The reason is that multiprocessing doesn\'t play well with the interactive python\ninterpreter, see [here](https://docs.python.org/3/library/multiprocessing.html#using-a-pool-of-workers)\nfor details.\n\n### Span-classification-based Named Entity Recognition\n\n```python\nfrom dataclasses import dataclass\n\nfrom pytorch_ie.annotations import LabeledSpan\nfrom pytorch_ie.auto import AutoPipeline\nfrom pytorch_ie.core import AnnotationList, annotation_field\nfrom pytorch_ie.documents import TextDocument\n\n@dataclass\nclass ExampleDocument(TextDocument):\n    entities: AnnotationList[LabeledSpan] = annotation_field(target="text")\n\ndocument = ExampleDocument(\n    "“Making a super tasty alt-chicken wing is only half of it,” said Po Bronson, general partner at SOSV and managing director of IndieBio."\n)\n\n# see below for the long version\nner_pipeline = AutoPipeline.from_pretrained("pie/example-ner-spanclf-conll03", device=-1, num_workers=0)\n\nner_pipeline(document)\n\nfor entity in document.entities.predictions:\n    print(f"{entity} -> {entity.label}")\n\n# Result:\n# IndieBio -> ORG\n# Po Bronson -> PER\n# SOSV -> ORG\n```\n\n<details>\n<summary>\nTo create the same pipeline as above without `AutoPipeline`\n</summary>\n\n```python\nfrom pytorch_ie.auto import AutoTaskModule, AutoModel\nfrom pytorch_ie.pipeline import Pipeline\n\nmodel_name_or_path = "pie/example-ner-spanclf-conll03"\nner_taskmodule = AutoTaskModule.from_pretrained(model_name_or_path)\nner_model = AutoModel.from_pretrained(model_name_or_path)\nner_pipeline = Pipeline(model=ner_model, taskmodule=ner_taskmodule, device=-1, num_workers=0)\n```\n\n</details>\n\n<details>\n<summary>\nOr, without `Auto` classes at all\n</summary>\n\n```python\nfrom pytorch_ie.pipeline import Pipeline\nfrom pytorch_ie.models import TransformerSpanClassificationModel\nfrom pytorch_ie.taskmodules import TransformerSpanClassificationTaskModule\n\nmodel_name_or_path = "pie/example-ner-spanclf-conll03"\nner_taskmodule = TransformerSpanClassificationTaskModule.from_pretrained(model_name_or_path)\nner_model = TransformerSpanClassificationModel.from_pretrained(model_name_or_path)\nner_pipeline = Pipeline(model=ner_model, taskmodule=ner_taskmodule, device=-1, num_workers=0)\n```\n\n</details>\n\n<details>\n<summary>\n\n### Text-classification-based Relation Extraction\n\n</summary>\n\n```python\nfrom dataclasses import dataclass\n\nfrom pytorch_ie.annotations import BinaryRelation, LabeledSpan\nfrom pytorch_ie.auto import AutoPipeline\nfrom pytorch_ie.core import AnnotationList, annotation_field\nfrom pytorch_ie.documents import TextDocument\n\n\n@dataclass\nclass ExampleDocument(TextDocument):\n    entities: AnnotationList[LabeledSpan] = annotation_field(target="text")\n    relations: AnnotationList[BinaryRelation] = annotation_field(target="entities")\n\ndocument = ExampleDocument(\n    "“Making a super tasty alt-chicken wing is only half of it,” said Po Bronson, general partner at SOSV and managing director of IndieBio."\n)\n\nre_pipeline = AutoPipeline.from_pretrained("pie/example-re-textclf-tacred", device=-1, num_workers=0)\n\nfor start, end, label in [(65, 75, "PER"), (96, 100, "ORG"), (126, 134, "ORG")]:\n    document.entities.append(LabeledSpan(start=start, end=end, label=label))\n\nre_pipeline(document, batch_size=2)\n\nfor relation in document.relations.predictions:\n    print(f"({relation.head} -> {relation.tail}) -> {relation.label}")\n\n# Result:\n# (Po Bronson -> SOSV) -> per:employee_of\n# (Po Bronson -> IndieBio) -> per:employee_of\n# (SOSV -> Po Bronson) -> org:top_members/employees\n# (IndieBio -> Po Bronson) -> org:top_members/employees\n```\n\n</details>\n\n## ⚡️ Examples: Training\n\n<details>\n\n<summary>\n\n### Span-classification-based Named Entity Recognition\n\n</summary>\n\n```python\nimport pytorch_lightning as pl\nfrom pytorch_lightning.callbacks import ModelCheckpoint\nfrom torch.utils.data import DataLoader\n\nimport datasets\nfrom pytorch_ie.models.transformer_span_classification import TransformerSpanClassificationModel\nfrom pytorch_ie.taskmodules.transformer_span_classification import (\n    TransformerSpanClassificationTaskModule,\n)\n\npl.seed_everything(42)\n\nmodel_output_path = "./model_output/"\nmodel_name = "bert-base-cased"\nnum_epochs = 10\nbatch_size = 32\n\n# Get the PIE dataset consisting of PIE Documents that will be used for training (and evaluation).\ndataset = datasets.load_dataset(\n    path="pie/conll2003",\n)\ntrain_docs, val_docs = dataset["train"], dataset["validation"]\n\nprint("train docs: ", len(train_docs))\nprint("val docs: ", len(val_docs))\n\n# Create a PIE taskmodule.\ntask_module = TransformerSpanClassificationTaskModule(\n    tokenizer_name_or_path=model_name,\n    max_length=128,\n)\n\n# Prepare the taskmodule with the training data. This may collect available labels etc.\n# The result of this should affect the state of the taskmodule config which will be\n# persisted (and can be loaded) later on.\ntask_module.prepare(train_docs)\n\n# Persist the taskmodule. This writes the taskmodule config as a json file into the\n# model_output_path directory. The config contains all constructor parameters to\n# re-create the taskmodule at this state (via AutoTaskmodule.from_pretrained(model_output_path)).\ntask_module.save_pretrained(model_output_path)\n\n# Use the taskmodule to encode the train and dev sets. This may use the text and\n# available annotations of the documents.\ntrain_dataset = task_module.encode(train_docs, encode_target=True, as_dataset=True)\nval_dataset = task_module.encode(val_docs, encode_target=True, as_dataset=True)\n\n# Create the dataloaders. Note that the taskmodule provides the collate function!\ntrain_dataloader = DataLoader(\n    train_dataset,\n    batch_size=batch_size,\n    shuffle=True,\n    collate_fn=task_module.collate,\n)\n\nval_dataloader = DataLoader(\n    val_dataset,\n    batch_size=batch_size,\n    shuffle=False,\n    collate_fn=task_module.collate,\n)\n\n# Create the PIE model. Note that we use the number of entries in the previously\n# collected label_to_id mapping to set the number of classes to predict.\nmodel = TransformerSpanClassificationModel(\n    model_name_or_path=model_name,\n    num_classes=len(task_module.label_to_id),\n    t_total=len(train_dataloader) * num_epochs,\n    learning_rate=1e-4,\n)\n\n# Optionally, set up a model checkpoint callback. See here for further information:\n# https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.callbacks.ModelCheckpoint.html\n# checkpoint_callback = ModelCheckpoint(\n#     monitor="val/f1",\n#     dirpath=model_output_path,\n#     filename="zs-ner-{epoch:02d}-val_f1-{val/f1:.2f}",\n#     save_top_k=1,\n#     mode="max",\n#     auto_insert_metric_name=False,\n#     save_weights_only=True,\n# )\n\n# Create the pytorch-lightning trainer. See here for further information:\n# https://pytorch-lightning.readthedocs.io/en/latest/api/pytorch_lightning.trainer.trainer.Trainer.html\ntrainer = pl.Trainer(\n    fast_dev_run=False,\n    max_epochs=num_epochs,\n    gpus=0,\n    checkpoint_callback=False,\n    # callbacks=[checkpoint_callback],\n    precision=32,\n)\n# Start the training.\ntrainer.fit(model, train_dataloader, val_dataloader)\n\n# Persist the trained model. This will save the model weights and the model config that allows\n# to re-create the model at this state (via AutoModel.from_pretrained(model_output_path)).\n# model.save_pretrained(model_output_path)\n```\n\n</details>\n\n## 📚 Datasets\n\nWe parse all datasets into a common format that can be loaded directly from the model hub via Huggingface 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.\n\n```python\nimport datasets\n\ndataset = datasets.load_dataset("pie/conll2003")\n\nprint(dataset["train"][0])\n# >>> CoNLL2003Document(text=\'EU rejects German call to boycott British lamb .\', id=\'0\', metadata={})\n\ndataset["train"][0].entities\n# >>> AnnotationList([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)])\n\nentity = dataset["train"][0].entities[1]\n\nprint(f"[{entity.start}, {entity.end}] {entity}")\n# >>> [11, 17] German\n```\n\n<details>\n<summary><b>How to create your own Pytorch-IE dataset</b></summary>\n\nPyTorch-IE datasets are built on top of Huggingface datasets. For instance, consider the\n[conll2003 from the Huggingface Hub](https://huggingface.co/datasets/conll2003) and especially their respective\n[dataset loading script](https://huggingface.co/datasets/conll2003/blob/main/conll2003.py). To create a PyTorch-IE\ndataset from that, you have to implement:\n\n1. A Document class. This will be the type of individual dataset examples.\n\n```python\n@dataclass\nclass CoNLL2003Document(TextDocument):\n    entities: AnnotationList[LabeledSpan] = annotation_field(target="text")\n```\n\nHere we derive from `TextDocument` that has a simple `text` string as base annotation target. The `CoNLL2003Document`\nadds one single annotation list called `entities` that consists of `LabeledSpan`s which reference the `text` field of\nthe document. You can add further annotation types by adding `AnnotationList` fields that may also reference (i.e.\n`target`) other annotations as you like. See [\'pytorch_ie.annotations`](src/pytorch_ie/annotations.py) for predefined\nannotation types.\n\n2. A dataset config. This is similar to\n   [creating a Huggingface dataset config](https://huggingface.co/docs/datasets/dataset_script#multiple-configurations).\n\n```python\nclass CoNLL2003Config(datasets.BuilderConfig):\n    """BuilderConfig for CoNLL2003"""\n\n    def __init__(self, **kwargs):\n        """BuilderConfig for CoNLL2003.\n        Args:\n          **kwargs: keyword arguments forwarded to super.\n        """\n        super().__init__(**kwargs)\n```\n\n3. A dataset builder class. This should inherit from\n   [`pytorch_ie.data.builder.GeneratorBasedBuilder`](src/pytorch_ie/data/builder.py) which is a wrapper around the\n   [Huggingface dataset builder class](https://huggingface.co/docs/datasets/v2.4.0/en/package_reference/builder_classes#datasets.GeneratorBasedBuilder)\n   with some utility functionality to work with PyTorch-IE `Documents`. The key elements to implement are: `DOCUMENT_TYPE`,\n   `BASE_DATASET_PATH`, and `_generate_document`.\n\n```python\nclass Conll2003(pytorch_ie.data.builder.GeneratorBasedBuilder):\n    # Specify the document type. This will be the class of individual dataset examples.\n    DOCUMENT_TYPE = CoNLL2003Document\n\n    # The Huggingface identifier that points to the base dataset. This may be any string that works\n    # as path with Huggingface `datasets.load_dataset`.\n    BASE_DATASET_PATH = "conll2003"\n\n    # The builder configs, see https://huggingface.co/docs/datasets/dataset_script for further information.\n    BUILDER_CONFIGS = [\n        CoNLL2003Config(\n            name="conll2003", version=datasets.Version("1.0.0"), description="CoNLL2003 dataset"\n        ),\n    ]\n\n    # [Optional] Define additional keyword arguments which will be passed to `_generate_document` below.\n    def _generate_document_kwargs(self, dataset):\n        return {"int_to_str": dataset.features["ner_tags"].feature.int2str}\n\n    # Define how a Pytorch-IE Document will be created from a Huggingface dataset example.\n    def _generate_document(self, example, int_to_str):\n        doc_id = example["id"]\n        tokens = example["tokens"]\n        ner_tags = [int_to_str(tag) for tag in example["ner_tags"]]\n\n        text, ner_spans = tokens_and_tags_to_text_and_labeled_spans(tokens=tokens, tags=ner_tags)\n\n        document = CoNLL2003Document(text=text, id=doc_id)\n\n        for span in sorted(ner_spans, key=lambda span: span.start):\n            document.entities.append(span)\n\n        return document\n```\n\nThe full script can be found here: [datasets/conll2003/conll2003.py](datasets/conll2003/conll2003.py). Note, that to\nload the dataset with `datasets.load_dataset`, the script has to be located in a directory with the same name (as it\nis the case for standard Huggingface dataset loading scripts).\n\n</details>\n\n<!-- github-only -->\n\n✨📚✨ [Read the full documentation](https://pytorch-ie.readthedocs.io/)\n\n## 🔧 Development Setup\n\n## 🏅 Acknowledgements\n\n-   This package is based on the [sourcery-ai/python-best-practices-cookiecutter](https://github.com/sourcery-ai/python-best-practices-cookiecutter) and [cjolowicz/cookiecutter-hypermodern-python](https://github.com/cjolowicz/cookiecutter-hypermodern-python) project templates.\n\n## 📃 Citation\n\nIf you find the framework useful please consider citing it:\n\n```bibtex\n@misc{alt2022pytorchie,\n    author={Christoph Alt, Arne Binder},\n    title = {PyTorch-IE: State-of-the-art Information Extraction in PyTorch},\n    year = {2022},\n    publisher = {GitHub},\n    journal = {GitHub repository},\n    howpublished = {\\url{https://github.com/ChristophAlt/pytorch-ie}}\n}\n```\n\n[license]: https://github.com/christophalt/pytorch-ie/blob/main/LICENSE\n',
    'author': 'Christoph Alt',
    'author_email': 'christoph.alt@posteo.de',
    'maintainer': None,
    'maintainer_email': None,
    'url': 'https://github.com/christophalt/pytorch-ie',
    'package_dir': package_dir,
    'packages': packages,
    'package_data': package_data,
    'install_requires': install_requires,
    'python_requires': '>=3.9,<4.0',
}


setup(**setup_kwargs)
