opik-optimizer 0.7.7__tar.gz → 0.8.0__tar.gz

opik_optimizer-0.8.0/PKG-INFO

@@ -0,0 +1,196 @@
+Metadata-Version: 2.4
+Name: opik_optimizer
+Version: 0.8.0
+Summary: Agent optimization with Opik
+Home-page: https://github.com/comet-ml/opik
+Author: Comet ML
+Author-email: support@comet.com
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Requires-Python: >=3.9,<3.13
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: opik>=1.7.17
+Requires-Dist: dspy<3,>=2.6.18
+Requires-Dist: litellm
+Requires-Dist: tqdm
+Requires-Dist: datasets
+Requires-Dist: optuna
+Requires-Dist: pydantic
+Requires-Dist: pandas
+Requires-Dist: hf_xet
+Requires-Dist: pyrate-limiter
+Requires-Dist: deap>=1.4.3
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-conv; extra == "dev"
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license-file
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# Opik Agent Optimizer
+
+[![PyPI version](https://img.shields.io/pypi/v/opik-optimizer.svg)](https://pypi.org/project/opik-optimizer/)
+[![Python versions](https://img.shields.io/pypi/pyversions/opik-optimizer.svg)](https://pypi.org/project/opik-optimizer/)
+[![Downloads](https://static.pepy.tech/badge/opik-optimizer)](https://pepy.tech/project/opik-optimizer)
+[![License](https://img.shields.io/github/license/comet-ml/opik)](https://github.com/comet-ml/opik/blob/main/LICENSE)
+
+The Opik Agent Optimizer refines your prompts to achieve better performance from your Large Language Models (LLMs). It supports a variety of optimization algorithms, including:
+
+* EvolutionaryOptimizer
+* FewShotBayesianOptimizer
+* MetaPromptOptimizer
+* MiproOptimizer
+
+Opik Optimizer is a component of the [Opik platform](https://github.com/comet-ml/opik), an open-source LLM evaluation platform by Comet.
+For more information about the broader Opik ecosystem, visit our [Website](https://www.comet.com/site/products/opik/) or [Documentation](https://www.comet.com/docs/opik/).
+
+## Quickstart
+
+Explore Opik Optimizer's capabilities with our interactive notebook:
+
+<a href="https://colab.research.google.com/github/comet-ml/opik/blob/main/sdks/opik_optimizer/notebooks/OpikOptimizerIntro.ipynb">
+  <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open in Colab"/>
+</a>
+
+## Setup
+
+To get started with Opik Optimizer, follow these steps:
+
+1.  **Install the package:**
+    ```bash
+    # using pip
+    pip install opik-optimizer
+
+    # using uv (faster)
+    uv pip install opik-optimizer
+    ```
+
+2.  **Configure Opik (Optional, for advanced features):**
+    If you plan to log optimization experiments to Comet or use Opik Datasets, you'll need to configure the Opik client:
+    ```bash
+    # Install the main Opik CLI (if not already installed)
+    pip install opik
+
+    # Configure your Comet API key and workspace
+    opik configure
+    # When prompted, enter your Opik API key and workspace details.
+    ```
+    Using Opik with Comet allows you to track your optimization runs, compare results, and manage datasets seamlessly.
+
+3.  **Set up LLM Provider API Keys:**
+    Ensure your environment variables are set for the LLM(s) you intend to use. For example, for OpenAI models:
+    ```bash
+    export OPENAI_API_KEY="your_openai_api_key"
+    ```
+    The optimizer utilizes LiteLLM, so you can configure keys for various providers as per LiteLLM's documentation.
+
+You'll typically need:
+
+*   An LLM model name (e.g., "gpt-4o-mini", "claude-3-haiku-20240307").
+*   An [Opik Dataset](https://www.comet.com/docs/opik/evaluation/manage_datasets/) (or a compatible local dataset/data generator).
+*   An [Opik Metric](https://www.comet.com/docs/opik/evaluation/metrics/overview/) (or a custom evaluation function).
+*   A starting prompt (template string).
+
+## Example
+
+Here's a brief example of how to use the `FewShotBayesianOptimizer`. We'll use a sample dataset provided by Opik.
+
+Available sample datasets for testing:
+*   `"tiny-test"`
+*   `"halu-eval-300"`
+*   `"hotpot-300"`
+
+```python
+from opik.evaluation.metrics import LevenshteinRatio
+from opik_optimizer import FewShotBayesianOptimizer
+from opik_optimizer.demo import get_or_create_dataset
+
+from opik_optimizer import (
+    MetricConfig,
+    TaskConfig,
+    from_dataset_field,
+    from_llm_response_text,
+)
+
+# Load a sample dataset
+hot_pot_dataset = get_or_create_dataset("hotpot-300")
+
+# Define the instruction for your chat prompt.
+# Input parameters from dataset examples will be interpolated into the full prompt.
+prompt_instruction = """
+Answer the question based on the provided context.
+"""
+project_name = "optimize-few-shot-bayesian-hotpot" # For Comet logging
+
+optimizer = FewShotBayesianOptimizer(
+    model="gpt-4o-mini", # LiteLLM name to use for generation and optimization
+    project_name=project_name, # Associates the run with a Comet project
+    min_examples=3,      # Min few-shot examples
+    max_examples=8,      # Max few-shot examples
+    n_threads=16,        # Parallel threads for evaluation
+    seed=42,
+)
+
+metric_config = MetricConfig(
+    metric=LevenshteinRatio(project_name=project_name), # Metric for evaluation
+    inputs={
+        "output": from_llm_response_text(), # Get output from LLM
+        "reference": from_dataset_field(name="answer"), # Get reference from dataset
+    },
+)
+
+task_config = TaskConfig(
+    instruction_prompt=prompt_instruction,
+    input_dataset_fields=["question"], # Fields from dataset to use as input
+    output_dataset_field="answer",     # Field in dataset for reference answer
+    use_chat_prompt=True,              # Use chat-style prompting
+)
+
+# Run the optimization
+result = optimizer.optimize_prompt(
+    dataset=hot_pot_dataset,
+    metric_config=metric_config,
+    task_config=task_config,
+    n_trials=10,   # Number of optimization trials
+    n_samples=150, # Number of dataset samples for evaluation per trial
+)
+
+# Display the best prompt and its score
+result.display()
+```
+The `result` object contains the optimized prompt, evaluation scores, and other details from the optimization process. If `project_name` is provided and Opik is configured, results will also be logged to your Comet workspace.
+
+## Development
+
+To contribute or use the Opik Optimizer from source:
+
+1.  **Clone the Opik repository:**
+    ```bash
+    git clone git@github.com:comet-ml/opik.git
+    ```
+2.  **Navigate to the optimizer's directory:**
+    ```bash
+    cd opik/sdks/opik_optimizer  # Adjust 'opik' if you cloned into a different folder name
+    ```
+3.  **Install in editable mode (with development dependencies):**
+    ```bash
+    pip install -e .[dev]
+    ```
+    The `[dev]` extra installs dependencies useful for development, such as `pytest`.
+
+## Requirements
+
+- Python `>=3.9,<3.13`
+- Opik API key (recommended for full functionality, configure via `opik configure`)
+- API key for your chosen LLM provider (e.g., OpenAI, Anthropic, Gemini), configured as per LiteLLM guidelines.

opik_optimizer-0.8.0/README.md

@@ -0,0 +1,156 @@
+# Opik Agent Optimizer
+
+[![PyPI version](https://img.shields.io/pypi/v/opik-optimizer.svg)](https://pypi.org/project/opik-optimizer/)
+[![Python versions](https://img.shields.io/pypi/pyversions/opik-optimizer.svg)](https://pypi.org/project/opik-optimizer/)
+[![Downloads](https://static.pepy.tech/badge/opik-optimizer)](https://pepy.tech/project/opik-optimizer)
+[![License](https://img.shields.io/github/license/comet-ml/opik)](https://github.com/comet-ml/opik/blob/main/LICENSE)
+
+The Opik Agent Optimizer refines your prompts to achieve better performance from your Large Language Models (LLMs). It supports a variety of optimization algorithms, including:
+
+* EvolutionaryOptimizer
+* FewShotBayesianOptimizer
+* MetaPromptOptimizer
+* MiproOptimizer
+
+Opik Optimizer is a component of the [Opik platform](https://github.com/comet-ml/opik), an open-source LLM evaluation platform by Comet.
+For more information about the broader Opik ecosystem, visit our [Website](https://www.comet.com/site/products/opik/) or [Documentation](https://www.comet.com/docs/opik/).
+
+## Quickstart
+
+Explore Opik Optimizer's capabilities with our interactive notebook:
+
+<a href="https://colab.research.google.com/github/comet-ml/opik/blob/main/sdks/opik_optimizer/notebooks/OpikOptimizerIntro.ipynb">
+  <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open in Colab"/>
+</a>
+
+## Setup
+
+To get started with Opik Optimizer, follow these steps:
+
+1.  **Install the package:**
+    ```bash
+    # using pip
+    pip install opik-optimizer
+
+    # using uv (faster)
+    uv pip install opik-optimizer
+    ```
+
+2.  **Configure Opik (Optional, for advanced features):**
+    If you plan to log optimization experiments to Comet or use Opik Datasets, you'll need to configure the Opik client:
+    ```bash
+    # Install the main Opik CLI (if not already installed)
+    pip install opik
+
+    # Configure your Comet API key and workspace
+    opik configure
+    # When prompted, enter your Opik API key and workspace details.
+    ```
+    Using Opik with Comet allows you to track your optimization runs, compare results, and manage datasets seamlessly.
+
+3.  **Set up LLM Provider API Keys:**
+    Ensure your environment variables are set for the LLM(s) you intend to use. For example, for OpenAI models:
+    ```bash
+    export OPENAI_API_KEY="your_openai_api_key"
+    ```
+    The optimizer utilizes LiteLLM, so you can configure keys for various providers as per LiteLLM's documentation.
+
+You'll typically need:
+
+*   An LLM model name (e.g., "gpt-4o-mini", "claude-3-haiku-20240307").
+*   An [Opik Dataset](https://www.comet.com/docs/opik/evaluation/manage_datasets/) (or a compatible local dataset/data generator).
+*   An [Opik Metric](https://www.comet.com/docs/opik/evaluation/metrics/overview/) (or a custom evaluation function).
+*   A starting prompt (template string).
+
+## Example
+
+Here's a brief example of how to use the `FewShotBayesianOptimizer`. We'll use a sample dataset provided by Opik.
+
+Available sample datasets for testing:
+*   `"tiny-test"`
+*   `"halu-eval-300"`
+*   `"hotpot-300"`
+
+```python
+from opik.evaluation.metrics import LevenshteinRatio
+from opik_optimizer import FewShotBayesianOptimizer
+from opik_optimizer.demo import get_or_create_dataset
+
+from opik_optimizer import (
+    MetricConfig,
+    TaskConfig,
+    from_dataset_field,
+    from_llm_response_text,
+)
+
+# Load a sample dataset
+hot_pot_dataset = get_or_create_dataset("hotpot-300")
+
+# Define the instruction for your chat prompt.
+# Input parameters from dataset examples will be interpolated into the full prompt.
+prompt_instruction = """
+Answer the question based on the provided context.
+"""
+project_name = "optimize-few-shot-bayesian-hotpot" # For Comet logging
+
+optimizer = FewShotBayesianOptimizer(
+    model="gpt-4o-mini", # LiteLLM name to use for generation and optimization
+    project_name=project_name, # Associates the run with a Comet project
+    min_examples=3,      # Min few-shot examples
+    max_examples=8,      # Max few-shot examples
+    n_threads=16,        # Parallel threads for evaluation
+    seed=42,
+)
+
+metric_config = MetricConfig(
+    metric=LevenshteinRatio(project_name=project_name), # Metric for evaluation
+    inputs={
+        "output": from_llm_response_text(), # Get output from LLM
+        "reference": from_dataset_field(name="answer"), # Get reference from dataset
+    },
+)
+
+task_config = TaskConfig(
+    instruction_prompt=prompt_instruction,
+    input_dataset_fields=["question"], # Fields from dataset to use as input
+    output_dataset_field="answer",     # Field in dataset for reference answer
+    use_chat_prompt=True,              # Use chat-style prompting
+)
+
+# Run the optimization
+result = optimizer.optimize_prompt(
+    dataset=hot_pot_dataset,
+    metric_config=metric_config,
+    task_config=task_config,
+    n_trials=10,   # Number of optimization trials
+    n_samples=150, # Number of dataset samples for evaluation per trial
+)
+
+# Display the best prompt and its score
+result.display()
+```
+The `result` object contains the optimized prompt, evaluation scores, and other details from the optimization process. If `project_name` is provided and Opik is configured, results will also be logged to your Comet workspace.
+
+## Development
+
+To contribute or use the Opik Optimizer from source:
+
+1.  **Clone the Opik repository:**
+    ```bash
+    git clone git@github.com:comet-ml/opik.git
+    ```
+2.  **Navigate to the optimizer's directory:**
+    ```bash
+    cd opik/sdks/opik_optimizer  # Adjust 'opik' if you cloned into a different folder name
+    ```
+3.  **Install in editable mode (with development dependencies):**
+    ```bash
+    pip install -e .[dev]
+    ```
+    The `[dev]` extra installs dependencies useful for development, such as `pytest`.
+
+## Requirements
+
+- Python `>=3.9,<3.13`
+- Opik API key (recommended for full functionality, configure via `opik configure`)
+- API key for your chosen LLM provider (e.g., OpenAI, Anthropic, Gemini), configured as per LiteLLM guidelines.

{opik_optimizer-0.7.7 → opik_optimizer-0.8.0}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name="opik_optimizer",
-    version="0.7.7",
+    version="0.8.0",
     description="Agent optimization with Opik",
     author="Comet ML",
     author_email="support@comet.com",
@@ -31,7 +31,6 @@ setup(
     # dev requirements
     extras_require={
         "dev": [
-            "adalflow",
             "pytest",
             "pytest-conv"
         ],

{opik_optimizer-0.7.7 → opik_optimizer-0.8.0}/src/opik_optimizer/__init__.py

@@ -23,6 +23,7 @@ from .optimization_config.mappers import (
 )
 
 from opik.evaluation.models.litellm import warning_filters
+from . import datasets
 
 warning_filters.add_warning_filters()
 
@@ -42,4 +43,5 @@ __all__ = [
     "from_llm_response_text",
     "OptimizationResult",
     "setup_logging",
+    "datasets",
 ]

{opik_optimizer-0.7.7 → opik_optimizer-0.8.0}/src/opik_optimizer/base_optimizer.py

@@ -4,15 +4,15 @@ import logging
 import time
 
 import litellm
+from . import _throttle
 from opik.rest_api.core import ApiError
 
 from pydantic import BaseModel
-from ._throttle import RateLimiter, rate_limited
 from .cache_config import initialize_cache
 from opik.evaluation.models.litellm import opik_monitor as opik_litellm_monitor
 from .optimization_config.configs import TaskConfig, MetricConfig
 
-limiter = RateLimiter(max_calls_per_second=8)
+_limiter = _throttle.get_rate_limiter_for_current_opik_installation()
 
 # Don't use unsupported params:
 litellm.drop_params = True
@@ -32,19 +32,21 @@ class OptimizationRound(BaseModel):
 
 
 class BaseOptimizer:
-    def __init__(self, model: str, project_name: Optional[str] = None, **model_kwargs):
+    def __init__(self, model: str, project_name: Optional[str] = None, verbose: int = 1, **model_kwargs):
         """
         Base class for optimizers.
 
         Args:
            model: LiteLLM model name
            project_name: Opik project name
+           verbose: Controls internal logging/progress bars (0=off, 1=on).
            model_kwargs: additional args for model (eg, temperature)
         """
         self.model = model
         self.reasoning_model = model
         self.model_kwargs = model_kwargs
         self.project_name = project_name
+        self.verbose = verbose
         self._history = []
         self.experiment_config = None
         self.llm_call_counter = 0
@@ -141,7 +143,7 @@ class BaseOptimizer:
         """
         self._history.append(round_data)
 
-
+            
     def update_optimization(self, optimization, status: str) -> None:
         """
         Update the optimization status

opik_optimizer-0.8.0/src/opik_optimizer/datasets/__init__.py

@@ -0,0 +1,27 @@
+from .hotpot_qa import hotpot_300, hotpot_500
+from .halu_eval import halu_eval_300
+from .tiny_test import tiny_test
+from .gsm8k import gsm8k
+from .ai2_arc import ai2_arc
+from .truthful_qa import truthful_qa
+from .cnn_dailymail import cnn_dailymail
+from .ragbench import ragbench_sentence_relevance
+from .election_questions import election_questions
+from .medhallu import medhallu
+from .rag_hallucinations import rag_hallucinations
+
+
+__all__ = [
+    "hotpot_300",
+    "hotpot_500",
+    "halu_eval_300",
+    "tiny_test",
+    "gsm8k",
+    "ai2_arc",
+    "truthful_qa",
+    "cnn_dailymail",
+    "ragbench_sentence_relevance",
+    "election_questions",
+    "medhallu",
+    "rag_hallucinations",
+]

opik_optimizer-0.8.0/src/opik_optimizer/datasets/ai2_arc.py

@@ -0,0 +1,44 @@
+import opik
+
+def ai2_arc(
+    test_mode: bool = False
+) -> opik.Dataset:
+    """
+    Dataset containing the first 300 samples of the AI2 ARC dataset.
+    """
+    dataset_name = "ai2_arc" if not test_mode else "ai2_arc_test"
+    nb_items = 300 if not test_mode else 5
+
+    client = opik.Opik()
+    dataset = client.get_or_create_dataset(dataset_name)
+    
+    items = dataset.get_items()
+    if len(items) == nb_items:
+        return dataset
+    elif len(items) != 0:
+        raise ValueError(f"Dataset {dataset_name} contains {len(items)} items, expected {nb_items}. We recommend deleting the dataset and re-creating it.")
+    elif len(items) == 0:
+        import datasets as ds
+
+        # Load data from file and insert into the dataset
+        download_config = ds.DownloadConfig(download_desc=False, disable_tqdm=True)
+        ds.disable_progress_bar()
+        hf_dataset = ds.load_dataset(
+            "ai2_arc", "ARC-Challenge",
+            streaming=True, download_config=download_config
+        )
+        
+        data = []
+        for i, item in enumerate(hf_dataset["train"]):
+            if i >= nb_items:
+                break
+            data.append({
+                "question": item["question"],
+                "answer": item["answerKey"],
+                "choices": item["choices"],
+            })
+        ds.enable_progress_bar()
+
+        dataset.insert(data)
+
+        return dataset

opik_optimizer-0.8.0/src/opik_optimizer/datasets/cnn_dailymail.py

@@ -0,0 +1,40 @@
+import opik
+
+def cnn_dailymail(
+    test_mode: bool = False
+) -> opik.Dataset:
+    """
+    Dataset containing the first 100 samples of the CNN Daily Mail dataset.
+    """
+    dataset_name = "cnn_dailymail" if not test_mode else "cnn_dailymail_test"
+    nb_items = 100 if not test_mode else 5
+
+    client = opik.Opik()
+    dataset = client.get_or_create_dataset(dataset_name)
+    
+    items = dataset.get_items()
+    if len(items) == nb_items:
+        return dataset
+    elif len(items) != 0:
+        raise ValueError(f"Dataset {dataset_name} contains {len(items)} items, expected {nb_items}. We recommend deleting the dataset and re-creating it.")
+    elif len(items) == 0:
+        import datasets as ds
+        
+        download_config = ds.DownloadConfig(download_desc=False, disable_tqdm=True)
+        ds.disable_progress_bar()
+        hf_dataset = ds.load_dataset("cnn_dailymail", "3.0.0", streaming=True, download_config=download_config)
+        
+        data = []
+        for i, item in enumerate(hf_dataset["validation"]):
+            if i >= nb_items:
+                break
+            data.append({
+                "article": item["article"],
+                "highlights": item["highlights"],
+            })
+        ds.enable_progress_bar()
+        
+        dataset.insert(data)
+        
+        return dataset
+    

opik_optimizer-0.8.0/src/opik_optimizer/datasets/election_questions.py

@@ -0,0 +1,36 @@
+import opik
+
+
+def election_questions(
+    test_mode: bool = False
+) -> opik.Dataset:
+    dataset_name = "election_questions" if not test_mode else "election_questions_test"
+    nb_items = 300 if not test_mode else 5
+
+    client = opik.Opik()
+    dataset = client.get_or_create_dataset(dataset_name)
+    
+    items = dataset.get_items()
+    if len(items) == nb_items:
+        return dataset
+    elif len(items) != 0:
+        raise ValueError(f"Dataset {dataset_name} contains {len(items)} items, expected {nb_items}. We recommend deleting the dataset and re-creating it.")
+    elif len(items) == 0:
+        import datasets as ds
+
+        # Load data from file and insert into the dataset
+        download_config = ds.DownloadConfig(download_desc=False, disable_tqdm=True)
+        ds.disable_progress_bar()
+        hf_dataset = ds.load_dataset("Anthropic/election_questions", download_config=download_config)
+        
+        data = [
+            {
+                "question": item["question"],
+                "label": item["label"]
+            }
+            for item in hf_dataset["test"].select(range(nb_items))
+        ]
+        ds.enable_progress_bar()
+        dataset.insert(data)
+
+        return dataset

opik_optimizer-0.8.0/src/opik_optimizer/datasets/gsm8k.py

@@ -0,0 +1,40 @@
+import opik
+
+def gsm8k(
+    test_mode: bool = False
+) -> opik.Dataset:
+    """
+    Dataset containing the first 300 samples of the GSM8K dataset.
+    """
+    dataset_name = "gsm8k" if not test_mode else "gsm8k_test"
+    nb_items = 300 if not test_mode else 5
+
+    client = opik.Opik()
+    dataset = client.get_or_create_dataset(dataset_name)
+    
+    items = dataset.get_items()
+    if len(items) == nb_items:
+        return dataset
+    elif len(items) != 0:
+        raise ValueError(f"Dataset {dataset_name} contains {len(items)} items, expected {nb_items}. We recommend deleting the dataset and re-creating it.")
+    elif len(items) == 0:
+        import datasets as ds
+
+        # Load data from file and insert into the dataset
+        download_config = ds.DownloadConfig(download_desc=False, disable_tqdm=True)
+        ds.disable_progress_bar()
+        hf_dataset = ds.load_dataset("gsm8k", "main", streaming=True, download_config=download_config)
+        
+        data = []
+        for i, item in enumerate(hf_dataset["train"]):
+            if i >= nb_items:
+                break
+            data.append({
+                "question": item["question"],
+                "answer": item["answer"],
+            })
+        ds.enable_progress_bar()
+
+        dataset.insert(data)
+
+        return dataset

opik_optimizer-0.8.0/src/opik_optimizer/datasets/halu_eval.py

@@ -0,0 +1,43 @@
+import opik
+
+def halu_eval_300(
+    test_mode: bool = False
+) -> opik.Dataset:
+    """
+    Dataset containing the first 300 samples of the HaluEval dataset.
+    """
+    dataset_name = "halu_eval_300" if not test_mode else "halu_eval_300_test"
+    nb_items = 300 if not test_mode else 5
+
+    client = opik.Opik()
+    dataset = client.get_or_create_dataset(dataset_name)
+    
+    items = dataset.get_items()
+    if len(items) == nb_items:
+        return dataset
+    elif len(items) != 0:
+        raise ValueError(f"Dataset {dataset_name} contains {len(items)} items, expected {nb_items}. We recommend deleting the dataset and re-creating it.")
+    elif len(items) == 0:
+        import pandas as pd
+
+        try:
+            df = pd.read_parquet(
+                "hf://datasets/pminervini/HaluEval/general/data-00000-of-00001.parquet"
+            )
+        except Exception:
+            raise Exception("Unable to download HaluEval; please try again") from None
+
+        sample_size = min(nb_items, len(df))
+        df_sampled = df.sample(n=sample_size, random_state=42)
+
+        dataset_records = [
+            {
+                "input": x["user_query"],
+                "llm_output": x["chatgpt_response"],
+                "expected_hallucination_label": x["hallucination"],
+            }
+            for x in df_sampled.to_dict(orient="records")
+        ]
+
+        dataset.insert(dataset_records)
+        return dataset