import pandas as pd
import phoenix as px
from phoenix.client import Client


def create_dataset():
    dataset_df = pd.read_json("./data/metadata.jsonl", lines=True)

    # Script should be run with a running phoenix server, if not uncomment:
    # _ = px.launch_app()
    px_client = Client()
    dataset = px_client.datasets.create_dataset(
        dataframe=dataset_df,
        name="gaia",
        input_keys=["Question"],
        output_keys=["Final answer"],
        metadata_keys=["task_id", "Annotator Metadata", "file_name"],
    )
    print(f"Dataset created: {dataset.id}")


if __name__ == "__main__":
    create_dataset()