HotpotQA

@synalinks_export("synalinks.datasets.hotpotqa.get_input_data_model")
def get_input_data_model():
    """
    Returns HotpotQA input data model for pipeline configurations.

    Returns:
        (DataModel): The HotpotQA input data model
    """
    return Question

@synalinks_export("synalinks.datasets.hotpotqa.get_knowledge_data_model")
def get_knowledge_data_model():
    """
    Returns HotpotQA knowledge data model for pipeline configurations.

    Returns:
        (DataModel): The HotpotQA knowledge data model
    """
    return Document

@synalinks_export("synalinks.datasets.hotpotqa.get_output_data_model")
def get_output_data_model():
    """
    Returns HotpotQA output data model for pipeline configurations.

    Returns:
        (DataModel): The HotpotQA output data model
    """
    return Answer

@synalinks_export("synalinks.datasets.hotpotqa.iterable_dataset")
def iterable_dataset(repeat=1, batch_size=1, limit=None, split="train"):
    """
    Streaming dataset for RL-style training.

    Returns:
        (HuggingFaceDataset): A streaming, iterable dataset.
    """
    return HuggingFaceDataset(
        path="hotpot_qa",
        name="fullwiki",
        split=split,
        streaming=True,
        input_data_model=Question,
        input_template=_QUESTION_TEMPLATE,
        output_data_model=Answer,
        output_template=_ANSWER_TEMPLATE,
        batch_size=batch_size,
        limit=limit,
        repeat=repeat,
        trust_remote_code=True,
    )

@synalinks_export("synalinks.datasets.hotpotqa.load_data")
def load_data():
    """
    Load and format data from HuggingFace.

    Example:

    ```python
    (x_train, y_train), (x_test, y_test) = synalinks.datasets.hotpotqa.load_data()
    ```

    Returns:
        (tuple): The train and test data ready for training
    """
    train_ds = HuggingFaceDataset(
        path="hotpot_qa",
        name="fullwiki",
        split="train",
        streaming=False,
        input_data_model=Question,
        input_template=_QUESTION_TEMPLATE,
        output_data_model=Answer,
        output_template=_ANSWER_TEMPLATE,
        batch_size=None,
        trust_remote_code=True,
    )
    test_ds = _HardOnlyHotpotQA(
        path="hotpot_qa",
        name="fullwiki",
        split="validation",
        streaming=False,
        input_data_model=Question,
        input_template=_QUESTION_TEMPLATE,
        output_data_model=Answer,
        output_template=_ANSWER_TEMPLATE,
        batch_size=None,
        trust_remote_code=True,
    )
    x_train, y_train = next(iter(train_ds))
    x_test, y_test = next(iter(test_ds))
    return (x_train, y_train), (x_test, y_test)

@synalinks_export("synalinks.datasets.hotpotqa.load_knowledge")
def load_knowledge():
    """
    Load and format data from HuggingFace.

    Example:

    ```python
    knowledge = synalinks.datasets.hotpotqa.load_knowledge()
    ```

    Returns:
        (np.ndarray): The data ready for knowledge ingestion
    """
    ds = _HotpotKnowledge(
        path="hotpot_qa",
        name="fullwiki",
        split="train",
        streaming=False,
        input_data_model=Document,
        input_template=_DOCUMENT_TEMPLATE,
        batch_size=None,
        trust_remote_code=True,
    )
    (documents,) = next(iter(ds))
    return documents