[BlenderBot2] clean up reply during interactive mode

projects/blenderbot2/agents/blenderbot2.py

@@ -13,6 +13,7 @@
 The Memory Decoder examines the context and generates memories to write to
 the long-term memory module.
 """
+import re
 import copy
 import torch
 import torch.nn
@@ -898,6 +899,38 @@ def compute_loss(
         else:
             return loss
 
+    def _clean_text(self, txt):
+        cleaned_txt = re.sub(r'_[\S]*unsafe_*', '', txt, flags=re.IGNORECASE)
+        return cleaned_txt.strip()
+
+    def self_observe(self, self_message: Message) -> None:
+        """
+        Observe one's own utterance. Override TorchAgent.self_observe with cleaned text.
+
+        :param self_message:
+            The message corresponding to the output from batch_act.
+        """
+        use_reply = self.opt.get('use_reply', 'label')
+
+        if (
+            self.observation['episode_done']  # last example in the episode
+            or use_reply == 'none'  # not including our own responses anyway
+            or (
+                use_reply == 'label'
+                and any([l in self.observation for l in ['labels', 'eval_labels']])
+            )  # has true label
+        ):
+            return super().self_observe(self_message)
+
+        # otherwise, we use the CLEANED last output the model generated
+        if self_message is not None:
+            last_reply = self_message['text']
+            clean_reply = self._clean_text(last_reply)
+            self.history.add_reply(clean_reply)
+            return
+
+        raise RuntimeError("Unexpected case in self_observe.")
+
 
 class BlenderBot2FidAgent(FidAgent, BlenderBot2RagAgent):
     model: BlenderBot2FidModel