Save only flag

Author: BabyChouSr

fastchat/serve/gradio_block_arena_vision.py

@@ -241,9 +241,9 @@ def add_text(state, model_selector, chat_input, context: Context, request: gr.Re
     )
 
     text_flagged, nsfw_flag, csam_flag = (
-        moderation_type_to_response_map["text_moderation"]["flagged"],
-        moderation_type_to_response_map["nsfw_moderation"]["flagged"],
-        moderation_type_to_response_map["csam_moderation"]["flagged"],
+        state.content_moderator.text_flagged,
+        state.content_moderator.nsfw_flagged,
+        state.content_moderator.csam_flagged,
     )
 
     if csam_flag:

fastchat/serve/gradio_block_arena_vision_anony.py

@@ -333,15 +333,13 @@ def add_text(
 
     # Use the first state to get the moderation response because this is based on user input so it is independent of the model
     moderation_image_input = images[0] if len(images) > 0 else None
-    moderation_type_to_response_map = states[
-        0
-    ].content_moderator.image_and_text_moderation_filter(
+    moderation_type_to_response_map = states[0].content_moderator.image_and_text_moderation_filter(
         moderation_image_input, text, model_list, do_moderation=True
     )
     text_flagged, nsfw_flag, csam_flag = (
-        moderation_type_to_response_map["text_moderation"]["flagged"],
-        moderation_type_to_response_map["nsfw_moderation"]["flagged"],
-        moderation_type_to_response_map["csam_moderation"]["flagged"],
+        states[0].content_moderator.text_flagged,
+        states[0].content_moderator.nsfw_flagged,
+        states[0].content_moderator.csam_flagged,
     )
 
     if csam_flag:

fastchat/serve/gradio_block_arena_vision_named.py

@@ -257,16 +257,13 @@ def add_text(
 
     # Use the first state to get the moderation response because this is based on user input so it is independent of the model
     moderation_image_input = images[0] if len(images) > 0 else None
-    moderation_type_to_response_map = states[
-        0
-    ].content_moderator.image_and_text_moderation_filter(
+    moderation_type_to_response_map = states[0].content_moderator.image_and_text_moderation_filter(
         moderation_image_input, text, model_list, do_moderation=False
     )
-
     text_flagged, nsfw_flag, csam_flag = (
-        moderation_type_to_response_map["text_moderation"]["flagged"],
-        moderation_type_to_response_map["nsfw_moderation"]["flagged"],
-        moderation_type_to_response_map["csam_moderation"]["flagged"],
+        states[0].content_moderator.text_flagged,
+        states[0].content_moderator.nsfw_flagged,
+        states[0].content_moderator.csam_flagged,
     )
 
     if csam_flag:

fastchat/serve/moderation/moderator.py

@@ -136,9 +136,11 @@ def image_moderation_filter(self, image: Image):
         self.nsfw_flagged = nsfw_flagged_map["flagged"]
         self.csam_flagged = csam_flagged_map["flagged"]
 
+        # We save only the boolean value instead of the entire response dictionary
+        # to save space. nsfw_flagged_map and csam_flagged_map will contain the whole dictionary
         return {
-            "nsfw_moderation": nsfw_flagged_map,
-            "csam_moderation": csam_flagged_map,
+            "nsfw_moderation": {"flagged": self.nsfw_flagged},
+            "csam_moderation": {"flagged": self.csam_flagged},
         }
 
     def _openai_moderation_filter(
@@ -205,12 +207,17 @@ def text_moderation_filter(
 
         moderation_response_map = {"flagged": False}
         if do_moderation:
+            # We save the entire response dictionary here
             moderation_response_map = self._openai_moderation_filter(
                 text, custom_thresholds
             )
             self.text_flagged = moderation_response_map["flagged"]
+        else:
+            self.text_flagged = False
 
-        return {"text_moderation": moderation_response_map}
+        # We only save whether the text was flagged or not instead of the entire response dictionary
+        # to save space. moderation_response_map will contain the whole dictionary
+        return {"text_moderation": {"flagged": self.text_flagged}}
 
     def image_and_text_moderation_filter(
         self, image: Image, text: str, model_list: List[str], do_moderation=True