-
Notifications
You must be signed in to change notification settings - Fork 6
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #264 from parea-ai/PAI-508-create-experiments-with…
…-sdk-visualized-in-ui feat: add experiments
- Loading branch information
Showing
13 changed files
with
296 additions
and
122 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1 @@ | ||
PAREA_OS_ENV_EXPERIMENT_UUID = "_PAREA_EXPERIMENT_UUID" |
Empty file.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,48 @@ | ||
import argparse | ||
import csv | ||
import importlib | ||
import os | ||
import sys | ||
import traceback | ||
from importlib import util | ||
|
||
from .experiment import _experiments | ||
|
||
|
||
def load_from_path(module_path): | ||
# Ensure the directory of user-provided script is in the system path | ||
dir_name = os.path.dirname(module_path) | ||
if dir_name not in sys.path: | ||
sys.path.insert(0, dir_name) | ||
|
||
module_name = os.path.basename(module_path) | ||
spec = importlib.util.spec_from_file_location(module_name, module_path) | ||
module = importlib.util.module_from_spec(spec) | ||
spec.loader.exec_module(module) | ||
|
||
if spec.name not in sys.modules: | ||
sys.modules[spec.name] = module | ||
|
||
|
||
def read_input_file(file_path) -> list[dict]: | ||
with open(file_path) as file: | ||
reader = csv.DictReader(file) | ||
inputs = list(reader) | ||
return inputs | ||
|
||
|
||
def experiment(args): | ||
parser = argparse.ArgumentParser() | ||
parser.add_argument("file", help="Path to the experiment", type=str) | ||
|
||
parsed_args = parser.parse_args(args) | ||
|
||
try: | ||
load_from_path(parsed_args.file) | ||
except Exception as e: | ||
print(f"Error loading function: {e}\n", file=sys.stderr) | ||
traceback.print_exc() | ||
sys.exit(1) | ||
|
||
for _experiment in _experiments: | ||
_experiment.run() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,93 @@ | ||
from typing import Callable, Dict, Iterator, List | ||
|
||
import asyncio | ||
import inspect | ||
import json | ||
import os | ||
import time | ||
|
||
from attrs import define, field | ||
from dotenv import load_dotenv | ||
from tqdm import tqdm | ||
|
||
from parea.client import Parea | ||
from parea.constants import PAREA_OS_ENV_EXPERIMENT_UUID | ||
from parea.schemas.models import CreateExperimentRequest, ExperimentSchema, ExperimentStatsSchema, TraceStatsSchema | ||
|
||
|
||
def calculate_avg_as_string(values: List[float]) -> str: | ||
if not values: | ||
return "N/A" | ||
values = [x for x in values if x is not None] | ||
avg = sum(values) / len(values) | ||
return f"{avg:.2f}" | ||
|
||
|
||
def calculate_avg_std_for_experiment(experiment_stats: ExperimentStatsSchema) -> Dict[str, str]: | ||
trace_stats: List[TraceStatsSchema] = experiment_stats.parent_trace_stats | ||
latency_values = [trace_stat.latency for trace_stat in trace_stats] | ||
input_tokens_values = [trace_stat.input_tokens for trace_stat in trace_stats] | ||
output_tokens_values = [trace_stat.output_tokens for trace_stat in trace_stats] | ||
total_tokens_values = [trace_stat.total_tokens for trace_stat in trace_stats] | ||
cost_values = [trace_stat.cost for trace_stat in trace_stats] | ||
score_name_to_values: Dict[str, List[float]] = {} | ||
for trace_stat in trace_stats: | ||
if trace_stat.scores: | ||
for score in trace_stat.scores: | ||
score_name_to_values[score.name] = score_name_to_values.get(score.name, []) + [score.score] | ||
|
||
return { | ||
"latency": calculate_avg_as_string(latency_values), | ||
"input_tokens": calculate_avg_as_string(input_tokens_values), | ||
"output_tokens": calculate_avg_as_string(output_tokens_values), | ||
"total_tokens": calculate_avg_as_string(total_tokens_values), | ||
"cost": calculate_avg_as_string(cost_values), | ||
**{score_name: calculate_avg_as_string(score_values) for score_name, score_values in score_name_to_values.items()}, | ||
} | ||
|
||
|
||
def async_wrapper(fn, **kwargs): | ||
return asyncio.run(fn(**kwargs)) | ||
|
||
|
||
def experiment(name: str, data: Iterator, func: Callable) -> ExperimentStatsSchema: | ||
"""Creates an experiment and runs the function on the data iterator.""" | ||
load_dotenv() | ||
|
||
if not (parea_api_key := os.getenv("PAREA_API_KEY")): | ||
raise ValueError("Please set the PAREA_API_KEY environment variable") | ||
p = Parea(api_key=parea_api_key) | ||
|
||
experiment_schema: ExperimentSchema = p.create_experiment(CreateExperimentRequest(name=name)) | ||
experiment_uuid = experiment_schema.uuid | ||
os.environ[PAREA_OS_ENV_EXPERIMENT_UUID] = experiment_uuid | ||
|
||
for data_input in tqdm(data): | ||
if inspect.iscoroutinefunction(func): | ||
asyncio.run(func(**data_input)) | ||
else: | ||
func(**data_input) | ||
time.sleep(5) # wait for all trace logs to be written to DB | ||
experiment_stats: ExperimentStatsSchema = p.get_experiment_stats(experiment_uuid) | ||
stat_name_to_avg_std = calculate_avg_std_for_experiment(experiment_stats) | ||
print(f"Experiment stats:\n{json.dumps(stat_name_to_avg_std, indent=2)}\n\n") | ||
print(f"View experiment & its traces at: https://app.parea.ai/experiments/{experiment_uuid}\n") | ||
return experiment_stats | ||
|
||
|
||
_experiments = [] | ||
|
||
|
||
@define | ||
class Experiment: | ||
name: str = field(init=True) | ||
data: Iterator[Dict] = field(init=True) | ||
func: Callable = field(init=True) | ||
experiment_stats: ExperimentStatsSchema = field(init=False, default=None) | ||
|
||
def __attrs_post_init__(self): | ||
global _experiments | ||
_experiments.append(self) | ||
|
||
def run(self): | ||
self.experiment_stats = experiment(self.name, self.data, self.func) |
Oops, something went wrong.