Jacklanchantin/online training

src/fairseq2/metrics/recorders/_wandb.py

@@ -45,6 +45,7 @@ def __init__(
         name: str,
         output_dir: Path,
         metric_descriptors: Provider[MetricDescriptor],
+        id: str | None = None,
     ) -> None:
         """
         :param project: The W&B project name.
@@ -60,7 +61,7 @@ def __init__(
             self._run = None
         else:
             self._run = wandb.init(
-                project=project, name=name, dir=output_dir.parent, resume="allow"
+                project=project, name=name, id=id, dir=output_dir.parent, resume="allow"
             )
 
         self._metric_descriptors = metric_descriptors
@@ -77,6 +78,12 @@ def record_metrics(
         if self._run is None:
             return
 
+        # try:
+        #     self._run.log({"_step": step_nr})  # Log to the specific step
+        #     self._run.step = step_nr  # Directly update the internal step counter
+        # except:
+        #     ...
+
         for name, value in values.items():
             try:
                 descriptor = self._metric_descriptors.get(name)
@@ -88,6 +95,8 @@ def record_metrics(
             else:
                 display_name = descriptor.display_name
 
+            display_name = run + "/" + display_name
+
             try:
                 self._run.log({display_name: value}, step=step_nr)
             except RuntimeError as ex:
@@ -112,6 +121,8 @@ class WandbRecorderConfig:
 
     run: str | None = None
 
+    id: str | None = None
+
     def validate(self) -> None:
         result = ValidationResult()
 
@@ -151,7 +162,11 @@ def create(self, output_dir: Path, config: object) -> MetricRecorder:
         wandb_dir = output_dir.joinpath("wandb")
 
         return WandbRecorder(
-            config.project, config.run, wandb_dir, self._metric_descriptors
+            config.project,
+            config.run,
+            wandb_dir,
+            self._metric_descriptors,
+            id=config.id,
         )
 
     @property

src/fairseq2/recipes/common/_trainer.py

@@ -96,6 +96,7 @@ def create_trainer(
         valid_data_readers=valid_data_readers,
         validate_after_n_steps=regime_section.validate_after_n_steps,
         validate_every_n_steps=regime_section.validate_every_n_steps,
+        validate_step_0=regime_section.validate_step_0,
         validate_after_n_data_epochs=regime_section.validate_after_n_data_epochs,
         validate_every_n_data_epochs=regime_section.validate_every_n_data_epochs,
         checkpoint_manager=checkpoint_manager,

src/fairseq2/recipes/config.py

@@ -219,6 +219,9 @@ class RegimeSection:
     validate_every_n_steps: int | None = None
     """The step interval at which to validate the model."""
 
+    validate_step_0: bool = False
+    """Validate before training"""
+
     validate_after_n_data_epochs: int = 0
 
     validate_every_n_data_epochs: int | None = None

src/fairseq2/recipes/lm/__init__.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -161,6 +161,14 @@
 from fairseq2.recipes.lm._online_finetune._rewards import (
     SkyworkVerifierHandler as SkyworkVerifierHandler,
 )
+
+from fairseq2.recipes.lm._online_finetune._rewards import (
+    AtheneVerifier as AtheneVerifier,
+)
+from fairseq2.recipes.lm._online_finetune._rewards import (
+    AtheneVerifierHandler as AtheneVerifierHandler,
+)
+
 from fairseq2.recipes.lm._online_finetune._rewards import (
     NuminaMathVerifier as NuminaMathVerifier,
 )

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