aiq.profiler.forecasting.models#

Submodules#

Classes#

ForecastingBaseModel

Abstract base class for all models in this package.

LinearModel

A linear regression model that conforms to the BaseModel interface.

RandomForestModel

A random forest regressor that predicts n_step token usage and call latency.

Package Contents#

class ForecastingBaseModel#

Bases: abc.ABC

Abstract base class for all models in this package.

abstractmethod fit(raw_stats)#

Train/fine-tune the model on the provided dataset.

abstractmethod predict(raw_stats) → numpy.ndarray#

Predict using the trained model. Returns a np.ndarray, shape = (N, 4).

class LinearModel#

Bases: aiq.profiler.forecasting.models.forecasting_base_model.ForecastingBaseModel

A linear regression model that conforms to the BaseModel interface.

model#
matrix_length = None#
fit(
raw_stats: list[list[aiq.data_models.intermediate_step.IntermediateStep]],
)#

X: shape (N, M) # M = matrix_length * 4 y: shape (N, 4)

predict(
raw_stats: list[list[aiq.data_models.intermediate_step.IntermediateStep]],
) → numpy.ndarray#

Predict using the fitted linear model. Returns shape (N, 4)

_prep_single(raw_stats) → numpy.ndarray#
_prep_for_model_training(raw_stats)#
_extract_token_usage_meta(all_requests_data)#
_preprocess_for_forecasting(arr: numpy.ndarray, matrix_length: int)#

Given a 2D NumPy array arr of shape (n_rows, 4), generate a list of (input_array, output_array) pairs for forecasting, each of shape:

  • input_array: (matrix_length, 4) after padding/trimming

  • output_array: (1, 4)

_flatten_features(x_list, y_list)#

x_list: list of arrays, each of shape (matrix_length, 4) y_list: list of arrays, each of shape (1, 4)

Returns:

x_flat: np.array of shape (N, matrix_length*4) y_flat: np.array of shape (N, 4)

class RandomForestModel#

Bases: aiq.profiler.forecasting.models.forecasting_base_model.ForecastingBaseModel

A random forest regressor that predicts n_step token usage and call latency.

model#
matrix_length = None#
fit(
raw_stats: list[list[aiq.data_models.intermediate_step.IntermediateStep]],
)#

X: shape (N, M) # M = matrix_length * 4 y: shape (N, 4)

predict(
raw_stats: list[list[aiq.data_models.intermediate_step.IntermediateStep]],
) → numpy.ndarray#

Predict using the fitted linear model. Returns shape (N, 4)

_prep_single(raw_stats) → numpy.ndarray#
_prep_for_model_training(raw_stats)#
_preprocess_for_forecasting(
arrays: list[numpy.ndarray],
n: int = 3,
k: int = 3,
) → list[tuple[numpy.ndarray, numpy.ndarray]]#

Preprocess a list of arrays where each array has shape (T, 3), with columns:

0: seconds_since_last_llm_call 1: input_prompt_tokens 2: output_prompt_tokens

For each row ‘i’ in each array, produce:
X: shape (n, 3)
-> The previous n calls up to row i (padded if needed).

For row i itself, set output_prompt_tokens=0 (simulate unknown current output).

Y: shape (k, 3)

-> The next k calls after row i (padded if needed).

Parameters#

arrayslist of np.ndarray

Each array is shape (T, 3).

nint

Number of past calls to include for the input context (window size).

kint

Number of future calls to include in the label (forecast horizon).

Returns#

sampleslist of (X, Y) tuples

Each X has shape (n, 3), each Y has shape (k, 3).

_extract_token_usage_meta(
all_requests_data: list[list[aiq.data_models.intermediate_step.IntermediateStep]],
)#
_flatten_features(x_list, y_list)#

X_list: list of arrays, each of shape (matrix_length, 4) y_list: list of arrays, each of shape (1, 4)

Returns:

X_flat: np.array of shape (N, matrix_length*4) y_flat: np.array of shape (N, 4)