From 887a9772de2799e1c1067a4ac04ae4aba734c093 Mon Sep 17 00:00:00 2001 From: Xi Yan Date: Thu, 13 Mar 2025 16:31:38 -0700 Subject: [PATCH 1/4] fix precommit --- llama_stack/apis/datasetio/datasetio.py | 4 +- llama_stack/apis/datasets/datasets.py | 1 + .../open-benchmark/open_benchmark.py | 60 +++++----------- llama_stack/templates/open-benchmark/run.yaml | 70 ++++++------------- llama_stack/templates/template.py | 3 + 5 files changed, 42 insertions(+), 96 deletions(-) diff --git a/llama_stack/apis/datasetio/datasetio.py b/llama_stack/apis/datasetio/datasetio.py index 2c6c8e9816..9ce90da015 100644 --- a/llama_stack/apis/datasetio/datasetio.py +++ b/llama_stack/apis/datasetio/datasetio.py @@ -55,6 +55,4 @@ async def iterrows( ... @webmethod(route="/datasets/{dataset_id}/rows", method="POST") - async def append_rows( - self, dataset_id: str, rows: List[Dict[str, Any]] - ) -> None: ... + async def append_rows(self, dataset_id: str, rows: List[Dict[str, Any]]) -> None: ... diff --git a/llama_stack/apis/datasets/datasets.py b/llama_stack/apis/datasets/datasets.py index 9fe3cd1e01..d830a87b22 100644 --- a/llama_stack/apis/datasets/datasets.py +++ b/llama_stack/apis/datasets/datasets.py @@ -13,6 +13,7 @@ from llama_stack.schema_utils import json_schema_type, register_schema, webmethod +@json_schema_type class DatasetPurpose(Enum): """ Purpose of the dataset. Each purpose has a required input data schema. diff --git a/llama_stack/templates/open-benchmark/open_benchmark.py b/llama_stack/templates/open-benchmark/open_benchmark.py index 2b40797f9a..a598d86600 100644 --- a/llama_stack/templates/open-benchmark/open_benchmark.py +++ b/llama_stack/templates/open-benchmark/open_benchmark.py @@ -6,7 +6,7 @@ from typing import Dict, List, Tuple -from llama_stack.apis.common.content_types import URL +from llama_stack.apis.datasets import DatasetPurpose, URIDataSource from llama_stack.apis.models.models import ModelType from llama_stack.distribution.datatypes import ( BenchmarkInput, @@ -171,60 +171,34 @@ def get_distribution_template() -> DistributionTemplate: DatasetInput( dataset_id="simpleqa", provider_id="huggingface", - url=URL(uri="https://huggingface.co/datasets/llamastack/simpleqa"), - metadata={ - "path": "llamastack/simpleqa", - "split": "train", - }, - dataset_schema={ - "input_query": {"type": "string"}, - "expected_answer": {"type": "string"}, - "chat_completion_input": {"type": "string"}, - }, + purpose=DatasetPurpose.eval_messages_answer, + source=URIDataSource( + uri="huggingface://llamastack/simpleqa?split=train", + ), ), DatasetInput( dataset_id="mmlu_cot", provider_id="huggingface", - url=URL(uri="https://huggingface.co/datasets/llamastack/mmlu_cot"), - metadata={ - "path": "llamastack/mmlu_cot", - "name": "all", - "split": "test", - }, - dataset_schema={ - "input_query": {"type": "string"}, - "expected_answer": {"type": "string"}, - "chat_completion_input": {"type": "string"}, - }, + purpose=DatasetPurpose.eval_messages_answer, + source=URIDataSource( + uri="huggingface://llamastack/mmlu_cot?split=test&name=all", + ), ), DatasetInput( dataset_id="gpqa_cot", provider_id="huggingface", - url=URL(uri="https://huggingface.co/datasets/llamastack/gpqa_0shot_cot"), - metadata={ - "path": "llamastack/gpqa_0shot_cot", - "name": "gpqa_main", - "split": "train", - }, - dataset_schema={ - "input_query": {"type": "string"}, - "expected_answer": {"type": "string"}, - "chat_completion_input": {"type": "string"}, - }, + purpose=DatasetPurpose.eval_messages_answer, + source=URIDataSource( + uri="huggingface://llamastack/gpqa_0shot_cot?split=test&name=gpqa_main", + ), ), DatasetInput( dataset_id="math_500", provider_id="huggingface", - url=URL(uri="https://huggingface.co/datasets/llamastack/math_500"), - metadata={ - "path": "llamastack/math_500", - "split": "test", - }, - dataset_schema={ - "input_query": {"type": "string"}, - "expected_answer": {"type": "string"}, - "chat_completion_input": {"type": "string"}, - }, + purpose=DatasetPurpose.eval_messages_answer, + source=URIDataSource( + uri="huggingface://llamastack/math_500?split=test", + ), ), ] diff --git a/llama_stack/templates/open-benchmark/run.yaml b/llama_stack/templates/open-benchmark/run.yaml index 5ef25435b7..3fab2cfda0 100644 --- a/llama_stack/templates/open-benchmark/run.yaml +++ b/llama_stack/templates/open-benchmark/run.yaml @@ -158,62 +158,32 @@ shields: - shield_id: meta-llama/Llama-Guard-3-8B vector_dbs: [] datasets: -- dataset_schema: - input_query: - type: string - expected_answer: - type: string - chat_completion_input: - type: string - url: - uri: https://huggingface.co/datasets/llamastack/simpleqa - metadata: - path: llamastack/simpleqa - split: train +- purpose: &id001 eval/messages-answer + source: + type: uri + uri: huggingface://llamastack/simpleqa?split=train + metadata: {} dataset_id: simpleqa provider_id: huggingface -- dataset_schema: - input_query: - type: string - expected_answer: - type: string - chat_completion_input: - type: string - url: - uri: https://huggingface.co/datasets/llamastack/mmlu_cot - metadata: - path: llamastack/mmlu_cot - name: all - split: test +- purpose: *id001 + source: + type: uri + uri: huggingface://llamastack/mmlu_cot?split=test&name=all + metadata: {} dataset_id: mmlu_cot provider_id: huggingface -- dataset_schema: - input_query: - type: string - expected_answer: - type: string - chat_completion_input: - type: string - url: - uri: https://huggingface.co/datasets/llamastack/gpqa_0shot_cot - metadata: - path: llamastack/gpqa_0shot_cot - name: gpqa_main - split: train +- purpose: *id001 + source: + type: uri + uri: huggingface://llamastack/gpqa_0shot_cot?split=test&name=gpqa_main + metadata: {} dataset_id: gpqa_cot provider_id: huggingface -- dataset_schema: - input_query: - type: string - expected_answer: - type: string - chat_completion_input: - type: string - url: - uri: https://huggingface.co/datasets/llamastack/math_500 - metadata: - path: llamastack/math_500 - split: test +- purpose: *id001 + source: + type: uri + uri: huggingface://llamastack/math_500?split=test + metadata: {} dataset_id: math_500 provider_id: huggingface scoring_fns: [] diff --git a/llama_stack/templates/template.py b/llama_stack/templates/template.py index a5c8e80bc6..92b1b534d6 100644 --- a/llama_stack/templates/template.py +++ b/llama_stack/templates/template.py @@ -11,6 +11,7 @@ import yaml from pydantic import BaseModel, Field +from llama_stack.apis.datasets import DatasetPurpose from llama_stack.apis.models.models import ModelType from llama_stack.distribution.datatypes import ( Api, @@ -214,7 +215,9 @@ def enum_representer(dumper, data): # Register YAML representer for ModelType yaml.add_representer(ModelType, enum_representer) + yaml.add_representer(DatasetPurpose, enum_representer) yaml.SafeDumper.add_representer(ModelType, enum_representer) + yaml.SafeDumper.add_representer(DatasetPurpose, enum_representer) for output_dir in [yaml_output_dir, doc_output_dir]: output_dir.mkdir(parents=True, exist_ok=True) From fb07642549561cf49f3212352895eb5c1d20c712 Mon Sep 17 00:00:00 2001 From: Xi Yan Date: Thu, 13 Mar 2025 16:33:31 -0700 Subject: [PATCH 2/4] clean up --- llama_stack/apis/datasets/datasets.py | 1 - 1 file changed, 1 deletion(-) diff --git a/llama_stack/apis/datasets/datasets.py b/llama_stack/apis/datasets/datasets.py index d830a87b22..9fe3cd1e01 100644 --- a/llama_stack/apis/datasets/datasets.py +++ b/llama_stack/apis/datasets/datasets.py @@ -13,7 +13,6 @@ from llama_stack.schema_utils import json_schema_type, register_schema, webmethod -@json_schema_type class DatasetPurpose(Enum): """ Purpose of the dataset. Each purpose has a required input data schema. From 7801d9cdf2753e030caba633ce75dc4ab68f4178 Mon Sep 17 00:00:00 2001 From: Xi Yan Date: Thu, 13 Mar 2025 16:34:29 -0700 Subject: [PATCH 3/4] fix --- llama_stack/apis/datasets/datasets.py | 3 ++- llama_stack/templates/open-benchmark/run.yaml | 8 ++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/llama_stack/apis/datasets/datasets.py b/llama_stack/apis/datasets/datasets.py index 9fe3cd1e01..bd89a55a6b 100644 --- a/llama_stack/apis/datasets/datasets.py +++ b/llama_stack/apis/datasets/datasets.py @@ -13,7 +13,8 @@ from llama_stack.schema_utils import json_schema_type, register_schema, webmethod -class DatasetPurpose(Enum): +@json_schema_type +class DatasetPurpose(str, Enum): """ Purpose of the dataset. Each purpose has a required input data schema. diff --git a/llama_stack/templates/open-benchmark/run.yaml b/llama_stack/templates/open-benchmark/run.yaml index 3fab2cfda0..c0043eb2c7 100644 --- a/llama_stack/templates/open-benchmark/run.yaml +++ b/llama_stack/templates/open-benchmark/run.yaml @@ -158,28 +158,28 @@ shields: - shield_id: meta-llama/Llama-Guard-3-8B vector_dbs: [] datasets: -- purpose: &id001 eval/messages-answer +- purpose: eval/messages-answer source: type: uri uri: huggingface://llamastack/simpleqa?split=train metadata: {} dataset_id: simpleqa provider_id: huggingface -- purpose: *id001 +- purpose: eval/messages-answer source: type: uri uri: huggingface://llamastack/mmlu_cot?split=test&name=all metadata: {} dataset_id: mmlu_cot provider_id: huggingface -- purpose: *id001 +- purpose: eval/messages-answer source: type: uri uri: huggingface://llamastack/gpqa_0shot_cot?split=test&name=gpqa_main metadata: {} dataset_id: gpqa_cot provider_id: huggingface -- purpose: *id001 +- purpose: eval/messages-answer source: type: uri uri: huggingface://llamastack/math_500?split=test From 59ff6ead2cd9f016a3f1360d23b8e03a19d704c8 Mon Sep 17 00:00:00 2001 From: Xi Yan Date: Thu, 13 Mar 2025 16:37:07 -0700 Subject: [PATCH 4/4] remove json --- llama_stack/apis/datasets/datasets.py | 1 - 1 file changed, 1 deletion(-) diff --git a/llama_stack/apis/datasets/datasets.py b/llama_stack/apis/datasets/datasets.py index bd89a55a6b..1dd29ad8e4 100644 --- a/llama_stack/apis/datasets/datasets.py +++ b/llama_stack/apis/datasets/datasets.py @@ -13,7 +13,6 @@ from llama_stack.schema_utils import json_schema_type, register_schema, webmethod -@json_schema_type class DatasetPurpose(str, Enum): """ Purpose of the dataset. Each purpose has a required input data schema.