scorebook 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

scorebook/inference/vertex.py

@@ -0,0 +1,295 @@
+"""
+Google Cloud Vertex AI batch inference implementation for Scorebook.
+
+This module provides utilities for running batch inference using Google Cloud
+Vertex AI Gemini models, supporting large-scale asynchronous processing. It handles
+API communication, request formatting, response processing, and Cloud Storage operations.
+"""
+
+import asyncio
+import json
+import os
+import tempfile
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Union
+
+import fsspec
+import pandas as pd
+from google import genai
+from google.cloud import storage
+from google.genai import types
+from tqdm.asyncio import tqdm
+
+
+async def responses(
+    items: List[
+        Union[
+            str,
+            List[str],
+            types.Content,
+            List[types.Content],
+            types.FunctionCall,
+            List[types.FunctionCall],
+            types.Part,
+            List[types.Part],
+        ]
+    ],
+    model: str,
+    client: Optional[genai.Client] = None,
+    project_id: Optional[str] = None,
+    location: str = "us-central1",
+    system_instruction: Optional[str] = None,
+    **hyperparameters: Any,
+) -> List[types.GenerateContentResponse]:
+    """Process multiple inference requests using Google Cloud Vertex AI.
+
+    This asynchronous function handles multiple inference requests,
+    manages the API communication, and processes the responses.
+
+    Args:
+        items: List of preprocessed items to process.
+        model: Gemini model ID to use (e.g., 'gemini-2.0-flash-001').
+        client: Optional Vertex AI client instance.
+        project_id: Google Cloud Project ID. If None, uses GOOGLE_CLOUD_PROJECT env var.
+        location: Google Cloud region (default: 'us-central1').
+        system_instruction: Optional system instruction to guide model behavior.
+        hyperparameters: Additional parameters for the requests.
+
+    Returns:
+        List of raw model responses.
+    """
+    if client is None:
+        if project_id is None:
+            project_id = os.getenv("GOOGLE_CLOUD_PROJECT")
+            if not project_id:
+                raise ValueError(
+                    "Project ID must be provided or set in GOOGLE_CLOUD_PROJECT "
+                    "environment variable"
+                )
+
+        client = genai.Client(
+            vertexai=True,
+            project=project_id,
+            location=location,
+            http_options=types.HttpOptions(api_version="v1"),
+        )
+
+    # Create config if system_instruction or hyperparameters are provided
+    config = None
+    if system_instruction or hyperparameters:
+        config_dict = {}
+        if system_instruction:
+            config_dict["system_instruction"] = system_instruction
+        if hyperparameters:
+            config_dict.update(hyperparameters)
+        config = types.GenerateContentConfig(**config_dict)
+
+    results = []
+    for item in items:
+        response = client.models.generate_content(
+            model=model,
+            contents=item,
+            config=config,
+        )
+        results.append(response)
+
+    return results
+
+
+async def batch(
+    items: List[Any],
+    model: str,
+    project_id: Optional[str] = None,
+    location: str = "us-central1",
+    input_bucket: Optional[str] = None,
+    output_bucket: Optional[str] = None,
+    **hyperparameters: Any,
+) -> List[Any]:
+    """Process multiple inference requests in batch using Google Cloud Vertex AI.
+
+    This asynchronous function handles batch processing of inference requests,
+    optimizing for cost and throughput using Google Cloud's batch prediction API.
+
+    Args:
+        items: List of preprocessed items to process.
+        model: Gemini model ID to use (e.g., 'gemini-2.0-flash-001').
+        project_id: Google Cloud Project ID. If None, uses GOOGLE_CLOUD_PROJECT env var.
+        location: Google Cloud region (default: 'us-central1').
+        input_bucket: GCS bucket for input data (required).
+        output_bucket: GCS bucket for output data (required).
+        hyperparameters: Additional parameters for the batch requests.
+
+    Returns:
+        A list of raw model responses.
+    """
+    # Set up project ID
+    if project_id is None:
+        project_id = os.getenv("GOOGLE_CLOUD_PROJECT")
+        if not project_id:
+            raise ValueError(
+                "Project ID must be provided or set in GOOGLE_CLOUD_PROJECT " "environment variable"
+            )
+
+    if not input_bucket or not output_bucket:
+        raise ValueError("Both input_bucket and output_bucket must be provided")
+
+    # Initialize client
+    client = genai.Client(vertexai=True, project=project_id, location=location)
+
+    # Upload batch data
+    input_uri = await _upload_batch(items, input_bucket, model, project_id, **hyperparameters)
+
+    # Start batch job
+    batch_job = await _start_batch_job(client, model, input_uri, output_bucket)
+
+    # Wait for completion with progress tracking
+    await _wait_for_completion(client, batch_job, len(items))
+
+    # Retrieve results
+    results = await _get_batch_results(batch_job)
+
+    return results
+
+
+async def _upload_batch(
+    items: List[Any], input_bucket: str, model: str, project_id: str, **hyperparameters: Any
+) -> str:
+    """Create a JSONL file from preprocessed items and upload to GCS for batch processing."""
+
+    # Create temp JSONL file
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f:
+        for item in items:
+            # Construct batch request in Vertex AI format
+            request_data: Dict[str, Any] = {
+                "request": {
+                    "contents": [
+                        {
+                            "role": "user",
+                            "parts": [
+                                {
+                                    "text": (
+                                        str(item)
+                                        if not isinstance(item, list)
+                                        else item[0]["content"]
+                                    )
+                                }
+                            ],
+                        }
+                    ]
+                }
+            }
+
+            # Only add generationConfig if hyperparameters are provided
+            if hyperparameters:
+                request_data["request"]["generationConfig"] = hyperparameters
+            f.write(json.dumps(request_data) + "\n")
+        file_path = f.name
+
+    # Upload to GCS using Python client
+    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+
+    # Parse bucket name and path from input_bucket
+    if input_bucket.startswith("gs://"):
+        bucket_path = input_bucket[5:]  # Remove 'gs://' prefix
+    else:
+        bucket_path = input_bucket
+
+    # Split bucket name and path
+    bucket_name = bucket_path.split("/")[0]
+    bucket_prefix = "/".join(bucket_path.split("/")[1:]) if "/" in bucket_path else ""
+
+    # Create blob path
+    blob_name = (
+        f"{bucket_prefix}/batch_input_{timestamp}.jsonl"
+        if bucket_prefix
+        else f"batch_input_{timestamp}.jsonl"
+    )
+
+    # Upload using Google Cloud Storage client
+    try:
+        gcs_client = storage.Client(project=project_id)
+        bucket = gcs_client.bucket(bucket_name)
+        blob = bucket.blob(blob_name)
+
+        with open(file_path, "rb") as f:
+            blob.upload_from_file(f)
+
+        input_uri = f"gs://{bucket_name}/{blob_name}"
+
+    except Exception as e:
+        raise Exception(f"Failed to upload file to GCS: {e}")
+    finally:
+        # Clean up temp file
+        os.unlink(file_path)
+
+    return input_uri
+
+
+async def _start_batch_job(
+    client: genai.Client, model: str, input_uri: str, output_bucket: str
+) -> Any:
+    """Start a batch prediction job."""
+    batch_job = client.batches.create(
+        model=model,
+        src=input_uri,
+        config=types.CreateBatchJobConfig(dest=output_bucket),
+    )
+    return batch_job
+
+
+async def _wait_for_completion(client: genai.Client, batch_job: Any, total_items: int) -> None:
+    """Wait for batch job completion with progress tracking."""
+    # Initialize progress bar
+    pbar = tqdm(total=total_items, desc="Batch processing", unit="requests")
+
+    while True:
+        # Refresh job status
+        batch_job = client.batches.get(name=batch_job.name)
+        state = batch_job.state
+
+        # Update progress bar
+        pbar.set_postfix(status=str(state).replace("JobState.JOB_STATE_", ""))
+        pbar.refresh()
+
+        if state.name == "JOB_STATE_SUCCEEDED":
+            pbar.n = pbar.total
+            pbar.set_postfix(status="COMPLETED")
+            break
+        elif state.name == "JOB_STATE_FAILED":
+            pbar.close()
+            error_msg = getattr(batch_job, "error", "Unknown error")
+            raise Exception(f"Batch job failed: {error_msg}")
+        elif state.name in ["JOB_STATE_CANCELLED", "JOB_STATE_PAUSED"]:
+            pbar.close()
+            raise Exception(f"Batch job was {state.name}")
+
+        # Wait before checking again
+        await asyncio.sleep(30)
+
+    pbar.close()
+
+
+async def _get_batch_results(batch_job: Any) -> List[str]:
+    """Download and parse batch results from GCS."""
+
+    # Set up GCS filesystem
+    fs = fsspec.filesystem("gcs")
+
+    # Find predictions file - the pattern is: dest_uri/prediction-model-*/predictions.jsonl
+    output_uri = batch_job.dest.gcs_uri.rstrip("/")
+    file_paths = fs.glob(f"{output_uri}/prediction-model-*/predictions.jsonl")
+
+    if not file_paths:
+        raise Exception("No predictions file found in output bucket")
+
+    # Load and parse results
+    df = pd.read_json(f"gs://{file_paths[0]}", lines=True)
+
+    results = []
+    for _, row in df.iterrows():
+        # Extract text content from successful responses
+        response = row["response"]
+        text_content = response["candidates"][0]["content"]["parts"][0]["text"]
+        results.append(text_content)
+
+    return results

scorebook/types/__init__.py

@@ -7,5 +7,6 @@ and evaluation results.
 
 from scorebook.types.eval_dataset import EvalDataset
 from scorebook.types.eval_result import EvalResult
+from scorebook.types.eval_run_spec import EvalRunSpec
 
-__all__ = ["EvalDataset", "EvalResult"]
+__all__ = ["EvalDataset", "EvalResult", "EvalRunSpec"]

scorebook/types/eval_dataset.py

@@ -4,6 +4,7 @@ import csv
 import json
 from typing import Any, Dict, Iterator, List, Optional, Type, Union
 
+import yaml
 from datasets import Dataset as HuggingFaceDataset
 from datasets import DatasetDict as HuggingFaceDatasetDict
 from datasets import load_dataset
@@ -21,6 +22,7 @@ class EvalDataset:
         label: str,
         metrics: Union[str, Type[MetricBase], List[Union[str, Type[MetricBase]]]],
         hf_dataset: HuggingFaceDataset,
+        prompt_template: Optional[str] = None,
     ):
         """
         Create a new scorebook evaluation dataset instance.
@@ -29,11 +31,13 @@ class EvalDataset:
         :param label: The label field of the dataset.
         :param metrics: The specified metrics associated with the dataset.
         :param hf_dataset: The dataset as a hugging face dataset object.
+        :param prompt_template: Optional prompt template for building prompts from dataset items.
         """
         self.name: str = name
         self.label: str = label
         self.metrics: List[MetricBase] = self._resolve_metrics(metrics)
         self._hf_dataset: Optional[HuggingFaceDataset] = hf_dataset
+        self.prompt_template: Optional[str] = prompt_template
 
     def __len__(self) -> int:
         """Return the number of items in the dataset."""
@@ -82,6 +86,12 @@ class EvalDataset:
             raise ValueError("Dataset is not initialized")
         return iter(self._hf_dataset)
 
+    def shuffle(self) -> None:
+        """Randomly shuffle the dataset items."""
+        if self._hf_dataset is None:
+            raise ValueError("Dataset is not initialized")
+        self._hf_dataset.shuffle()
+
     @property
     def items(self) -> List[Any]:
         """Return a list of all examples in the dataset."""
@@ -286,6 +296,52 @@ class EvalDataset:
 
         return cls(name=path, label=label, metrics=metrics, hf_dataset=hf_dataset)
 
+    @classmethod
+    def from_yaml(cls, file_path: str) -> "EvalDataset":
+        """Instantiate an EvalDataset from a YAML file.
+
+        The YAML file should contain configuration for loading a dataset, including:
+        - name: Name of the dataset or Hugging Face dataset path
+        - label: The field used as the evaluation label
+        - metrics: List of metrics to evaluate
+        - split: Optional split name to load
+        - template: Optional prompt template
+
+        Returns:
+            An EvalDataset instance configured according to the YAML file.
+
+        Raises:
+            ValueError: If the YAML file is invalid or missing required fields.
+        """
+        path = validate_path(file_path, expected_suffix=".yaml")
+
+        try:
+            with path.open("r", encoding="utf-8") as f:
+                config = yaml.safe_load(f)
+        except yaml.YAMLError as e:
+            raise ValueError(f"Invalid YAML in {file_path}: {e}") from e
+
+        # Validate required fields
+        required_fields = ["name", "label", "metrics"]
+        missing_fields = [field for field in required_fields if field not in config]
+        if missing_fields:
+            raise ValueError(f"Missing required fields in YAML config: {', '.join(missing_fields)}")
+
+        # Load the dataset from Hugging Face
+        dataset = cls.from_huggingface(
+            path=config["name"],
+            label=config["label"],
+            metrics=config["metrics"],
+            split=config.get("split"),  # Optional field
+            name=config.get("config"),  # Optional field
+        )
+
+        # Add template if provided
+        if "template" in config:
+            dataset.prompt_template = config["template"]
+
+        return dataset
+
     @staticmethod
     def _resolve_metrics(
         metrics: Union[

scorebook/types/eval_result.py

@@ -42,6 +42,7 @@ class EvalResult:
             result = {
                 "item_id": idx,
                 "dataset_name": self.eval_dataset.name,
+                "inference_output": self.inference_outputs[idx],
                 **{
                     metric: self.metric_scores[metric]["item_scores"][idx]
                     for metric in metric_names
@@ -74,13 +75,13 @@ class EvalResult:
     def to_dict(self) -> Dict[str, Any]:
         """Return a dictionary representing the evaluation results."""
         return {
-            "aggregate": [
+            "aggregate_results": [
                 {
                     **getattr(self.eval_dataset, "hyperparams", {}),
                     **self.aggregate_scores,
                 }
             ],
-            "per_sample": [item for item in self.item_scores],
+            "item_results": [item for item in self.item_scores],
         }
 
     def to_csv(self, file_path: str) -> None:
@@ -113,7 +114,10 @@ class EvalResult:
                 writer.writerow(row)
 
     def to_json(self, file_path: str) -> None:
-        """Save evaluation results to a JSON file in structured format (Option 2)."""
+        """Save evaluation results to a JSON file in a structured format.
+
+        The JSON file will contain both aggregate & item results, produced by the to_dict() method.
+        """
         Path(file_path).parent.mkdir(parents=True, exist_ok=True)
         with open(file_path, "w") as f:
             json.dump(self.to_dict(), f, indent=2)

scorebook/types/eval_run_spec.py

@@ -0,0 +1,28 @@
+"""Evaluation run specification types for Scorebook."""
+
+from typing import Any, Dict, List, NamedTuple
+
+from scorebook.types import EvalDataset
+
+
+class EvalRunSpec(NamedTuple):
+    """Represents a single evaluation run configuration."""
+
+    dataset_idx: int
+    eval_dataset: EvalDataset
+    items: List[Dict[str, Any]]
+    labels: List[Any]
+    hyperparams: Dict[str, Any]
+    hp_idx: int
+
+    def __str__(self) -> str:
+        """Return a formatted string summary of the evaluation run specification."""
+        hyperparams_str = ", ".join([f"{k}={v}" for k, v in self.hyperparams.items()])
+
+        return (
+            f"EvalRunSpec(dataset_idx={self.dataset_idx},"
+            f"  hp_idx={self.hp_idx},"
+            f"  dataset_name='{self.eval_dataset.name}',"
+            f"  hyperparams=[{hyperparams_str}]"
+            f")"
+        )

scorebook/types/inference_pipeline.py

@@ -57,7 +57,7 @@ class InferencePipeline:
             List of processed outputs after running through the complete pipeline
         """
         if self.preprocessor:
-            input_items = [self.preprocessor(item) for item in items]
+            input_items = [self.preprocessor(item, **hyperparameters) for item in items]
         else:
             input_items = items
 
@@ -67,7 +67,10 @@ class InferencePipeline:
             inference_outputs = self.inference_function(input_items, **hyperparameters)
 
         if self.postprocessor:
-            return [self.postprocessor(inference_output) for inference_output in inference_outputs]
+            return [
+                self.postprocessor(inference_output, **hyperparameters)
+                for inference_output in inference_outputs
+            ]
         else:
             return cast(List[Any], inference_outputs)
 

scorebook/utils/__init__.py

@@ -1,8 +1,9 @@
 """Utility functions and common helpers for the Scorebook framework."""
 
 from scorebook.utils.async_utils import is_awaitable
+from scorebook.utils.build_prompt import build_prompt
 from scorebook.utils.io_helpers import validate_path
 from scorebook.utils.progress_bars import evaluation_progress
 from scorebook.utils.transform_helpers import expand_dict
 
-__all__ = ["is_awaitable", "validate_path", "expand_dict", "evaluation_progress"]
+__all__ = ["is_awaitable", "validate_path", "expand_dict", "evaluation_progress", "build_prompt"]

scorebook/utils/build_prompt.py

@@ -0,0 +1,52 @@
+"""
+Module for building prompt strings using Jinja2 templating.
+
+Provides functionality to render prompts from templates with custom filters
+and global variables, using strict undefined handling for better error detection.
+"""
+
+from typing import Any, Dict, Optional
+
+from jinja2 import BaseLoader, Environment, StrictUndefined
+
+from scorebook.utils.jinja_helpers import default_jinja_filters, default_jinja_globals
+
+
+def build_prompt(
+    prompt_template: str,
+    prompt_args: Dict[str, Any],
+    filters: Optional[Dict[str, Any]] = None,
+    globals_dict: Optional[Dict[str, Any]] = None,
+) -> str:
+    """
+    Build a prompt string from a template and arguments.
+
+    Args:
+        prompt_template: Jinja2 template string
+        prompt_args: Dictionary of arguments to pass to the template
+        filters: Dictionary of Jinja2 filters. Defaults to default_jinja_filters().
+        globals_dict: Dictionary of global functions/variables. Defaults to default_jinja_globals().
+
+    Returns:
+        str: Rendered prompt string
+    """
+
+    # Use defaults if not provided
+    filters = filters or default_jinja_filters()
+    globals_dict = globals_dict or default_jinja_globals()
+
+    # Create a Jinja2 environment with strict undefined handling
+    env = Environment(
+        loader=BaseLoader(),
+        undefined=StrictUndefined,
+        trim_blocks=True,
+        lstrip_blocks=True,
+    )
+
+    # Add filters and globals
+    env.filters.update(filters)
+    env.globals.update(globals_dict)
+
+    # Render the template
+    template = env.from_string(prompt_template)
+    return str(template.render(**prompt_args))

scorebook/utils/jinja_helpers.py

@@ -0,0 +1,146 @@
+"""Jinja2 template helper functions for Scorebook."""
+
+import json
+import re
+from typing import Any, Dict, List
+
+# Helper functions for use in Jinja templates
+
+
+def number_to_letter(index: int, uppercase: bool = True) -> str:
+    """Convert a number to a letter (0->A, 1->B, etc.).
+
+    Args:
+        index: The number to convert to a letter (0-based index, must be 0-25)
+        uppercase: If True, returns uppercase letter; if False, returns lowercase
+
+    Returns:
+        str: A letter from A-Z (or a-z if uppercase is False)
+
+    Raises:
+        ValueError: If index is less than 0 or greater than 25
+    """
+    if not 0 <= index <= 25:
+        raise ValueError("Index must be between 0 and 25 inclusive")
+    letter = chr(65 + index)
+    return letter if uppercase else letter.lower()
+
+
+def letter_to_number(letter: str) -> int:
+    """Convert a letter to a number (A->0, B->1, etc.).
+
+    Args:
+        letter: A single letter character (A-Z or a-z)
+
+    Returns:
+        int: The zero-based position of the letter in the alphabet
+
+    Raises:
+        ValueError: If the input is not a single letter character
+    """
+    if not letter.isalpha() or len(letter) != 1:
+        raise ValueError("Input must be a single letter character")
+    return ord(letter.upper()) - 65
+
+
+def format_list(items: List[Any], separator: str = ", ", last_separator: str = " and ") -> str:
+    """Format a list with proper separators and conjunction.
+
+    Examples:
+        format_list(["a", "b", "c"]) -> "a, b and c"
+        format_list(["a", "b"]) -> "a and b"
+        format_list(["a"]) -> "a"
+    """
+    if not items:
+        return ""
+    if len(items) == 1:
+        return str(items[0])
+    if len(items) == 2:
+        return f"{items[0]}{last_separator}{items[1]}"
+    return f"{separator.join(str(item) for item in items[:-1])}{last_separator}{items[-1]}"
+
+
+def truncate_text(text: str, max_length: int, suffix: str = "...") -> str:
+    """Truncate text to a maximum length with optional suffix."""
+    if len(text) <= max_length:
+        return text
+    return text[: max_length - len(suffix)] + suffix
+
+
+def format_number(number: float, precision: int = 2) -> str:
+    """Format a number with specified decimal places."""
+    return f"{number:.{precision}f}"
+
+
+def extract_initials(text: str) -> str:
+    """Extract initials from a text string.
+
+    Examples:
+        extract_initials("John Doe") -> "JD"
+        extract_initials("Machine Learning Model") -> "MLM"
+    """
+    words = re.findall(r"\b[A-Za-z]", text)
+    return "".join(words).upper()
+
+
+def json_pretty(obj: Any, indent: int = 2) -> str:
+    """Pretty-print an object as JSON."""
+    return json.dumps(obj, indent=indent, ensure_ascii=False)
+
+
+def percentage(value: float, total: float, precision: int = 1) -> str:
+    """Calculate and format a percentage.
+
+    Examples:
+        percentage(25, 100) -> "25.0%"
+        percentage(1, 3, 2) -> "33.33%"
+    """
+    if total == 0:
+        return "0.0%"
+    pct = (value / total) * 100
+    return f"{pct:.{precision}f}%"
+
+
+def ordinal(n: int) -> str:
+    """Convert number to ordinal format like 1st, 2nd, 3rd, etc."""
+    if 10 <= n % 100 <= 20:
+        suffix = "th"
+    else:
+        suffix = {1: "st", 2: "nd", 3: "rd"}.get(n % 10, "th")
+    return f"{n}{suffix}"
+
+
+def default_jinja_globals() -> Dict[str, Any]:
+    """Get default global functions for Jinja templates."""
+    return {
+        "number_to_letter": number_to_letter,
+        "letter_to_number": letter_to_number,
+        "format_list": format_list,
+        "truncate_text": truncate_text,
+        "format_number": format_number,
+        "extract_initials": extract_initials,
+        "json_pretty": json_pretty,
+        "percentage": percentage,
+        "ordinal": ordinal,
+        "max": max,
+        "min": min,
+        "len": len,
+        "abs": abs,
+        "round": round,
+        "sum": sum,
+        "sorted": sorted,
+        "enumerate": enumerate,
+        "zip": zip,
+        "range": range,
+    }
+
+
+def default_jinja_filters() -> Dict[str, Any]:
+    """Get default filters for Jinja templates."""
+    return {
+        "chr": chr,
+        "ord": ord,
+        "abs": abs,
+        "round": round,
+        "len": len,
+    }