ogx_open_client 1.0.3.dev0__py3-none-any.whl

ogx_client/lib/cli/datasets/register.py

@@ -0,0 +1,78 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+import base64
+import json
+import mimetypes
+import os
+from typing import Literal
+
+import click
+import yaml
+
+from ..common.utils import handle_client_errors
+
+
+def data_url_from_file(file_path: str) -> str:
+    if not os.path.exists(file_path):
+        raise FileNotFoundError(f"File not found: {file_path}")
+
+    with open(file_path, "rb") as file:
+        file_content = file.read()
+
+    base64_content = base64.b64encode(file_content).decode("utf-8")
+    mime_type, _ = mimetypes.guess_type(file_path)
+
+    data_url = f"data:{mime_type};base64,{base64_content}"
+    return data_url
+
+
+@click.command("register")
+@click.help_option("-h", "--help")
+@click.option("--dataset-id", required=True, help="Id of the dataset")
+@click.option(
+    "--purpose",
+    type=click.Choice(["eval/question-answer", "eval/messages-answer"]),
+    help="Purpose of the dataset",
+    required=True,
+)
+@click.option("--metadata", type=str, help="Metadata of the dataset")
+@click.option("--url", type=str, help="URL of the dataset", required=False)
+@click.option(
+    "--dataset-path", required=False, help="Local file path to the dataset. If specified, upload dataset via URL"
+)
+@click.pass_context
+@handle_client_errors("register dataset")
+def register(
+    ctx,
+    dataset_id: str,
+    purpose: Literal["eval/question-answer", "eval/messages-answer"],
+    metadata: str | None,
+    url: str | None,
+    dataset_path: str | None,
+):
+    """Create a new dataset"""
+    client = ctx.obj["client"]
+
+    if metadata:
+        try:
+            metadata = json.loads(metadata)
+        except json.JSONDecodeError as err:
+            raise click.BadParameter("Metadata must be valid JSON") from err
+
+    if dataset_path:
+        url = data_url_from_file(dataset_path)
+    else:
+        if not url:
+            raise click.BadParameter("URL is required when dataset path is not specified")
+
+    response = client.datasets.register(
+        dataset_id=dataset_id,
+        source={"uri": url},
+        metadata=metadata,
+        purpose=purpose,
+    )
+    if response:
+        click.echo(yaml.dump(response.dict()))

ogx_client/lib/cli/datasets/unregister.py

@@ -0,0 +1,20 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+import click
+
+from ..common.utils import handle_client_errors
+
+
+@click.command("unregister")
+@click.help_option("-h", "--help")
+@click.argument("dataset-id", required=True)
+@click.pass_context
+@handle_client_errors("unregister dataset")
+def unregister(ctx, dataset_id: str):
+    """Remove a dataset"""
+    client = ctx.obj["client"]
+    client.datasets.unregister(dataset_id=dataset_id)
+    click.echo(f"Dataset '{dataset_id}' unregistered successfully")

ogx_client/lib/cli/eval/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from .eval import eval
+
+__all__ = ["eval"]

ogx_client/lib/cli/eval/eval.py

@@ -0,0 +1,22 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+
+import click
+
+from .run_benchmark import run_benchmark
+from .run_scoring import run_scoring
+
+
+@click.group()
+@click.help_option("-h", "--help")
+def eval():
+    """Run evaluation tasks."""
+
+
+# Register subcommands
+eval.add_command(run_benchmark)
+eval.add_command(run_scoring)

ogx_client/lib/cli/eval/run_benchmark.py

@@ -0,0 +1,193 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import json
+import os
+
+import click
+from rich import print as rprint
+from tqdm.rich import tqdm
+
+from ..common.utils import create_bar_chart
+from .utils import (
+    aggregate_accuracy,
+    aggregate_average,
+    aggregate_categorical_count,
+    aggregate_median,
+    aggregate_weighted_average,
+)
+
+
+@click.command("run-benchmark")
+@click.help_option("-h", "--help")
+@click.argument("benchmark-ids", nargs=-1, required=True)
+@click.option(
+    "--model-id",
+    required=True,
+    help="model id to run the benchmark eval on",
+    default=None,
+    type=str,
+)
+@click.option(
+    "--output-dir",
+    required=True,
+    help="Path to the dump eval results output directory",
+)
+@click.option(
+    "--num-examples",
+    required=False,
+    help="Number of examples to evaluate on, useful for debugging",
+    default=None,
+    type=int,
+)
+@click.option(
+    "--temperature",
+    required=False,
+    help="temperature in the sampling params to run generation",
+    default=0.0,
+    type=float,
+)
+@click.option(
+    "--max-tokens",
+    required=False,
+    help="max-tokens in the sampling params to run generation",
+    default=4096,
+    type=int,
+)
+@click.option(
+    "--top-p",
+    required=False,
+    help="top-p in the sampling params to run generation",
+    default=0.9,
+    type=float,
+)
+@click.option(
+    "--repeat-penalty",
+    required=False,
+    help="repeat-penalty in the sampling params to run generation",
+    default=1.0,
+    type=float,
+)
+@click.option(
+    "--visualize",
+    is_flag=True,
+    default=False,
+    help="Visualize evaluation results after completion",
+)
+@click.pass_context
+def run_benchmark(
+    ctx,
+    benchmark_ids: tuple[str, ...],
+    model_id: str,
+    output_dir: str,
+    num_examples: int | None,
+    temperature: float,
+    max_tokens: int,
+    top_p: float,
+    repeat_penalty: float,
+    visualize: bool,
+):
+    """Run a evaluation benchmark task"""
+
+    client = ctx.obj["client"]
+
+    for benchmark_id in benchmark_ids:
+        benchmark = client.benchmarks.retrieve(benchmark_id=benchmark_id)
+        scoring_functions = benchmark.scoring_functions
+        dataset_id = benchmark.dataset_id
+
+        results = client.datasets.iterrows(dataset_id=dataset_id, limit=-1 if num_examples is None else num_examples)
+
+        output_res = {}
+
+        for i, r in enumerate(tqdm(results.data)):
+            eval_res = client.eval.evaluate_rows(
+                benchmark_id=benchmark_id,
+                input_rows=[r],
+                scoring_functions=scoring_functions,
+                benchmark_config={
+                    "type": "benchmark",
+                    "eval_candidate": {
+                        "type": "model",
+                        "model": model_id,
+                        "sampling_params": {
+                            "temperature": temperature,
+                            "max_tokens": max_tokens,
+                            "top_p": top_p,
+                            "repeat_penalty": repeat_penalty,
+                        },
+                    },
+                },
+            )
+            for k in r.keys():
+                if k not in output_res:
+                    output_res[k] = []
+                output_res[k].append(r[k])
+
+            for k in eval_res.generations[0].keys():
+                if k not in output_res:
+                    output_res[k] = []
+                output_res[k].append(eval_res.generations[0][k])
+
+            for scoring_fn in scoring_functions:
+                if scoring_fn not in output_res:
+                    output_res[scoring_fn] = []
+                output_res[scoring_fn].append(eval_res.scores[scoring_fn].score_rows[0])
+
+                aggregation_functions = client.scoring_functions.retrieve(
+                    scoring_fn_id=scoring_fn
+                ).params.aggregation_functions
+
+                # only output the aggregation result for the last row
+                if i == len(results.data) - 1:
+                    for aggregation_function in aggregation_functions:
+                        scoring_results = output_res[scoring_fn]
+                        if aggregation_function == "categorical_count":
+                            output_res[scoring_fn].append(aggregate_categorical_count(scoring_results))
+                        elif aggregation_function == "average":
+                            output_res[scoring_fn].append(aggregate_average(scoring_results))
+                        elif aggregation_function == "weighted_average":
+                            output_res[scoring_fn].append(aggregate_weighted_average(scoring_results))
+                        elif aggregation_function == "median":
+                            output_res[scoring_fn].append(aggregate_median(scoring_results))
+                        elif aggregation_function == "accuracy":
+                            output_res[scoring_fn].append(aggregate_accuracy(scoring_results))
+                        else:
+                            raise NotImplementedError(
+                                f"Aggregation function {aggregation_function} is not supported yet"
+                            )
+
+        # Create output directory if it doesn't exist
+        os.makedirs(output_dir, exist_ok=True)
+        # Save results to JSON file
+        output_file = os.path.join(output_dir, f"{benchmark_id}_results.json")
+        with open(output_file, "w") as f:
+            json.dump(output_res, f, indent=2)
+
+        rprint(f"[green]✓[/green] Results saved to: [blue]{output_file}[/blue]!\n")
+
+        if visualize:
+            for scoring_fn in scoring_functions:
+                aggregation_functions = client.scoring_functions.retrieve(
+                    scoring_fn_id=scoring_fn
+                ).params.aggregation_functions
+
+                for aggregation_function in aggregation_functions:
+                    res = output_res[scoring_fn]
+                    assert len(res) > 0 and "score" in res[0]
+                    if aggregation_function == "categorical_count":
+                        scores = [str(r["score"]) for r in res]
+                        unique_scores = sorted(set(scores))
+                        counts = [scores.count(s) for s in unique_scores]
+                        create_bar_chart(
+                            counts,
+                            unique_scores,
+                            title=f"{scoring_fn}-{aggregation_function}",
+                        )
+                    else:
+                        raise NotImplementedError(
+                            f"Aggregation function {aggregation_function} ius not supported for visualization yet"
+                        )

ogx_client/lib/cli/eval/run_scoring.py

@@ -0,0 +1,116 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import json
+import os
+
+import click
+import pandas
+from rich import print as rprint
+from tqdm.rich import tqdm
+
+
+@click.command("run-scoring")
+@click.help_option("-h", "--help")
+@click.argument("scoring-function-ids", nargs=-1, required=True)
+@click.option(
+    "--dataset-id",
+    required=False,
+    help="Pre-registered dataset_id to score (from ogx-client datasets list)",
+)
+@click.option(
+    "--dataset-path",
+    required=False,
+    help="Path to the dataset file to score",
+    type=click.Path(exists=True),
+)
+@click.option(
+    "--scoring-params-config",
+    required=False,
+    help="Path to the scoring params config file in JSON format",
+    type=click.Path(exists=True),
+)
+@click.option(
+    "--num-examples",
+    required=False,
+    help="Number of examples to evaluate on, useful for debugging",
+    default=None,
+    type=int,
+)
+@click.option(
+    "--output-dir",
+    required=True,
+    help="Path to the dump eval results output directory",
+)
+@click.option(
+    "--visualize",
+    is_flag=True,
+    default=False,
+    help="Visualize evaluation results after completion",
+)
+@click.pass_context
+def run_scoring(
+    ctx,
+    scoring_function_ids: tuple[str, ...],
+    dataset_id: str | None,
+    dataset_path: str | None,
+    scoring_params_config: str | None,
+    num_examples: int | None,
+    output_dir: str,
+):
+    """Run scoring from application datasets"""
+    # one of dataset_id or dataset_path is required
+    if dataset_id is None and dataset_path is None:
+        raise click.BadParameter("Specify either dataset_id (pre-registered dataset) or dataset_path (local file)")
+
+    client = ctx.obj["client"]
+
+    scoring_params = dict.fromkeys(scoring_function_ids)
+    if scoring_params_config:
+        with open(scoring_params_config) as f:
+            scoring_params = json.load(f)
+
+    output_res = {}
+
+    if dataset_id is not None:
+        dataset = client.datasets.retrieve(dataset_id=dataset_id)
+        if not dataset:
+            click.BadParameter(f"Dataset {dataset_id} not found. Please register using ogx-client datasets register")
+
+        # TODO: this will eventually be replaced with jobs polling from server vis score_bath
+        # For now, get all datasets rows via datasets API
+        results = client.datasets.iterrows(dataset_id=dataset_id, limit=-1 if num_examples is None else num_examples)
+        rows = results.rows
+
+    if dataset_path is not None:
+        df = pandas.read_csv(dataset_path)
+        rows = df.to_dict(orient="records")
+        if num_examples is not None:
+            rows = rows[:num_examples]
+
+    for r in tqdm(rows):
+        score_res = client.scoring.score(
+            input_rows=[r],
+            scoring_functions=scoring_params,
+        )
+        for k in r.keys():
+            if k not in output_res:
+                output_res[k] = []
+            output_res[k].append(r[k])
+
+        for fn_id in scoring_function_ids:
+            if fn_id not in output_res:
+                output_res[fn_id] = []
+            output_res[fn_id].append(score_res.results[fn_id].score_rows[0])
+
+    # Create output directory if it doesn't exist
+    os.makedirs(output_dir, exist_ok=True)
+    output_file = os.path.join(output_dir, f"{dataset_path or dataset_id}_score_results.csv")
+    df = pandas.DataFrame(output_res)
+    df.to_csv(output_file, index=False)
+    print(df)
+
+    rprint(f"[green]✓[/green] Results saved to: [blue]{output_file}[/blue]!\n")

ogx_client/lib/cli/eval/utils.py

@@ -0,0 +1,59 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import statistics
+from typing import Any
+
+
+def aggregate_categorical_count(
+    scoring_results: list[dict[str, bool | float | str | list[object] | object | None]],
+) -> dict[str, Any]:
+    scores = [str(r["score"]) for r in scoring_results]
+    unique_scores = sorted(set(scores))
+    return {"categorical_count": {s: scores.count(s) for s in unique_scores}}
+
+
+def aggregate_average(
+    scoring_results: list[dict[str, bool | float | str | list[object] | object | None]],
+) -> dict[str, Any]:
+    return {
+        "average": sum(result["score"] for result in scoring_results if result["score"] is not None)
+        / len([_ for _ in scoring_results if _["score"] is not None]),
+    }
+
+
+def aggregate_weighted_average(
+    scoring_results: list[dict[str, bool | float | str | list[object] | object | None]],
+) -> dict[str, Any]:
+    return {
+        "weighted_average": sum(
+            result["score"] * result["weight"]
+            for result in scoring_results
+            if result["score"] is not None and result["weight"] is not None
+        )
+        / sum(result["weight"] for result in scoring_results if result["weight"] is not None),
+    }
+
+
+def aggregate_median(
+    scoring_results: list[dict[str, bool | float | str | list[object] | object | None]],
+) -> dict[str, Any]:
+    scores = [r["score"] for r in scoring_results if r["score"] is not None]
+    median = statistics.median(scores) if scores else None
+    return {"median": median}
+
+
+def aggregate_accuracy(
+    scoring_results: list[dict[str, bool | float | str | list[object] | object | None]],
+) -> dict[str, Any]:
+    num_correct = sum(result["score"] for result in scoring_results)
+    avg_score = num_correct / len(scoring_results)
+
+    return {
+        "accuracy": avg_score,
+        "num_correct": num_correct,
+        "num_total": len(scoring_results),
+    }

ogx_client/lib/cli/eval_tasks/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from .eval_tasks import eval_tasks
+
+__all__ = ["eval_tasks"]

ogx_client/lib/cli/eval_tasks/eval_tasks.py

@@ -0,0 +1,65 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+
+import json
+
+import click
+import yaml
+
+from ..common.utils import handle_client_errors
+from .list import list_eval_tasks
+
+
+@click.group()
+@click.help_option("-h", "--help")
+def eval_tasks():
+    """Manage evaluation tasks."""
+
+
+@eval_tasks.command()
+@click.help_option("-h", "--help")
+@click.option("--eval-task-id", required=True, help="ID of the eval task")
+@click.option("--dataset-id", required=True, help="ID of the dataset to evaluate")
+@click.option("--scoring-functions", required=True, multiple=True, help="Scoring functions to use for evaluation")
+@click.option("--provider-id", help="Provider ID for the eval task", default=None)
+@click.option("--provider-eval-task-id", help="Provider's eval task ID", default=None)
+@click.option("--metadata", type=str, help="Metadata for the eval task in JSON format")
+@click.pass_context
+@handle_client_errors("register eval task")
+def register(
+    ctx,
+    eval_task_id: str,
+    dataset_id: str,
+    scoring_functions: tuple[str, ...],
+    provider_id: str | None,
+    provider_eval_task_id: str | None,
+    metadata: str | None,
+):
+    """Register a new eval task"""
+    client = ctx.obj["client"]
+
+    if metadata:
+        try:
+            metadata = json.loads(metadata)
+        except json.JSONDecodeError as err:
+            raise click.BadParameter("Metadata must be valid JSON") from err
+
+    response = client.eval_tasks.register(
+        eval_task_id=eval_task_id,
+        dataset_id=dataset_id,
+        scoring_functions=scoring_functions,
+        provider_id=provider_id,
+        provider_eval_task_id=provider_eval_task_id,
+        metadata=metadata,
+    )
+    if response:
+        click.echo(yaml.dump(response.dict()))
+
+
+# Register subcommands
+eval_tasks.add_command(list_eval_tasks)
+eval_tasks.add_command(register)

ogx_client/lib/cli/eval_tasks/list.py

@@ -0,0 +1,35 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import click
+from rich.console import Console
+from rich.table import Table
+
+from ..common.utils import handle_client_errors
+
+
+@click.command("list")
+@click.help_option("-h", "--help")
+@click.pass_context
+@handle_client_errors("list eval tasks")
+def list_eval_tasks(ctx):
+    """Show available eval tasks on distribution endpoint"""
+
+    client = ctx.obj["client"]
+    console = Console()
+    headers = []
+    eval_tasks_list_response = client.eval_tasks.list()
+    if eval_tasks_list_response and len(eval_tasks_list_response) > 0:
+        headers = sorted(eval_tasks_list_response[0].__dict__.keys())
+
+    if eval_tasks_list_response:
+        table = Table()
+        for header in headers:
+            table.add_column(header)
+
+        for item in eval_tasks_list_response:
+            table.add_row(*[str(getattr(item, header)) for header in headers])
+        console.print(table)

ogx_client/lib/cli/inference/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from .inference import inference
+
+__all__ = ["inference"]