utils

Attributes

TASK_NAME

ALL_SUPPORTED_VALUES

devices

Classes

LengthSampler

A class for sampling text lengths within a given range.

Functions

get_device(→ str)

Determines the device for model computation.

get_nvidia_smi_info(→ Union[str, None])

Fetches NVIDIA GPU details using the nvidia-smi command.

save_results_to_json(→ None)

Saves results to a JSON file.

clean_and_trim_to_last_sentence(→ List[str])

Cleans and trims decoded outputs to the last complete sentence.

get_prompts_from_Anthropic_harmless()

Retrieves harmless prompts from the Anthropic dataset.

get_prompts_from_imdb()

Retrieves prompts from the IMDB dataset, filtered by length.

cal_humor_probabilities(sentences, model, tokenizer)

Calculates humor probabilities and raw scores.

cal_positive_sentiment(sentences, model, tokenizer)

Calculates positive sentiment probabilities and scores.

cal_harmless_probabilities(→ Tuple[List[float], ...)

Calculates harmlessness probabilities and scores.

cal_gpt2_harmless_probabilities(→ Tuple[List[float], ...)

Calculates probabilities and raw scores indicating the harmlessness of a response.

cal_gpt2_helpful_probabilities(→ Tuple[List[float], ...)

Calculates probabilities and raw scores indicating the helpfulness of a response.

compute_rep_n(→ float)

Calculates the unique fraction of n-grams within a sentence.

cal_diversity(→ float)

Calculates diversity score by computing unique n-grams for n=2,3,4.

cal_log_perplexity(→ List[float])

Calculates log perplexity for a list of sentences.

cal_coherence(→ List[float])

Calculates coherence between prompts and continuations using cosine similarity.

get_reward(→ Union[List[float], None])

Calculates reward scores based on the specified value (e.g., humor, coherence).

convert_ppo_modelname_to_huggingface_valid(→ str)

Converts a PPO-trained model name to a valid Hugging Face model path format.

get_model_and_tokenizer(model_name)

Fetches a model and tokenizer based on the model name.

Module Contents

utils.TASK_NAME = 'conversation'
utils.ALL_SUPPORTED_VALUES = ['humor', 'gpt2-helpful', 'gpt2-harmless', 'diversity', 'coherence', 'perplexity']
utils.get_device() str

Determines the device for model computation.

Returns:

Device type, either “cuda” if available, else “cpu”.

Return type:

str

utils.get_nvidia_smi_info() str | None

Fetches NVIDIA GPU details using the nvidia-smi command.

Returns:

Output from nvidia-smi command if successful, None otherwise.

Return type:

str or None

utils.devices
utils.save_results_to_json(results: list, file_path: str) None

Saves results to a JSON file.

Parameters:
  • results (list) – List of dictionaries containing prompt and generated text.

  • file_path (str) – Path to save the JSON file.

utils.clean_and_trim_to_last_sentence(prompts: List[str], decoded_outputs: List[str]) List[str]

Cleans and trims decoded outputs to the last complete sentence.

Parameters:
  • prompts (List[str]) – List of input prompts.

  • decoded_outputs (List[str]) – List of decoded model outputs.

Returns:

Cleaned and trimmed outputs.

Return type:

List[str]

class utils.LengthSampler(min_length, max_length)

A class for sampling text lengths within a given range.

min_length
max_length
__call__()

Generates a random length within the specified range.

Returns:

Random length within the min and max range.

Return type:

int

trim_to_word_boundary(text, length)

Trims text to the last word boundary within a given length.

Parameters:
  • text (str) – Text to be trimmed.

  • length (int) – Maximum character length for trimming.

Returns:

Text trimmed to the last word boundary.

Return type:

str

utils.get_prompts_from_Anthropic_harmless()

Retrieves harmless prompts from the Anthropic dataset. Importantly, the generation is non-random everytime being called.

Returns:

List of harmless prompts.

Return type:

List[str]

utils.get_prompts_from_imdb()

Retrieves prompts from the IMDB dataset, filtered by length. Importantly, the generation is non-random everytime being called.

Returns:

List of filtered IMDB prompts.

Return type:

List[str]

utils.cal_humor_probabilities(sentences, model, tokenizer)

Calculates humor probabilities and raw scores.

Parameters:
  • sentences (List[str]) – List of sentences, whose size/batchsize is limited by memory

  • model – Model for humor classification.

  • tokenizer – Tokenizer for input processing.

Returns:

Humor probabilities and raw scores.

Return type:

Tuple[List[float], List[float]]

utils.cal_positive_sentiment(sentences, model, tokenizer)

Calculates positive sentiment probabilities and scores.

Parameters:
  • sentences (List[str]) – List of sentences.

  • model – Model for sentiment classification.

  • tokenizer – Tokenizer for input processing.

Returns:

Positive probabilities and scores.

Return type:

Tuple[List[float], List[float]]

utils.cal_harmless_probabilities(sentences: List[str], model, tokenizer) Tuple[List[float], List[float]]

Calculates harmlessness probabilities and scores.

Parameters:
  • sentences (List[str]) – List of sentences.

  • model – Model for harmlessness classification.

  • tokenizer – Tokenizer for input processing.

Returns:

Harmless probabilities and scores.

Return type:

Tuple[List[float], List[float]]

utils.cal_gpt2_harmless_probabilities(prompts: List[str], continuations: List[str], model, tokenizer) Tuple[List[float], List[float]]

Calculates probabilities and raw scores indicating the harmlessness of a response.

Parameters:
  • prompts (List[str]) – List of initial prompt strings.

  • continuations (List[str]) – List of continuation (response) strings.

  • model – Model for evaluating harmlessness.

  • tokenizer – Tokenizer for processing the inputs.

Returns:

Harmlessness probabilities and raw scores.

Return type:

Tuple[List[float], List[float]]

utils.cal_gpt2_helpful_probabilities(prompts: List[str], continuations: List[str], model, tokenizer) Tuple[List[float], List[float]]

Calculates probabilities and raw scores indicating the helpfulness of a response.

Parameters:
  • prompts (List[str]) – List of prompt strings.

  • continuations (List[str]) – List of continuation (response) strings.

  • model – Model for evaluating helpfulness.

  • tokenizer – Tokenizer for processing the inputs.

Returns:

Helpfulness probabilities and raw scores.

Return type:

Tuple[List[float], List[float]]

utils.compute_rep_n(sentence: str, n: int) float

Calculates the unique fraction of n-grams within a sentence. This function is a subroutine of cal_diversity

Parameters:
  • sentence (str) – Sentence to analyze.

  • n (int) – Length of the n-grams to compute.

Returns:

Fraction of unique n-grams in the sentence.

Return type:

float

utils.cal_diversity(sentence: str) float

Calculates diversity score by computing unique n-grams for n=2,3,4.

Parameters:

sentence (str) – Sentence to analyze.

Returns:

Diversity score, product of unique n-grams for each n in (2, 3, 4).

Return type:

float

utils.cal_log_perplexity(sentences: List[str], model, tokenizer) List[float]

Calculates log perplexity for a list of sentences.

Parameters:
  • sentences (List[str]) – List of sentences.

  • model – Model to evaluate perplexity.

  • tokenizer – Tokenizer for processing the inputs.

Returns:

List of log perplexity values for each sentence.

Return type:

List[float]

utils.cal_coherence(prompts: List[str], continuations: List[str], model, tokenizer) List[float]

Calculates coherence between prompts and continuations using cosine similarity.

Parameters:
  • prompts (List[str]) – List of prompt sentences.

  • continuations (List[str]) – List of continuation sentences. Equal length as prompts as each prompt corresponds to each sentence.

  • model – Model for sentence embedding.

  • tokenizer – Tokenizer for processing the inputs.

Returns:

List of cosine similarity scores for each prompt-continuation pair.

Return type:

List[float]

utils.get_reward(sentences: List[str], value: str, model=None, tokenizer=None, prompts: List[str] = None, use_score: bool = True) List[float] | None

Calculates reward scores based on the specified value (e.g., humor, coherence).

Parameters:
  • sentences (List[str]) – List of sentences for evaluation. Each consists of a prompt and a continuation.

  • value (str) – Type of reward (e.g., “humor”, “positive”, etc.).

  • model – Model used for evaluation, if applicable.

  • tokenizer – Tokenizer used for model input processing.

  • prompts (List[str], optional) – List of prompts if needed for evaluation. Required to compute the gpt2_harmless, gpt2_helpful, coherence.

  • use_score (bool, optional) – If True, returns raw scores; otherwise, probabilities.

Returns:

List of reward scores or probabilities, or None if the value is invalid.

Return type:

Union[List[float], None]

Example

>>> rewards = get_reward(sentences, "humor")
>>> print("Average Reward:", np.mean(rewards))
utils.convert_ppo_modelname_to_huggingface_valid(ppo_model_name: str) str

Converts a PPO-trained model name to a valid Hugging Face model path format.

Parameters:

ppo_model_name (str) – Model name to be converted.

Returns:

Converted model name.

Return type:

str

utils.get_model_and_tokenizer(model_name: str)

Fetches a model and tokenizer based on the model name.

Parameters:

model_name (str) – Name of the model.

Returns:

The specified model and tokenizer.

Return type:

Tuple[transformers.PreTrainedModel, transformers.PreTrainedTokenizer]