chore: import upstream snapshot with attribution

test_unstructured/staging/test_baseplate.py

@@ -0,0 +1,73 @@
+from test_unstructured.unit_utils import assign_hash_ids
+from unstructured.documents.coordinates import PixelSpace
+from unstructured.documents.elements import (
+    CoordinatesMetadata,
+    ElementMetadata,
+    NarrativeText,
+    Title,
+)
+from unstructured.staging.baseplate import stage_for_baseplate
+
+
+def test_stage_for_baseplate():
+    points = (
+        (545.0150947570801, 226.5191650390625),
+        (545.0150947570801, 254.7656043600921),
+        (704.879451751709, 254.7656043600921),
+        (704.879451751709, 226.5191650390625),
+    )
+    system = PixelSpace(width=1700, height=2200)
+    coordinates_metadata = CoordinatesMetadata(points=points, system=system)
+    metadata = ElementMetadata(filename="fox.pdf", coordinates=coordinates_metadata)
+    elements = assign_hash_ids(
+        [
+            Title("A Wonderful Story About A Fox", metadata=metadata),
+            NarrativeText(
+                "A fox ran into the chicken coop and the chickens flew off!", metadata=metadata
+            ),
+        ]
+    )
+
+    rows = stage_for_baseplate(elements)
+    assert rows == {
+        "rows": [
+            {
+                "data": {
+                    "element_id": "933d9fce18f44b09f4ec6975f470a0d7",
+                    "text": "A Wonderful Story About A Fox",
+                    "type": "Title",
+                },
+                "metadata": {
+                    "filename": "fox.pdf",
+                    "coordinates_points": (
+                        (545.0150947570801, 226.5191650390625),
+                        (545.0150947570801, 254.7656043600921),
+                        (704.879451751709, 254.7656043600921),
+                        (704.879451751709, 226.5191650390625),
+                    ),
+                    "coordinates_system": "PixelSpace",
+                    "coordinates_layout_width": 1700,
+                    "coordinates_layout_height": 2200,
+                },
+            },
+            {
+                "data": {
+                    "element_id": "d154ec57da0d7d4439aaed8ec6546f6e",
+                    "text": "A fox ran into the chicken coop and the chickens flew off!",
+                    "type": "NarrativeText",
+                },
+                "metadata": {
+                    "filename": "fox.pdf",
+                    "coordinates_points": (
+                        (545.0150947570801, 226.5191650390625),
+                        (545.0150947570801, 254.7656043600921),
+                        (704.879451751709, 254.7656043600921),
+                        (704.879451751709, 226.5191650390625),
+                    ),
+                    "coordinates_system": "PixelSpace",
+                    "coordinates_layout_width": 1700,
+                    "coordinates_layout_height": 2200,
+                },
+            },
+        ],
+    }

test_unstructured/staging/test_datasaur.py

@@ -0,0 +1,58 @@
+import pytest
+
+from unstructured.documents.elements import Text
+from unstructured.staging import datasaur
+
+
+def test_stage_for_datasaur():
+    elements = [Text("Text 1"), Text("Text 2"), Text("Text 3")]
+    result = datasaur.stage_for_datasaur(elements)
+    assert result[0]["text"] == "Text 1"
+    assert result[0]["entities"] == []
+    assert result[1]["text"] == "Text 2"
+    assert result[1]["entities"] == []
+    assert result[2]["text"] == "Text 3"
+    assert result[2]["entities"] == []
+
+
+def test_stage_for_datasaur_with_entities():
+    elements = [Text("Text 1"), Text("Text 2"), Text("Text 3")]
+    entities = [[{"text": "Matt", "type": "PER", "start_idx": 11, "end_idx": 15}], [], []]
+
+    result = datasaur.stage_for_datasaur(elements, entities=entities)
+    assert result[0]["text"] == "Text 1"
+    assert result[0]["entities"] == entities[0]
+    assert result[1]["text"] == "Text 2"
+    assert result[1]["entities"] == entities[1]
+    assert result[2]["text"] == "Text 3"
+    assert result[2]["entities"] == entities[2]
+
+
+def test_datasaur_raises_with_missing_entity_text():
+    with pytest.raises(ValueError):
+        elements = [Text("Text 1"), Text("Text 2"), Text("Text 3")]
+        datasaur.stage_for_datasaur(elements, entities=[{"bad_key": "text"}])
+
+
+def test_datasaur_raises_with_missing_key():
+    entities = [[{"text": "Matt", "type": "PER", "start_idx": 11}], [], []]
+
+    with pytest.raises(ValueError):
+        elements = [Text("Text 1"), Text("Text 2"), Text("Text 3")]
+        datasaur.stage_for_datasaur(elements, entities=entities)
+
+
+def test_datasaur_raises_with_bad_type():
+    entities = [[{"text": "Matt", "type": "PER", "start_idx": 11, "end_idx": "15"}], [], []]
+
+    with pytest.raises(ValueError):
+        elements = [Text("Text 1"), Text("Text 2"), Text("Text 3")]
+        datasaur.stage_for_datasaur(elements, entities=entities)
+
+
+def test_datasaur_raises_with_wrong_length():
+    entities = [[{"text": "Matt", "type": "PER", "start_idx": 11, "end_idx": 15}], []]
+
+    with pytest.raises(ValueError):
+        elements = [Text("Text 1"), Text("Text 2"), Text("Text 3")]
+        datasaur.stage_for_datasaur(elements, entities=entities)

test_unstructured/staging/test_huggingface.py

@@ -0,0 +1,77 @@
+import pytest
+
+from unstructured.documents.elements import Text, Title
+from unstructured.staging import huggingface
+
+
+class MockTokenizer:
+    model_max_length = 20
+
+    def tokenize(self, text):
+        return text.split(" ")
+
+
+def test_stage_for_transformers():
+    title_element = (Title(text="Here is a wonderful story"),)
+    elements = [title_element, Text(text="hello " * 20 + "there " * 20)]
+
+    tokenizer = MockTokenizer()
+
+    chunk_elements = huggingface.stage_for_transformers(elements, tokenizer, buffer=10)
+
+    hello_chunk = Text(("hello " * 10).strip())
+    there_chunk = Text(("there " * 10).strip())
+
+    assert chunk_elements == [
+        title_element,
+        hello_chunk,
+        hello_chunk,
+        there_chunk,
+        there_chunk,
+    ]
+
+
+def test_chunk_by_attention_window():
+    text = "hello " * 20 + "there " * 20
+    tokenizer = MockTokenizer()
+    chunks = huggingface.chunk_by_attention_window(text, tokenizer, buffer=10)
+
+    hello_chunk = ("hello " * 10).strip()
+    there_chunk = ("there " * 10).strip()
+    assert chunks == [hello_chunk, hello_chunk, there_chunk, there_chunk]
+
+
+def test_chunk_by_attention_window_no_buffer():
+    text = "hello " * 20 + "there " * 20
+    tokenizer = MockTokenizer()
+    chunks = huggingface.chunk_by_attention_window(text, tokenizer, buffer=0)
+
+    hello_chunk = ("hello " * 20).strip()
+    there_chunk = ("there " * 20).strip()
+    assert chunks == [hello_chunk, there_chunk]
+
+
+def test_chunk_by_attention_window_raises_with_negative_buffer():
+    text = "hello " * 20 + "there " * 20
+    tokenizer = MockTokenizer()
+    with pytest.raises(ValueError):
+        huggingface.chunk_by_attention_window(text, tokenizer, buffer=-10)
+
+
+def test_chunk_by_attention_window_raises_if_buffer_too_big():
+    text = "hello " * 20 + "there " * 20
+    tokenizer = MockTokenizer()
+    with pytest.raises(ValueError):
+        # NOTE(robinson) - The buffer exceeds the max input size of 20
+        huggingface.chunk_by_attention_window(text, tokenizer, buffer=40)
+
+
+def test_chunk_by_attention_window_raises_if_chunk_exceeds_window():
+    text = "hello " * 100 + "."
+    tokenizer = MockTokenizer()
+    with pytest.raises(ValueError):
+
+        def split_function(text):
+            return text.split(".")
+
+        huggingface.chunk_by_attention_window(text, tokenizer, split_function=split_function)

test_unstructured/staging/test_label_box.py

@@ -0,0 +1,143 @@
+import os
+
+import pytest
+
+from unstructured.documents.elements import NarrativeText, Title
+from unstructured.staging import label_box
+
+
+@pytest.fixture()
+def elements():
+    return [Title(text="Title 1"), NarrativeText(text="Narrative 1")]
+
+
+@pytest.fixture()
+def output_directory(tmp_path):
+    return str(tmp_path)
+
+
+@pytest.fixture()
+def nonexistent_output_directory(tmp_path):
+    return os.path.join(str(tmp_path), "nonexistent_dir")
+
+
+@pytest.fixture()
+def url_prefix():
+    return "https://storage.googleapis.com/labelbox-sample-datasets/nlp"
+
+
+@pytest.mark.parametrize(
+    ("attachments", "raises_error"),
+    [
+        (
+            [
+                {"type": "RAW_TEXT", "value": "Description Text"},
+                {"type": "IMAGE", "value": "Image label", "ignored_value": 123},
+            ],
+            False,
+        ),
+        ([{"type": "INVALID_TYPE", "value": "Description Text"}], True),
+        ([{"type": "RAW_TEXT", "value": 1}], True),
+        ([{"type": "RAW_TEXT"}], True),
+        ([{"value": "My text label"}], True),
+    ],
+)
+def test_validate_attachments(attachments, raises_error):
+    if raises_error:
+        with pytest.raises(ValueError):
+            label_box._validate_attachments(attachments, 0)
+    else:
+        label_box._validate_attachments(attachments, 0)
+
+
+attachment = {"type": "RAW_TEXT", "value": "Text description."}
+
+
+@pytest.mark.parametrize(
+    (
+        (
+            "external_ids",
+            "attachments",
+            "output_directory_fixture",
+            "create_directory",
+            "raises",
+            "exception_class",
+        )
+    ),
+    [
+        (None, None, "output_directory", True, False, None),
+        (["id1", "id2"], None, "output_directory", True, False, None),
+        (["id1"], None, "output_directory", True, True, ValueError),
+        (None, [[attachment], [attachment]], "output_directory", True, False, None),
+        (None, [[attachment]], "output_directory", True, True, ValueError),
+        (["id1", "id2"], [[attachment] * 2, [attachment]], "output_directory", True, False, None),
+        (
+            ["id1", "id2"],
+            [[attachment] * 2, [attachment]],
+            "nonexistent_output_directory",
+            True,
+            False,
+            None,
+        ),
+        (
+            ["id1", "id2"],
+            [[attachment] * 2, [attachment]],
+            "nonexistent_output_directory",
+            False,
+            True,
+            FileNotFoundError,
+        ),
+    ],
+)
+def test_stage_for_label_box(
+    elements,
+    url_prefix,
+    external_ids,
+    attachments,
+    output_directory_fixture,
+    create_directory,
+    raises,
+    exception_class,
+    request,
+):
+    output_directory = request.getfixturevalue(output_directory_fixture)
+    if raises:
+        with pytest.raises(exception_class):
+            label_box.stage_for_label_box(
+                elements,
+                output_directory,
+                url_prefix,
+                external_ids=external_ids,
+                attachments=attachments,
+                create_directory=create_directory,
+            )
+    else:
+        config = label_box.stage_for_label_box(
+            elements,
+            output_directory,
+            url_prefix,
+            external_ids=external_ids,
+            attachments=attachments,
+            create_directory=create_directory,
+        )
+        assert len(config) == len(elements)
+        for index, (element_config, element) in enumerate(zip(config, elements)):
+            print(element_config)
+
+            if external_ids:
+                assert element_config["externalId"] == external_ids[index]
+            else:
+                assert element_config["externalId"] == element.id
+
+            if attachments:
+                assert element_config["attachments"] == [
+                    {"type": attachment["type"], "value": attachment["value"]}
+                    for attachment in attachments[index]
+                ]
+
+            assert element_config["data"].startswith(url_prefix)
+            assert element_config["data"].endswith(f"{element_config['externalId']}.txt")
+
+            output_filepath = os.path.join(output_directory, f"{element_config['externalId']}.txt")
+            with open(output_filepath) as data_file:
+                assert data_file.read().strip() == element.text.strip()

test_unstructured/staging/test_label_studio.py

@@ -0,0 +1,309 @@
+from __future__ import annotations
+
+import pytest
+
+from test_unstructured.unit_utils import assign_hash_ids
+from unstructured.documents.elements import Element, NarrativeText, Title
+from unstructured.staging import label_studio
+
+
+@pytest.fixture()
+def elements():
+    return [Title(text="Title 1"), NarrativeText(text="Narrative 1")]
+
+
+def test_convert_to_label_studio_data(elements: list[Element]):
+    label_studio_data = label_studio.stage_for_label_studio(elements)
+
+    assert label_studio_data[0]["data"]["text"] == "Title 1"
+    assert "ref_id" in label_studio_data[0]["data"]
+
+    assert label_studio_data[1]["data"]["text"] == "Narrative 1"
+    assert "ref_id" in label_studio_data[1]["data"]
+
+
+def test_specify_text_name(elements: list[Element]):
+    label_studio_data = label_studio.stage_for_label_studio(elements, text_field="random_text")
+    assert "random_text" in label_studio_data[0]["data"]
+    assert label_studio_data[0]["data"]["random_text"] == "Title 1"
+
+
+def test_specify_id_name(elements: list[Element]):
+    label_studio_data = label_studio.stage_for_label_studio(elements, id_field="random_id")
+    assert "random_id" in label_studio_data[0]["data"]
+
+
+def test_created_annotation():
+    annotation = label_studio.LabelStudioAnnotation(
+        result=[
+            label_studio.LabelStudioResult(
+                type="choices",
+                value={"choices": ["Positive"]},
+                from_name="sentiment",
+                to_name="text",
+            ),
+        ],
+    )
+
+    annotation.to_dict() == {
+        "result": [
+            {
+                "type": "choices",
+                "value": {"choices": ["Positive"]},
+                "from_name": "sentiment",
+                "id": None,
+                "to_name": "text",
+                "hidden": False,
+                "read_only": False,
+            },
+        ],
+        "was_canceled": False,
+    }
+
+
+@pytest.mark.parametrize(
+    ("score", "raises", "exception"),
+    [
+        (None, True, ValueError),
+        (-0.25, True, ValueError),
+        (0, False, None),
+        (0.5, False, None),
+        (1, False, None),
+        (1.25, True, ValueError),
+    ],
+)
+def test_init_prediction(score: float | None, raises: bool, exception: Exception | None):
+    result = [
+        label_studio.LabelStudioResult(
+            type="choices",
+            value={"choices": ["Positive"]},
+            from_name="sentiment",
+            to_name="text",
+        ),
+    ]
+
+    if raises:
+        with pytest.raises(exception):
+            label_studio.LabelStudioPrediction(result=result, score=score)
+    else:
+        prediction = label_studio.LabelStudioPrediction(result=result, score=score)
+        prediction.to_dict() == {
+            "result": [
+                {
+                    "type": "choices",
+                    "value": {"choices": ["Positive"]},
+                    "from_name": "sentiment",
+                    "id": None,
+                    "to_name": "text",
+                    "hidden": False,
+                    "read_only": False,
+                },
+            ],
+            "was_canceled": False,
+            "score": score,
+        }
+
+
+def test_stage_with_annotation():
+    elements = assign_hash_ids([NarrativeText(text="A big brown bear")])
+    annotations = [
+        label_studio.LabelStudioAnnotation(
+            result=[
+                label_studio.LabelStudioResult(
+                    type="choices",
+                    value={"choices": ["Positive"]},
+                    from_name="sentiment",
+                    to_name="text",
+                ),
+            ],
+        ),
+    ]
+    label_studio_data = label_studio.stage_for_label_studio(elements, [annotations])
+    assert label_studio_data == [
+        {
+            "data": {"text": "A big brown bear", "ref_id": "2812a3676591a479c5425789f9c0156f"},
+            "annotations": [
+                {
+                    "result": [
+                        {
+                            "type": "choices",
+                            "value": {"choices": ["Positive"]},
+                            "from_name": "sentiment",
+                            "id": None,
+                            "to_name": "text",
+                            "hidden": False,
+                            "read_only": False,
+                        },
+                    ],
+                    "was_canceled": False,
+                },
+            ],
+        },
+    ]
+
+
+def test_stage_with_prediction():
+    elements = assign_hash_ids([NarrativeText(text="A big brown bear")])
+
+    prediction = [
+        label_studio.LabelStudioPrediction(
+            result=[
+                label_studio.LabelStudioResult(
+                    type="choices",
+                    value={"choices": ["Positive"]},
+                    from_name="sentiment",
+                    to_name="text",
+                ),
+            ],
+            score=0.98,
+        ),
+    ]
+    label_studio_data = label_studio.stage_for_label_studio(elements, predictions=[prediction])
+    assert label_studio_data == [
+        {
+            "data": {"text": "A big brown bear", "ref_id": "2812a3676591a479c5425789f9c0156f"},
+            "predictions": [
+                {
+                    "result": [
+                        {
+                            "type": "choices",
+                            "value": {"choices": ["Positive"]},
+                            "from_name": "sentiment",
+                            "id": None,
+                            "to_name": "text",
+                            "hidden": False,
+                            "read_only": False,
+                        },
+                    ],
+                    "was_canceled": False,
+                    "score": 0.98,
+                },
+            ],
+        },
+    ]
+
+
+def test_stage_with_annotation_for_ner():
+    elements = assign_hash_ids([NarrativeText(text="A big brown bear")])
+
+    annotations = [
+        label_studio.LabelStudioAnnotation(
+            result=[
+                label_studio.LabelStudioResult(
+                    type="labels",
+                    value={"start": 12, "end": 16, "text": "bear", "labels": ["PER"]},
+                    from_name="label",
+                    to_name="text",
+                ),
+            ],
+        ),
+    ]
+    label_studio_data = label_studio.stage_for_label_studio(elements, [annotations])
+    assert label_studio_data == [
+        {
+            "data": {"text": "A big brown bear", "ref_id": "2812a3676591a479c5425789f9c0156f"},
+            "annotations": [
+                {
+                    "result": [
+                        {
+                            "type": "labels",
+                            "value": {"start": 12, "end": 16, "text": "bear", "labels": ["PER"]},
+                            "from_name": "label",
+                            "id": None,
+                            "to_name": "text",
+                            "hidden": False,
+                            "read_only": False,
+                        },
+                    ],
+                    "was_canceled": False,
+                },
+            ],
+        },
+    ]
+
+
+def test_stage_with_annotation_raises_with_mismatched_lengths():
+    element = NarrativeText(text="A big brown bear")
+    annotations = [
+        label_studio.LabelStudioAnnotation(
+            result=[
+                label_studio.LabelStudioResult(
+                    type="choices",
+                    value={"choices": ["Positive"]},
+                    from_name="sentiment",
+                    to_name="text",
+                ),
+            ],
+        ),
+    ]
+    with pytest.raises(ValueError):
+        label_studio.stage_for_label_studio([element], [annotations, annotations])
+
+
+def test_stage_with_prediction_raises_with_mismatched_lengths():
+    element = NarrativeText(text="A big brown bear")
+    prediction = [
+        label_studio.LabelStudioPrediction(
+            result=[
+                label_studio.LabelStudioResult(
+                    type="choices",
+                    value={"choices": ["Positive"]},
+                    from_name="sentiment",
+                    to_name="text",
+                ),
+            ],
+            score=0.82,
+        ),
+    ]
+    with pytest.raises(ValueError):
+        label_studio.stage_for_label_studio([element], predictions=[prediction, prediction])
+
+
+def test_stage_with_annotation_raises_with_invalid_type():
+    with pytest.raises(ValueError):
+        label_studio.LabelStudioResult(
+            type="bears",
+            value={"bears": ["Positive"]},
+            from_name="sentiment",
+            to_name="text",
+        )
+
+
+def test_stage_with_reviewed_annotation():
+    elements = assign_hash_ids([NarrativeText(text="A big brown bear")])
+    annotations = [
+        label_studio.LabelStudioAnnotation(
+            result=[
+                label_studio.LabelStudioResult(
+                    type="choices",
+                    value={"choices": ["Positive"]},
+                    from_name="sentiment",
+                    to_name="text",
+                ),
+            ],
+            reviews=[label_studio.LabelStudioReview(created_by={"user_id": 1}, accepted=True)],
+        ),
+    ]
+    label_studio_data = label_studio.stage_for_label_studio(elements, [annotations])
+    assert label_studio_data == [
+        {
+            "data": {"text": "A big brown bear", "ref_id": "2812a3676591a479c5425789f9c0156f"},
+            "annotations": [
+                {
+                    "result": [
+                        {
+                            "type": "choices",
+                            "value": {"choices": ["Positive"]},
+                            "from_name": "sentiment",
+                            "to_name": "text",
+                            "id": None,
+                            "hidden": False,
+                            "read_only": False,
+                        },
+                    ],
+                    "reviews": [{"created_by": {"user_id": 1}, "accepted": True, "id": None}],
+                    "was_canceled": False,
+                },
+            ],
+        },
+    ]

test_unstructured/staging/test_prodigy.py

@@ -0,0 +1,119 @@
+import csv
+import os
+
+import pytest
+
+from unstructured.documents.elements import NarrativeText, Title
+from unstructured.staging import prodigy
+
+
+@pytest.fixture()
+def elements():
+    return [Title(text="Title 1"), NarrativeText(text="Narrative 1")]
+
+
+@pytest.fixture()
+def valid_metadata():
+    return [{"score": 0.1}, {"category": "paragraph"}]
+
+
+@pytest.fixture()
+def metadata_with_id():
+    return [{"score": 0.1}, {"id": 1, "category": "paragraph"}]
+
+
+@pytest.fixture()
+def metadata_with_invalid_length():
+    return [{"score": 0.1}, {"category": "paragraph"}, {"type": "text"}]
+
+
+@pytest.fixture()
+def output_csv_file(tmp_path):
+    return os.path.join(tmp_path, "prodigy_data.csv")
+
+
+def test_validate_prodigy_metadata(elements):
+    validated_metadata = prodigy._validate_prodigy_metadata(elements, metadata=None)
+    assert len(validated_metadata) == len(elements)
+    assert all(not data for data in validated_metadata)
+
+
+def test_validate_prodigy_metadata_with_valid_metadata(elements, valid_metadata):
+    validated_metadata = prodigy._validate_prodigy_metadata(elements, metadata=valid_metadata)
+    assert len(validated_metadata) == len(elements)
+
+
+@pytest.mark.parametrize(
+    ("invalid_metadata_fixture", "exception_message"),
+    [
+        ("metadata_with_id", 'The key "id" is not allowed with metadata parameter at index: 1'),
+        (
+            "metadata_with_invalid_length",
+            "The length of the metadata parameter does not match with"
+            " the length of the elements parameter.",
+        ),
+    ],
+)
+def test_validate_prodigy_metadata_with_invalid_metadata(
+    elements,
+    invalid_metadata_fixture,
+    exception_message,
+    request,
+):
+    invalid_metadata = request.getfixturevalue(invalid_metadata_fixture)
+    with pytest.raises(ValueError) as validation_exception:
+        prodigy._validate_prodigy_metadata(elements, invalid_metadata)
+    assert str(validation_exception.value) == exception_message
+
+
+def test_convert_to_prodigy_data(elements):
+    prodigy_data = prodigy.stage_for_prodigy(elements)
+
+    assert len(prodigy_data) == len(elements)
+
+    assert prodigy_data[0]["text"] == "Title 1"
+    assert "meta" in prodigy_data[0]
+    assert "id" in prodigy_data[0]["meta"]
+    assert prodigy_data[0]["meta"]["id"] == elements[0].id
+
+    assert prodigy_data[1]["text"] == "Narrative 1"
+    assert "meta" in prodigy_data[1]
+    assert "id" in prodigy_data[1]["meta"]
+    assert prodigy_data[1]["meta"]["id"] == elements[1].id
+
+
+def test_convert_to_prodigy_data_with_valid_metadata(elements, valid_metadata):
+    prodigy_data = prodigy.stage_for_prodigy(elements, valid_metadata)
+
+    assert len(prodigy_data) == len(elements)
+
+    assert prodigy_data[0]["text"] == "Title 1"
+    assert "meta" in prodigy_data[0]
+    assert prodigy_data[0]["meta"] == {"id": elements[0].id, **valid_metadata[0]}
+
+    assert prodigy_data[1]["text"] == "Narrative 1"
+    assert "meta" in prodigy_data[1]
+    assert prodigy_data[1]["meta"] == {"id": elements[1].id, **valid_metadata[1]}
+
+
+def test_stage_csv_for_prodigy(elements, output_csv_file):
+    with open(output_csv_file, "w+") as csv_file:
+        prodigy_csv_string = prodigy.stage_csv_for_prodigy(elements)
+        csv_file.write(prodigy_csv_string)
+
+    fieldnames = ["text", "id"]
+    with open(output_csv_file) as csv_file:
+        csv_rows = csv.DictReader(csv_file)
+        assert all(set(row.keys()) == set(fieldnames) for row in csv_rows)
+
+
+def test_stage_csv_for_prodigy_with_metadata(elements, valid_metadata, output_csv_file):
+    with open(output_csv_file, "w+") as csv_file:
+        prodigy_csv_string = prodigy.stage_csv_for_prodigy(elements, valid_metadata)
+        csv_file.write(prodigy_csv_string)
+
+    fieldnames = {"text", "id"}.union(*(data.keys() for data in valid_metadata))
+    fieldnames = [fieldname.lower() for fieldname in fieldnames]
+    with open(output_csv_file) as csv_file:
+        csv_rows = csv.DictReader(csv_file)
+        assert all(set(row.keys()) == set(fieldnames) for row in csv_rows)

test_unstructured/staging/test_weaviate.py

@@ -0,0 +1,66 @@
+import contextlib
+import json
+import os
+
+import pytest
+
+# NOTE(robinson) - allows tests that do not require the weaviate client to
+# run for the docker container
+with contextlib.suppress(ModuleNotFoundError):
+    import weaviate
+
+from unstructured.partition.json import partition_json
+from unstructured.staging.weaviate import (
+    create_unstructured_weaviate_class,
+    stage_for_weaviate,
+)
+
+is_in_docker = os.path.exists("/.dockerenv")
+is_in_ci = os.getenv("CI", "").lower() not in {"", "false", "f", "0"}
+
+
+def test_stage_for_weaviate():
+    element_dict = {
+        "element_id": "015301d4f56aa4b20ec10ac889d2343f",
+        "text": "LayoutParser: A Unified Toolkit for Deep Learning Based Document Image Analysis",
+        "type": "Title",
+        "metadata": {
+            "filename": "layout-parser-paper-fast.pdf",
+            "filetype": "application/json",
+            "page_number": 1,
+            "coordinates": {
+                "points": (
+                    (157.62199999999999, 114.23496279999995),
+                    (157.62199999999999, 146.5141628),
+                    (457.7358962799999, 146.5141628),
+                    (457.7358962799999, 114.23496279999995),
+                ),
+                "system": "PixelSpace",
+                "layout_width": 324,
+                "layout_height": 450,
+            },
+        },
+    }
+
+    elements = partition_json(text=json.dumps([element_dict]))
+    data = stage_for_weaviate(elements)
+    assert data[0] == {
+        "filename": "layout-parser-paper-fast.pdf",
+        "filetype": "application/json",
+        "page_number": 1,
+        "text": "LayoutParser: A Unified Toolkit for Deep Learning Based Document Image Analysis",
+        "category": "Title",
+    }
+
+
+@pytest.mark.skipif(not is_in_ci, reason="Integration test that depends on having secret keys")
+@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
+def test_weaviate_schema_is_valid():
+    unstructured_class = create_unstructured_weaviate_class()
+    class_name = unstructured_class["class"]
+    client = weaviate.connect_to_embedded()
+    try:
+        client.collections.delete(class_name)
+        client.collections.create_from_dict(unstructured_class)
+    finally:
+        client.close()