llama-stack-mirror/tests/integration/datasets/test_script.py
2025-03-15 16:15:45 -07:00

21 lines
591 B
Python

from llama_stack_client import LlamaStackClient
from rich.pretty import pprint
def test_register_dataset():
client = LlamaStackClient(base_url="http://localhost:8321")
dataset = client.datasets.register(
purpose="eval/messages-answer",
source={
"type": "uri",
"uri": "huggingface://datasets/llamastack/simpleqa?split=train",
},
)
dataset_id = dataset.identifier
pprint(dataset)
rows = client.datasets.iterrows(dataset_id=dataset_id, limit=10)
pprint(rows)
if __name__ == "__main__":
test_register_dataset()