# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.

from ...datasets.interface import Row
from ...preprocessor.interface import IPreprocessor


class SimpleQAPreprocessor(IPreprocessor):
    def __init__(self, preprocessor_id: str):
        super().__init__(preprocessor_id)

    def preprocess_row(self, row: Row) -> Row:
        return {
            "input_query": row["input_query"],
            "expected_answer": row["expected_answer"],
            "messages": [
                {
                    "role": "user",
                    "content": row["input_query"],
                }
            ],
        }
