From be4f395032930f8ba9b7a21da6d8a9644396a631 Mon Sep 17 00:00:00 2001
From: Xi Yan <xiyan@meta.com>
Date: Tue, 15 Oct 2024 10:17:45 -0700
Subject: [PATCH] full evals / full scoring flow

---
 docs/resources/llama-stack-spec.html          | 74 ++++++++---------
 docs/resources/llama-stack-spec.yaml          | 39 ++++-----
 llama_stack/apis/evals/client.py              | 82 +++++++++----------
 .../registry/datasets/dataset_wrappers.py     |  2 +-
 4 files changed, 88 insertions(+), 109 deletions(-)

diff --git a/docs/resources/llama-stack-spec.html b/docs/resources/llama-stack-spec.html
index ac75dbf04..7787001ff 100644
--- a/docs/resources/llama-stack-spec.html
+++ b/docs/resources/llama-stack-spec.html
@@ -21,7 +21,7 @@
     "info": {
         "title": "[DRAFT] Llama Stack Specification",
         "version": "0.0.1",
-        "description": "This is the specification of the llama stack that provides\n                a set of endpoints and their corresponding interfaces that are tailored to\n                best leverage Llama Models. The specification is still in draft and subject to change.\n                Generated at 2024-10-15 00:44:26.278642"
+        "description": "This is the specification of the llama stack that provides\n                a set of endpoints and their corresponding interfaces that are tailored to\n                best leverage Llama Models. The specification is still in draft and subject to change.\n                Generated at 2024-10-15 10:15:15.195382"
     },
     "servers": [
         {
@@ -5805,23 +5805,13 @@
             "RunEvalTaskRequest": {
                 "type": "object",
                 "properties": {
-                    "model": {
-                        "type": "string"
-                    },
-                    "task": {
-                        "type": "string"
-                    },
-                    "dataset": {
-                        "type": "string"
-                    },
                     "eval_task_config": {
                         "$ref": "#/components/schemas/EvaluateTaskConfig"
                     }
                 },
                 "additionalProperties": false,
                 "required": [
-                    "model",
-                    "task"
+                    "eval_task_config"
                 ]
             },
             "EvalResult": {
@@ -6238,49 +6228,49 @@
     ],
     "tags": [
         {
-            "name": "Models"
-        },
-        {
-            "name": "BatchInference"
-        },
-        {
-            "name": "Inspect"
-        },
-        {
-            "name": "Evals"
-        },
-        {
-            "name": "Safety"
-        },
-        {
-            "name": "Shields"
-        },
-        {
-            "name": "Telemetry"
-        },
-        {
-            "name": "Agents"
-        },
-        {
-            "name": "Memory"
-        },
-        {
-            "name": "SyntheticDataGeneration"
+            "name": "Inference"
         },
         {
             "name": "PostTraining"
         },
         {
-            "name": "Datasets"
+            "name": "Agents"
         },
         {
             "name": "MemoryBanks"
         },
+        {
+            "name": "Inspect"
+        },
+        {
+            "name": "Models"
+        },
+        {
+            "name": "Safety"
+        },
+        {
+            "name": "Evals"
+        },
+        {
+            "name": "BatchInference"
+        },
+        {
+            "name": "Shields"
+        },
+        {
+            "name": "SyntheticDataGeneration"
+        },
+        {
+            "name": "Telemetry"
+        },
         {
             "name": "RewardScoring"
         },
         {
-            "name": "Inference"
+            "name": "Datasets"
+        },
+        {
+            "name": "Memory"
         },
         {
             "name": "BuiltinTool",
diff --git a/docs/resources/llama-stack-spec.yaml b/docs/resources/llama-stack-spec.yaml
index ab54c4c09..d601435d7 100644
--- a/docs/resources/llama-stack-spec.yaml
+++ b/docs/resources/llama-stack-spec.yaml
@@ -1785,17 +1785,10 @@ components:
     RunEvalTaskRequest:
       additionalProperties: false
       properties:
-        dataset:
-          type: string
         eval_task_config:
           $ref: '#/components/schemas/EvaluateTaskConfig'
-        model:
-          type: string
-        task:
-          type: string
       required:
-      - model
-      - task
+      - eval_task_config
       type: object
     RunScorerRequest:
       additionalProperties: false
@@ -2686,7 +2679,7 @@ info:
   description: "This is the specification of the llama stack that provides\n     \
     \           a set of endpoints and their corresponding interfaces that are tailored\
     \ to\n                best leverage Llama Models. The specification is still in\
-    \ draft and subject to change.\n                Generated at 2024-10-15 00:44:26.278642"
+    \ draft and subject to change.\n                Generated at 2024-10-15 10:15:15.195382"
   title: '[DRAFT] Llama Stack Specification'
   version: 0.0.1
 jsonSchemaDialect: https://json-schema.org/draft/2020-12/schema
@@ -3787,21 +3780,21 @@ security:
 servers:
 - url: http://any-hosted-llama-stack.com
 tags:
-- name: Models
-- name: BatchInference
-- name: Inspect
-- name: Evals
-- name: Safety
-- name: Shields
-- name: Telemetry
-- name: Agents
-- name: Memory
-- name: SyntheticDataGeneration
-- name: PostTraining
-- name: Datasets
-- name: MemoryBanks
-- name: RewardScoring
 - name: Inference
+- name: PostTraining
+- name: Agents
+- name: MemoryBanks
+- name: Inspect
+- name: Models
+- name: Safety
+- name: Evals
+- name: BatchInference
+- name: Shields
+- name: SyntheticDataGeneration
+- name: Telemetry
+- name: RewardScoring
+- name: Datasets
+- name: Memory
 - description: <SchemaDefinition schemaRef="#/components/schemas/BuiltinTool" />
   name: BuiltinTool
 - description: <SchemaDefinition schemaRef="#/components/schemas/CompletionMessage"
diff --git a/llama_stack/apis/evals/client.py b/llama_stack/apis/evals/client.py
index 1e76812c6..4756a570a 100644
--- a/llama_stack/apis/evals/client.py
+++ b/llama_stack/apis/evals/client.py
@@ -119,52 +119,48 @@ async def run_main(host: str, port: int, eval_dataset_path: str = ""):
         cprint(f"{k}: {v}", "green")
 
     # Scoring Task
-    # # 1. register huggingface dataset
-    # response = await dataset_client.create_dataset(
-    #     dataset_def=HuggingfaceDatasetDef(
-    #         identifier="Llama-3.1-8B-Instruct-evals__mmlu_pro__details",
-    #         dataset_path="meta-llama/Llama-3.1-8B-Instruct-evals",
-    #         dataset_name="Llama-3.1-8B-Instruct-evals__mmlu_pro__details",
-    #         rename_columns_map={
-    #             "output_parsed_answer": "generated_answer",
-    #             "input_correct_responses": "expected_answer",
-    #         },
-    #         kwargs={"split": "latest"},
-    #     )
-    # )
-    # cprint(response, "cyan")
+    # 1. register huggingface dataset
+    response = await dataset_client.create_dataset(
+        dataset_def=HuggingfaceDatasetDef(
+            identifier="Llama-3.1-8B-Instruct-evals__mmlu_pro__details",
+            dataset_path="meta-llama/Llama-3.1-8B-Instruct-evals",
+            dataset_name="Llama-3.1-8B-Instruct-evals__mmlu_pro__details",
+            rename_columns_map={
+                "output_parsed_answer": "generated_answer",
+                "input_correct_responses": "expected_answer",
+            },
+            kwargs={"split": "latest"},
+        )
+    )
+    cprint(response, "cyan")
 
-    # # register custom dataset from file path
-    # response = await dataset_client.create_dataset(
-    #     dataset_def=CustomDatasetDef(
-    #         identifier="rag-evals",
-    #         url=data_url_from_file(eval_dataset_path),
-    #         rename_columns_map={
-    #             "query": "input_query",
-    #         },
-    #     )
-    # )
-    # cprint(response, "cyan")
+    # register custom dataset from file path
+    response = await dataset_client.create_dataset(
+        dataset_def=CustomDatasetDef(
+            identifier="rag-evals",
+            url=data_url_from_file(eval_dataset_path),
+        )
+    )
+    cprint(response, "cyan")
 
-    # # 2. run evals on the registered dataset
-    # response = await client.run_scorer(
-    #     dataset_config=EvaluateDatasetConfig(
-    #         dataset_identifier="rag-evals",
-    #         # dataset_identifier="Llama-3.1-8B-Instruct-evals__mmlu_pro__details",
-    #         row_limit=10,
-    #     ),
-    #     eval_scoring_config=EvaluateScoringConfig(
-    #         scorer_config_list=[
-    #             EvaluateSingleScorerConfig(scorer_name="accuracy"),
-    #             EvaluateSingleScorerConfig(
-    #                 scorer_name="braintrust::answer-correctness"
-    #             ),
-    #         ]
-    #     ),
-    # )
+    # 2. run evals on the registered dataset
+    response = await client.run_scorer(
+        dataset_config=EvaluateDatasetConfig(
+            dataset_identifier="rag-evals",
+            row_limit=10,
+        ),
+        eval_scoring_config=EvaluateScoringConfig(
+            scorer_config_list=[
+                EvaluateSingleScorerConfig(scorer_name="accuracy"),
+                EvaluateSingleScorerConfig(
+                    scorer_name="braintrust::answer-correctness"
+                ),
+            ]
+        ),
+    )
 
-    # for k, v in response.eval_result.metrics.items():
-    #     cprint(f"{k}: {v}", "green")
+    for k, v in response.eval_result.metrics.items():
+        cprint(f"{k}: {v}", "green")
 
 
 def main(host: str, port: int, eval_dataset_path: str = ""):
diff --git a/llama_stack/distribution/registry/datasets/dataset_wrappers.py b/llama_stack/distribution/registry/datasets/dataset_wrappers.py
index 93cbd9ab2..6c9af5887 100644
--- a/llama_stack/distribution/registry/datasets/dataset_wrappers.py
+++ b/llama_stack/distribution/registry/datasets/dataset_wrappers.py
@@ -67,7 +67,7 @@ class CustomDataset(BaseDataset[DictSample]):
             raise ValueError(f"Unsupported file type: {self.config.url}")
 
         if n_samples is not None:
-            df = df.sample(n=n_samples)
+            df = df.sample(n=min(n_samples, len(df)))
 
         self.dataset = Dataset.from_pandas(df)
         if self.config.rename_columns_map: