Jacklanchantin/online training

src/fairseq2/recipes/lm/_online_finetune/_common.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -598,3 +598,14 @@
     for rollout in rollouts[0].outputs[:num_rollouts]:
         rollout_text = rollout.text
         log.info(f"{split_name} Rollout: {rollout_text}")
+
+
+def get_rollout_lengths(rollouts: List[SequenceData]):
+    """Get the lengths of the rollouts."""
+    rollout_lengths = []
+    for rollout in rollouts:
+        for sample in rollout.outputs:
+            token_ids = sample.token_ids
+            token_ids_len = len(token_ids)
+            rollout_lengths.append(token_ids_len)
+    return rollout_lengths

src/fairseq2/recipes/lm/_online_finetune/_diversity_metrics.py

@@ -0,0 +1,100 @@
+import string as string_lib
+from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
+import gzip
+import torch
+
+
+def get_compression_ratio(strings):
+
+    flattened_generation = " ".join(strings)
+    original_byte_size = len(bytes(flattened_generation, "UTF-8"))
+    compressed_bytes_size = len(gzip.compress(bytes(flattened_generation, "UTF-8")))
+
+    cr = compressed_bytes_size / original_byte_size
+    cr_tensor = torch.Tensor([cr])
+    return cr_tensor
+
+
+def get_self_bleu_score(strings):
+    # Create a translation table to remove punctuation
+    translator = str.maketrans("", "", string_lib.punctuation)
+
+    # Preprocess the strings: convert to lowercase and remove punctuation
+    cleaned_strings = [s.lower().translate(translator) for s in strings]
+
+    # Tokenize the cleaned strings into lists of words
+    tokenized_strings = [s.split() for s in cleaned_strings]
+
+    # Initialize a dictionary to store BLEU scores
+    bleu_scores = []
+
+    # Calculate BLEU scores for all pairs of strings
+    for i in range(len(tokenized_strings)):
+        for j in range(i + 1, len(tokenized_strings)):
+            # Use smoothing to handle cases where there are no n-grams in common
+            smoothie = SmoothingFunction().method4
+            bleu = sentence_bleu(
+                [tokenized_strings[i]],
+                tokenized_strings[j],
+                smoothing_function=smoothie,
+            )
+
+            # Store the BLEU score
+            bleu_scores.append(bleu)
+
+    mean_bleu_score = sum(bleu_scores) / len(bleu_scores) if bleu_scores else 0.0
+    mean_bleu_score_tensor = torch.Tensor([mean_bleu_score])
+    return mean_bleu_score_tensor
+
+
+def get_unique_1grams(strings):
+
+    # Initialize an empty set to store unique 1-grams
+    unique_words = set()
+    total_words = 0
+
+    # Create a translation table to remove punctuation
+    translator = str.maketrans("", "", string_lib.punctuation)
+
+    # Iterate over each string in the list
+    for string in strings:
+        # Convert the string to lowercase and remove punctuation
+        cleaned_string = string.lower().translate(translator)
+
+        # Split the cleaned string into words (1-grams) and update the set
+        words = cleaned_string.split()
+        total_words += len(words)
+        unique_words.update(words)
+
+    # Return the set of unique 1-grams
+    num_unique_1grams = len(unique_words)
+    num_unique_1grams_norm = len(unique_words) / total_words if total_words > 0 else 0
+    num_unique_1grams_tensor = torch.Tensor([num_unique_1grams])
+    num_unique_1grams_norm = torch.Tensor([num_unique_1grams_norm])
+    return num_unique_1grams_tensor, num_unique_1grams_norm
+
+
+def extract_logprobs(data):
+    logprobs = []
+    for item in data:
+        for key, logprob in item.items():
+            logprobs.append(logprob.logprob)
+    return logprobs
+
+
+def get_entropy(rollouts):
+    batch_sum_logprobs = []
+    batch_sum_logprobs_per_tok = []
+    for rollout_idx in range(len(rollouts[0].outputs)):
+        logprobs = extract_logprobs(rollouts[0].outputs[rollout_idx].logprobs)
+
+        sum_logprobs = -sum(logprobs)
+        sum_logprobs_per_tok = -sum(logprobs) / len(logprobs)
+
+        batch_sum_logprobs.append(sum_logprobs)
+        batch_sum_logprobs_per_tok.append(sum_logprobs_per_tok)
+
+    entropy = sum(batch_sum_logprobs) / len(batch_sum_logprobs)
+    entropy_norm = sum(batch_sum_logprobs_per_tok) / len(batch_sum_logprobs_per_tok)
+
+    return entropy, entropy_norm

src/fairseq2/recipes/lm/_online_finetune/_grpo.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -22,6 +22,7 @@
 from fairseq2.recipes.lm._online_finetune._common import (
     compute_token_level_entropy,
     log_rollouts,
+    get_rollout_lengths,
 )
 
 from fairseq2.context import RuntimeContext
@@ -172,6 +173,14 @@
             log_rollouts(prompt_batch, rollouts, "Valid")
         reward_output = self._reward.process_rollouts(rollouts, prompt_batch)
         avg_reward = torch.tensor(reward_output["rewards"]).float().mean()
+
+        rollout_lengths = get_rollout_lengths(rollouts)
+        avg_rollout_length = torch.tensor(rollout_lengths).float().mean()
+        avg_reward_len_norm = avg_reward / avg_rollout_length
+
+        self._metric_bag.update_avg_rollout_length(avg_rollout_length)
+        self._metric_bag.update_avg_reward_len_norm(avg_reward_len_norm)
+
         self._metric_bag.update_avg_reward(avg_reward)
         self._metric_bag.update_batch_metrics(prompt_batch)
         # returning dummy loss since trainer expects it
@@ -386,6 +395,12 @@
         )
         self.register_metric("grpo_loss", Mean(device=gang.device), persistent=False)
         self.register_metric("avg_reward", Mean(device=gang.device), persistent=False)
+        self.register_metric(
+            "avg_rollout_length", Mean(device=gang.device), persistent=False
+        )
+        self.register_metric(
+            "avg_reward_len_norm", Mean(device=gang.device), persistent=False
+        )
         self.register_metric(
             "logit_entropy", Mean(device=gang.device), persistent=False
         )
@@ -426,6 +441,14 @@
     def update_avg_reward(self, avg_reward):
         self.avg_reward.update(avg_reward, weight=1)
 
+    @torch.inference_mode()
+    def update_avg_rollout_length(self, avg_rollout_length):
+        self.avg_rollout_length.update(avg_rollout_length, weight=1)
+
+    @torch.inference_mode()
+    def update_avg_reward_len_norm(self, avg_reward_len_norm):
+        self.avg_reward_len_norm.update(avg_reward_len_norm, weight=1)
+
     @torch.inference_mode()
     def update_batch_metrics(self, batch: PreferenceBatch):
         num_examples = batch.batch_size