Skip to content

Commit

Permalink
Implemented test for dataset creation through from_documents() method
Browse files Browse the repository at this point in the history
  • Loading branch information
kai-car committed Aug 9, 2024
1 parent e134555 commit 83fffc1
Showing 1 changed file with 17 additions and 0 deletions.
17 changes: 17 additions & 0 deletions tests/unit/core/test_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import numpy
import pytest
import torch
from pytorch_ie import Document
from pytorch_ie.annotations import BinaryRelation, Label, LabeledSpan, Span
from pytorch_ie.core import AnnotationList, annotation_field
from pytorch_ie.core.taskmodule import (
Expand Down Expand Up @@ -431,3 +432,19 @@ def test_dataset_with_taskmodule(

for document in train_dataset:
assert not document["entities"].predictions


def test_pie_dataset_from_documents(documents):
assert len(documents) == 8
assert all(isinstance(doc, TextBasedDocument) for doc in documents)

dataset_from_documents = Dataset.from_documents(documents)

assert isinstance(dataset_from_documents, Dataset)

assert len(dataset_from_documents) == 8
assert all(isinstance(doc, TextBasedDocument) for doc in dataset_from_documents)
assert all(doc1.id == doc2.id for doc1, doc2 in zip(documents, dataset_from_documents))

assert hasattr(dataset_from_documents, "document_type")

0 comments on commit 83fffc1

Please sign in to comment.