Merge pull request #33 from LyaaaaaGames/sfw

Implemented the banned_words feature.

Author: Lyaaaaaaaaaaaaaaa

server/generator.py

@@ -40,6 +40,10 @@
 #--  - 31/01/2024 Lyaaaaa
 #--    - generate_text now longer receives memory and context as parameters.
 #--        They are embedded in the prompt parameter by the client.
+#--
+#--  - 07/05/2024 Lyaaaaa
+#--    - Updated generate_text to now be able to censor generation. The words
+#--        passed in p_banned_words parameters won't be generated anymore.
 #------------------------------------------------------------------------------
 
 from model        import Model
@@ -48,16 +52,24 @@
 import logger
 
 class Generator(Model):
-
 #------------------------------------------------------------------------------
 #-- generate_text
 #------------------------------------------------------------------------------
   def generate_text(self,
-                    p_prompt     = None,
-                    p_parameters = None):
+                    p_prompt       = None,
+                    p_parameters   = None,
+                    p_banned_words = []):
 
     model_input    = self._Tokenizer(p_prompt, return_tensors = "pt")
 
+    if p_banned_words:
+      banned_words_ids = self._Tokenizer(
+        p_banned_words,
+        add_special_tokens=False
+        ).input_ids
+
+      p_parameters["bad_words_ids"] = banned_words_ids
+
     if self.is_cuda_available:
       logger.log.info("Loading inputs to GPU")
       model_input.to("cuda")

server/model.py

@@ -203,6 +203,11 @@
 #--      - Added a log message to display the model's name and its path.
 #--      - Added a log message to display if cuda is supported.
 #--
+#--
+#--  - 07/05/2024 Lyaaaaa
+#--    - Updated _load_tokens to set add_prefix_space to True. It is needed
+#--        for using bad_words_ids parameter for generation.
+#--
 #--  - 08/05/2024 Lyaaaaa
 #--    - Updated _load_model to force (if config.TORCH_DTYPE_SAFETY is True)
 #--        torch_dtype to be set to float32 if cuda isn't available.
@@ -302,7 +307,10 @@ def _load(self):
 #------------------------------------------------------------------------------
   def _load_tokens(self):
     try:
-      self._Tokenizer = AutoTokenizer.from_pretrained(self._model_path)
+      self._Tokenizer = AutoTokenizer.from_pretrained(
+        self._model_path,
+        add_prefix_space=True
+        )
     except Exception as e:
       logger.log.error("Error loading tokens in " + self._model_path)
       logger.log.error(e)

server/server.py

@@ -118,6 +118,10 @@
 #--  - 31/01/2024 Lyaaaaa
 #--    - generate_text now longer receives memory and context as parameters.
 #--        They are embedded in the prompt parameter by the client.
+#--
+#--  - 07/05/2024 Lyaaaaa
+#--    - Updated handle_request and generation case to receive a banned_words
+#--        parameter and pass it to generator.generate_text
 #------------------------------------------------------------------------------
 
 import asyncio
@@ -177,10 +181,11 @@ def handle_request(p_websocket, p_data : dict):
   request = p_data['request']
 
   if request == Request.TEXT_GENERATION.value:
-    prompt     = p_data['prompt']
-    parameters = p_data['parameters']
+    prompt       = p_data['prompt']
+    parameters   = p_data['parameters']
+    banned_words = p_data['banned_words']
 
-    generated_text = generator.generate_text(prompt, parameters)
+    generated_text = generator.generate_text(prompt, parameters, banned_words)
 
     p_data["generated_text"] = generated_text