Feature/track raw score

prompting/validators/event.py

@@ -47,13 +47,25 @@ class EventSchema:
         List[float]
     ]  # Output vector of the dahoas reward model
     blacklist_filter: Optional[List[float]]  # Output vector of the blacklist filter
+    blacklist_filter_matched_ngram: Optional[
+        List[str]
+    ]  # Output vector of the blacklist filter
+    blacklist_filter_significance_score: Optional[
+        List[float]
+    ]  # Output vector of the blacklist filter
     nsfw_filter: Optional[List[float]]  # Output vector of the nsfw filter
     reciprocate_reward_model: Optional[
         List[float]
     ]  # Output vector of the reciprocate reward model
     diversity_reward_model: Optional[
         List[float]
     ]  # Output vector of the diversity reward model
+    diversity_reward_model_historic: Optional[
+        List[float]
+    ]  # Output vector of the diversity reward model
+    diversity_reward_model_batch: Optional[
+        List[float]
+    ]  # Output vector of the diversity reward model
     dpo_reward_model: Optional[List[float]]  # Output vector of the dpo reward model
     rlhf_reward_model: Optional[List[float]]  # Output vector of the rlhf reward model
     prompt_reward_model: Optional[
@@ -68,6 +80,7 @@ class EventSchema:
         List[float]
     ]  # Output vector of the dahoas reward model
     nsfw_filter_normalized: Optional[List[float]]  # Output vector of the nsfw filter
+    nsfw_filter_score: Optional[List[float]]  # Output vector of the nsfw filter
     reciprocate_reward_model_normalized: Optional[
         List[float]
     ]  # Output vector of the reciprocate reward model
@@ -86,6 +99,12 @@ class EventSchema:
     relevance_filter_normalized: Optional[
         List[float]
     ]  # Output vector of the relevance scoring reward model
+    relevance_filter_bert_score: Optional[
+        List[float]
+    ]  # Output vector of the relevance scoring reward model
+    relevance_filter_mpnet_score: Optional[
+        List[float]
+    ]  # Output vector of the relevance scoring reward model
     task_validator_filter_normalized: Optional[List[float]]
 
     # Weights data
@@ -106,6 +125,8 @@ def from_dict(event_dict: dict, disable_log_rewards: bool) -> "EventSchema":
                 RewardModelType.reciprocate.value
             ),
             "diversity_reward_model": event_dict.get(RewardModelType.diversity.value),
+            "diversity_reward_model_historic": event_dict.get(RewardModelType.diversity.value + '_historic'),
+            "diversity_reward_model_batch": event_dict.get(RewardModelType.diversity.value + '_batch'),
             "dpo_reward_model": event_dict.get(RewardModelType.dpo.value),
             "rlhf_reward_model": event_dict.get(RewardModelType.rlhf.value),
             "prompt_reward_model": event_dict.get(RewardModelType.prompt.value),
@@ -136,6 +157,19 @@ def from_dict(event_dict: dict, disable_log_rewards: bool) -> "EventSchema":
             "prompt_reward_model_normalized": event_dict.get(
                 RewardModelType.prompt.value + "_normalized"
             ),
+            "blacklist_filter_matched_ngram": event_dict.get(
+                RewardModelType.blacklist.value + "_matched_ngram"
+            ),
+            "blacklist_filter_significance_score": event_dict.get(
+                RewardModelType.blacklist.value + "_significance_score"
+            ),
+            "relevance_filter_bert_score": event_dict.get(
+                RewardModelType.relevance.value + "_bert_score"
+            ),
+            "relevance_filter_mpnet_score": event_dict.get(
+                RewardModelType.relevance.value + "_mpnet_score"
+            ),
+            "nsfw_filter_score": event_dict.get(RewardModelType.nsfw.value + "_score"),
         }
 
         # Logs warning that expected data was not set properly

prompting/validators/forward.py

@@ -99,7 +99,9 @@ async def run_step(
     )
 
     # Update blacklist with completions so that n-gram filtering can be applied
-    self.blacklist.add([response.completion for response in responses if response.completion])
+    self.blacklist.add(
+        [response.completion for response in responses if response.completion]
+    )
 
     # Restrict the format of acceptable followup completions.
     for response in responses:
@@ -122,19 +124,19 @@ async def run_step(
         self.device
     )
     for weight_i, reward_fn_i in zip(self.reward_weights, self.reward_functions):
-        reward_i, reward_i_normalized = reward_fn_i.apply(prompt, responses, name)
+        reward_i_normalized, reward_event = reward_fn_i.apply(prompt, responses, name)
         rewards += weight_i * reward_i_normalized.to(self.device)
         if not self.config.neuron.disable_log_rewards:
-            event[reward_fn_i.name] = reward_i.tolist()
-            event[reward_fn_i.name + "_normalized"] = reward_i_normalized.tolist()
+            event = {**event, **reward_event}
         bt.logging.trace(str(reward_fn_i.name), reward_i_normalized.tolist())
 
     for masking_fn_i in self.masking_functions:
-        mask_i, mask_i_normalized = masking_fn_i.apply(base_prompt, responses, name)
+        mask_i_normalized, reward_event = masking_fn_i.apply(
+            base_prompt, responses, name
+        )
         rewards *= mask_i_normalized.to(self.device)  # includes diversity
         if not self.config.neuron.disable_log_rewards:
-            event[masking_fn_i.name] = mask_i.tolist()
-            event[masking_fn_i.name + "_normalized"] = mask_i_normalized.tolist()
+            event = {**event, **reward_event}
         bt.logging.trace(str(masking_fn_i.name), mask_i_normalized.tolist())
 
     # Train the gating model based on the predicted scores and the actual rewards.

prompting/validators/mock.py

@@ -59,7 +59,7 @@ def __init__(self, mock_name: str = "MockReward"):
 
     def apply(self, prompt: str, completion: List[str], name: str) -> torch.FloatTensor:
         mock_reward = torch.tensor([1 for _ in completion], dtype=torch.float32)
-        return mock_reward, mock_reward
+        return mock_reward, {}
 
     def reset(self):
         return self