notdiamond.toolkit

notdiamond.toolkit.custom_router

class notdiamond.toolkit.custom_router.CustomRouter(language: str = 'english', maximize: bool = True, api_key: str | None = None)[source]

Bases: object

Implementation of CustomRouter class, used to train custom routers using custom datasets.

Parameters:

  • language (str)

  • maximize (bool)

  • api_key (str | None)

language

The language of the dataset in lowercase. Defaults to “english”.

Type:

str

maximize

Whether higher score is better. Defaults to true.

Type:

bool

api_key

The NotDiamond API key. If not specified, will try to find it in the environment variable NOTDIAMOND_API_KEY.

Type:

Optional[str], optional

eval(dataset: Dict[str | LLMConfig, DataFrame], prompt_column: str, response_column: str, score_column: str, preference_id: str) → Tuple[DataFrame, DataFrame][source]

Method to evaluate a custom router using provided dataset.

Parameters:

  • dataset (Dict[str, pandas.DataFrame]) – The dataset to train a custom router. Each key in the dictionary should be in the form of <provider>/<model>.

  • prompt_column (str) – The column name in each DataFrame corresponding to the prompts used to evaluate the LLM.

  • response_column (str) – The column name in each DataFrame corresponding to the response given by the LLM for a given prompt.

  • score_column (str) – The column name in each DataFrame corresponding to the score given to the response from the LLM.

  • preference_id (str) – The preference_id associated with the custom router returned from .fit().

Raises:

  • ApiError – When the NotDiamond API fails

  • ValueError – When parsing the provided dataset fails

  • UnsupportedLLMProvider – When a provider specified in the dataset is not supported.

Returns:

eval_results_df: A DataFrame containing all the prompts, responses of each provider

(indicated by column <provider>/<model>/response), scores of each provider (indicated by column <provider>/<model>/score), and notdiamond custom router response and score (indicated by column notdiamond/response and notdiamond/score).

eval_stats_df: A DataFrame containing the “Best Average Provider” computed from the

provided dataset, the “Best Provider Average Score” achieved by the “Best Average Provider”, and the “Not Diamond Average Score” achieved through custom router.

Return type:

Tuple[pandas.DataFrame, pandas.DataFrame]

fit(dataset: Dict[str | LLMConfig, DataFrame], prompt_column: str, response_column: str, score_column: str, preference_id: str | None = None, nd_api_url: str | None = 'https://not-diamond-server.onrender.com') → str[source]

Method to train a custom router using provided dataset.

Parameters:

  • dataset (Dict[str, pandas.DataFrame]) – The dataset to train a custom router. Each key in the dictionary should be in the form of <provider>/<model>.

  • prompt_column (str) – The column name in each DataFrame corresponding to the prompts used to evaluate the LLM.

  • response_column (str) – The column name in each DataFrame corresponding to the response given by the LLM for a given prompt.

  • score_column (str) – The column name in each DataFrame corresponding to the score given to the response from the LLM.

  • preference_id (Optional[str], optional) – If specified, the custom router associated with the preference_id will be updated with the provided dataset.

  • nd_api_url (Optional[str], optional) – The URL of the NotDiamond API. Defaults to prod.

Raises:

  • ApiError – When the NotDiamond API fails

  • ValueError – When parsing the provided dataset fails

  • UnsupportedLLMProvider – When a provider specified in the dataset is not supported.

Returns:

preference_id: the preference_id associated with the custom router.

Use this preference_id in your routing calls to use the custom router.

Return type:

str