forked from phoenix-oss/llama-stack-mirror
bcb13c492f
# What does this PR do? - revamp and clean up datasets/scoring/eval integration tests - closes https://github.com/meta-llama/llama-stack/issues/1396 [//]: # (If resolving an issue, uncomment and update the line below) [//]: # (Closes #[issue-number]) ## Test Plan **dataset** ``` LLAMA_STACK_BASE_URL=http://localhost:8321 pytest -v tests/integration/datasetio/ ``` <img width="842" alt="image" src="https://github.com/user-attachments/assets/88fc2b6a-b496-47bf-bc0c-8fea48ba36ff" /> **scoring** ``` LLAMA_STACK_CONFIG=fireworks pytest -v tests/integration/scoring --text-model meta-llama/Llama-3.1-8B-Instruct --judge-model meta-llama/Llama-3.1-8B-Instruct ``` <img width="851" alt="image" src="https://github.com/user-attachments/assets/50f46415-b44c-4c37-a6c3-076f2767adb3" /> **eval** ``` LLAMA_STACK_CONFIG=fireworks pytest -v tests/integration/eval --text-model meta-llama/Llama-3.1-8B-Instruct --judge-model meta-llama/Llama-3.1-8B-Instruct ``` <img width="841" alt="image" src="https://github.com/user-attachments/assets/8eb1c65c-3b39-4d66-8ff4-f471ca783e49" /> [//]: # (## Documentation)
101 lines
3.2 KiB
Python
101 lines
3.2 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# 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 mimetypes
|
|
import os
|
|
from pathlib import Path
|
|
|
|
# How to run this test:
|
|
#
|
|
# LLAMA_STACK_CONFIG="template-name" pytest -v tests/integration/datasetio
|
|
|
|
|
|
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
|
|
|
|
|
|
def register_dataset(llama_stack_client, for_generation=False, for_rag=False, dataset_id="test_dataset"):
|
|
if for_rag:
|
|
test_file = Path(os.path.abspath(__file__)).parent / "test_rag_dataset.csv"
|
|
else:
|
|
test_file = Path(os.path.abspath(__file__)).parent / "test_dataset.csv"
|
|
test_url = data_url_from_file(str(test_file))
|
|
|
|
if for_generation:
|
|
dataset_schema = {
|
|
"expected_answer": {"type": "string"},
|
|
"input_query": {"type": "string"},
|
|
"chat_completion_input": {"type": "chat_completion_input"},
|
|
}
|
|
elif for_rag:
|
|
dataset_schema = {
|
|
"expected_answer": {"type": "string"},
|
|
"input_query": {"type": "string"},
|
|
"generated_answer": {"type": "string"},
|
|
"context": {"type": "string"},
|
|
}
|
|
else:
|
|
dataset_schema = {
|
|
"expected_answer": {"type": "string"},
|
|
"input_query": {"type": "string"},
|
|
"generated_answer": {"type": "string"},
|
|
}
|
|
|
|
dataset_providers = [x for x in llama_stack_client.providers.list() if x.api == "datasetio"]
|
|
dataset_provider_id = dataset_providers[0].provider_id
|
|
|
|
llama_stack_client.datasets.register(
|
|
dataset_id=dataset_id,
|
|
dataset_schema=dataset_schema,
|
|
url=dict(uri=test_url),
|
|
provider_id=dataset_provider_id,
|
|
)
|
|
|
|
|
|
def test_register_unregister_dataset(llama_stack_client):
|
|
register_dataset(llama_stack_client)
|
|
response = llama_stack_client.datasets.list()
|
|
assert isinstance(response, list)
|
|
assert len(response) == 1
|
|
assert response[0].identifier == "test_dataset"
|
|
|
|
llama_stack_client.datasets.unregister("test_dataset")
|
|
response = llama_stack_client.datasets.list()
|
|
assert isinstance(response, list)
|
|
assert len(response) == 0
|
|
|
|
|
|
def test_get_rows_paginated(llama_stack_client):
|
|
register_dataset(llama_stack_client)
|
|
response = llama_stack_client.datasetio.get_rows_paginated(
|
|
dataset_id="test_dataset",
|
|
rows_in_page=3,
|
|
)
|
|
assert isinstance(response.rows, list)
|
|
assert len(response.rows) == 3
|
|
assert response.next_page_token == "3"
|
|
|
|
# iterate over all rows
|
|
response = llama_stack_client.datasetio.get_rows_paginated(
|
|
dataset_id="test_dataset",
|
|
rows_in_page=2,
|
|
page_token=response.next_page_token,
|
|
)
|
|
assert isinstance(response.rows, list)
|
|
assert len(response.rows) == 2
|
|
assert response.next_page_token == "5"
|