opik-optimizer 0.7.7__py3-none-any.whl → 0.8.0__py3-none-any.whl

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/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/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/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/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/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/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/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

opik_optimizer/datasets/hotpot_qa.py

@@ -0,0 +1,67 @@
+import opik
+from importlib.resources import files
+import json
+
+def hotpot_300(
+    test_mode: bool = False
+) -> opik.Dataset:
+    """
+    Dataset containing the first 300 samples of the HotpotQA dataset.
+    """
+    dataset_name = "hotpot_300" if not test_mode else "hotpot_300_test"
+    nb_items = 300 if not test_mode else 3
+
+    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:
+        # Load data from file and insert into the dataset
+        json_content = (files('opik_optimizer') / 'data' / 'hotpot-500.json').read_text(encoding='utf-8')
+        all_data = json.loads(json_content)
+        trainset = all_data[:nb_items]
+
+        data = []
+        for row in reversed(trainset):
+            data.append(row)
+
+        dataset.insert(data)
+        return dataset
+
+def hotpot_500(
+    test_mode: bool = False
+) -> opik.Dataset:
+    """
+    Dataset containing the first 500 samples of the HotpotQA dataset.
+    """
+    dataset_name = "hotpot_500" if not test_mode else "hotpot_500_test"
+    nb_items = 500 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:
+        # Load data from file and insert into the dataset
+        json_content = (files('opik_optimizer') / 'data' / 'hotpot-500.json').read_text(encoding='utf-8')
+        all_data = json.loads(json_content)
+        trainset = all_data[:nb_items]
+
+        data = []
+        for row in reversed(trainset):
+            data.append(row)
+
+        dataset.insert(data)
+        return dataset
+
+    
+
+    

opik_optimizer/datasets/medhallu.py

@@ -0,0 +1,39 @@
+import opik
+
+def medhallu(
+    test_mode: bool = False
+) -> opik.Dataset:
+    dataset_name = "medhallu" if not test_mode else "medhallu_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("UTAustin-AIHealth/MedHallu", "pqa_labeled", download_config=download_config)
+        
+        data = [
+            {
+                "question": item["Question"],
+                "knowledge": item["Knowledge"],
+                "ground_truth": item["Ground Truth"],
+                "hallucinated_answer": item["Hallucinated Answer"],
+                "difficulty_level": item["Difficulty Level"],
+                "hallucination_category": item["Category of Hallucination"],
+            }
+            for item in hf_dataset["train"].select(range(nb_items))
+        ]
+        ds.enable_progress_bar()
+        dataset.insert(data)
+
+        return dataset

opik_optimizer/datasets/rag_hallucinations.py

@@ -0,0 +1,41 @@
+import opik
+
+def rag_hallucinations(
+    test_mode: bool = False
+) -> opik.Dataset:
+    """
+    Dataset containing the first 300 samples of the RAG Hallucinations dataset.
+    """
+    dataset_name = "rag_hallucination" if not test_mode else "rag_hallucination_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("aporia-ai/rag_hallucinations", download_config=download_config)
+        
+        data = [
+            {
+                "context": item["context"],
+                "question": item["question"],
+                "answer": item["answer"],
+                "is_hallucination": item["is_hallucination"],
+            }
+            for item in hf_dataset["train"].select(range(nb_items))
+        ]
+        ds.enable_progress_bar()
+        
+        dataset.insert(data)
+
+        return dataset

opik_optimizer/datasets/ragbench.py

@@ -0,0 +1,40 @@
+import opik
+
+def ragbench_sentence_relevance(
+    test_mode: bool = False
+) -> opik.Dataset:
+    """
+    Dataset containing the first 300 samples of the RAGBench sentence relevance dataset.
+    """
+    dataset_name = "ragbench_sentence_relevance" if not test_mode else "ragbench_sentence_relevance_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("wandb/ragbench-sentence-relevance-balanced", download_config=download_config)
+        
+        data = [
+            {
+                "question": item["question"],
+                "sentence": item["sentence"],
+                "label": item["label"],
+            }
+            for item in hf_dataset["train"].select(range(nb_items))
+        ]
+        ds.enable_progress_bar()
+        
+        dataset.insert(data)
+
+        return dataset

opik_optimizer/datasets/tiny_test.py

@@ -0,0 +1,57 @@
+import opik
+
+TINY_TEST_ITEMS = [
+        {
+            "text": "What is the capital of France?",
+            "label": "Paris",
+            "metadata": {
+                "context": "France is a country in Europe. Its capital is Paris."
+            },
+        },
+        {
+            "text": "Who wrote Romeo and Juliet?",
+            "label": "William Shakespeare",
+            "metadata": {
+                "context": "Romeo and Juliet is a famous play written by William Shakespeare."
+            },
+        },
+        {
+            "text": "What is 2 + 2?",
+            "label": "4",
+            "metadata": {"context": "Basic arithmetic: 2 + 2 equals 4."},
+        },
+        {
+            "text": "What is the largest planet in our solar system?",
+            "label": "Jupiter",
+            "metadata": {
+                "context": "Jupiter is the largest planet in our solar system."
+            },
+        },
+        {
+            "text": "Who painted the Mona Lisa?",
+            "label": "Leonardo da Vinci",
+            "metadata": {"context": "The Mona Lisa was painted by Leonardo da Vinci."},
+        },
+    ]
+
+def tiny_test(
+    test_mode: bool = False
+) -> opik.Dataset:
+    """
+    Dataset containing the first 5 samples of the HotpotQA dataset.
+    """
+    dataset_name = "tiny_test" if not test_mode else "tiny_test_test"
+    nb_items = len(TINY_TEST_ITEMS)
+
+    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:
+        dataset.insert(TINY_TEST_ITEMS)
+        return dataset
+

opik_optimizer/datasets/truthful_qa.py

@@ -0,0 +1,107 @@
+import opik
+
+def truthful_qa(
+    test_mode: bool = False
+) -> opik.Dataset:
+    """
+    Dataset containing the first 300 samples of the TruthfulQA dataset.
+    """
+    dataset_name = "truthful_qa" if not test_mode else "truthful_qa_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()
+        
+        gen_dataset = ds.load_dataset("truthful_qa", "generation", download_config=download_config)
+        mc_dataset = ds.load_dataset("truthful_qa", "multiple_choice", download_config=download_config)
+        
+        data = []
+        for gen_item, mc_item in zip(
+            gen_dataset["validation"], mc_dataset["validation"]
+        ):
+            if len(data) >= nb_items:
+                break
+                
+            # Get correct answers from both configurations
+            correct_answers = set(gen_item["correct_answers"])
+            if "mc1_targets" in mc_item:
+                correct_answers.update(
+                    [
+                        choice
+                        for choice, label in zip(
+                            mc_item["mc1_targets"]["choices"],
+                            mc_item["mc1_targets"]["labels"],
+                        )
+                        if label == 1
+                    ]
+                )
+            if "mc2_targets" in mc_item:
+                correct_answers.update(
+                    [
+                        choice
+                        for choice, label in zip(
+                            mc_item["mc2_targets"]["choices"],
+                            mc_item["mc2_targets"]["labels"],
+                        )
+                        if label == 1
+                    ]
+                )
+
+            # Get all possible answers
+            all_answers = set(
+                gen_item["correct_answers"] + gen_item["incorrect_answers"]
+            )
+            if "mc1_targets" in mc_item:
+                all_answers.update(mc_item["mc1_targets"]["choices"])
+            if "mc2_targets" in mc_item:
+                all_answers.update(mc_item["mc2_targets"]["choices"])
+
+            # Create a single example with all necessary fields
+            example = {
+                "question": gen_item["question"],
+                "answer": gen_item["best_answer"],
+                "choices": list(all_answers),
+                "correct_answer": gen_item["best_answer"],
+                "input": gen_item["question"],  # For AnswerRelevance metric
+                "output": gen_item["best_answer"],  # For output_key requirement
+                "context": gen_item.get("source", ""),  # Use source as context
+                "type": "TEXT",  # Set type to TEXT as required by Opik
+                "category": gen_item["category"],
+                "source": "MANUAL",  # Set source to MANUAL as required by Opik
+                "correct_answers": list(
+                    correct_answers
+                ),  # Keep track of all correct answers
+                "incorrect_answers": gen_item[
+                    "incorrect_answers"
+                ],  # Keep track of incorrect answers
+            }
+
+            # Ensure all required fields are present
+            required_fields = [
+                "question",
+                "answer",
+                "choices",
+                "correct_answer",
+                "input",
+                "output",
+                "context",
+            ]
+            if all(field in example and example[field] for field in required_fields):
+                data.append(example)
+        ds.enable_progress_bar()
+        
+        dataset.insert(data)
+
+        return dataset