Function prepare_huggingface_dataset

tests/test-dataset.py:205–264

Download and prepare a dataset from Hugging Face

(dataset_id, base_destination)

def prepare_huggingface_dataset(dataset_id, base_destination):
    """Download and prepare a dataset from Hugging Face"""
    print(f"Loading dataset {dataset_id} from Hugging Face with streaming...")

    # Find dataset configuration by dataset_id
    dataset_config = None
    for name, config in datasets.items():
        if config["dataset_id"] == dataset_id:
            dataset_config = config
            break
    
    if not dataset_config:
        raise ValueError(f"No configuration found for dataset {dataset_id}")

    try:
        # Stream the dataset and process in chunks
        dataset = load_dataset(dataset_id, split='train', streaming=True)
        chunk_size = 200_000
        chunk = []
        file_index = 0

        for i, example in enumerate(dataset):
            chunk.append(example)
            if (i + 1) % chunk_size == 0:
                # Convert chunk to DataFrame
                chunk_df = pd.DataFrame(chunk)

                # Ensure columns
                if "id" not in chunk_df.columns:
                    offset = chunk_size * file_index
                    chunk_df["id"] = range(offset, offset + len(chunk_df))

                if "dense_values" not in chunk_df.columns:
                    # Use the predefined embedding column from configuration
                    embedding_col = dataset_config["embeddings"]
                    if embedding_col not in chunk_df.columns:
                        raise ValueError(
                            f"Configured embedding column '{embedding_col}' not found in dataset {dataset_id}. Available columns: {chunk_df.columns.tolist()}"
                        )
                    chunk_df["dense_values"] = chunk_df[embedding_col]

                # Save the chunk as parquet
                destination = f"{base_destination}/test{file_index}.parquet"
                print(f"\nSaving chunk to {destination}")
                chunk_df.to_parquet(destination)

                # Clear chunk for next batch
                chunk.clear()
                file_index += 1

        # Save any remaining data
        if chunk:
            chunk_df = pd.DataFrame(chunk)
            destination = f"{base_destination}/test{file_index}.parquet"
            print(f"\nSaving remaining chunk to {destination}")
            chunk_df.to_parquet(destination)

    except Exception as e: