Datasets

pb.datasets.from_file(
    path: str,                    # Path to the file to upload
    name: str,                    # Name for the dataset
) -> Dataset

# Upload a JSONL file as a dataset
dataset = pb.datasets.from_file(
    path="./data/news_articles.jsonl",
    name="tldr_news",
)

# Upload a CSV file as a dataset
dataset = pb.datasets.from_file(
    path="./data/customer_qa.csv",
    name="customer_questions",
)

pb.datasets.from_pandas_dataframe(
    df: pd.DataFrame,             # Pandas DataFrame to upload
    name: str,                    # Name for the dataset
) -> Dataset

import pandas as pd

# Create a sample DataFrame
data = {
    "prompt": ["Summarize this article:", "Explain this concept:"],
    "completion": ["This is a summary.", "This is an explanation."]
}
df = pd.DataFrame(data)

# Upload the DataFrame as a dataset
dataset = pb.datasets.from_pandas_dataframe(
    df=df,
    name="sample_dataset",  
)

pb.datasets.get(
    dataset_ref: str             # Name of the dataset
) -> Dataset

# Get a dataset by name
dataset = pb.datasets.get("tldr_news")

# Print dataset details
print(f"Dataset: {dataset.name}")
print(f"Description: {dataset.description}")
print(f"Number of rows: {dataset.num_rows}")
print(f"Created at: {dataset.created_at}")

pb.datasets.download(
    dataset_ref: str,     # Name of the dataset
    dest: str = None,     # Local destination path
) -> str

# Download a dataset as JSONL
pb.datasets.download(
    dataset_ref="tldr_news",
    dest="./downloaded_datasets/",
)
print(f"Downloaded dataset to: {download_path}")

pb.datasets.augment(
    config: dict | AugmentationConfig,  # The configuration for the augmentation
    dataset: Dataset,                   # The dataset to augment
    name: str,                          # The name of the generated dataset, must be unique in the connection
    openai_api_key: str,                # The OpenAI API key. If not provided, will be read from `OPENAI_API_KEY`
) -> Dataset

from predibase import AugmentationConfig

# Create an augmentation configuration
config = AugmentationConfig(
    base_model="gpt-4-turbo",  # Required: The OpenAI model to use
    num_samples_to_generate=500,  # Optional: Number of synthetic examples to generate
    num_seed_samples=10,  # Optional: Number of seed samples to use
    augmentation_strategy="mixture_of_agents",  # Optional: Augmentation strategy
    task_context="Generate diverse examples for customer service questions"  # Optional: Task context
)

# Get the source dataset
source_dataset = pb.datasets.get("customer_questions")

# Augment the dataset
augmented_dataset = pb.datasets.augment(
    config=config,
    dataset=source_dataset,
    name="augmented_customer_questions",
    openai_api_key="your-openai-api-key"  # Optional: If not set in environment
)

# Print dataset details
print(f"Created augmented dataset: {augmented_dataset.name}")
print(f"Number of rows: {augmented_dataset.num_rows}")
print(f"Created at: {augmented_dataset.created_at}")