Merge branch 'Develop' into sfw

Author: Lyaaaaaaaaaaaaaaa

server/config.py

@@ -47,6 +47,9 @@
 #--
 #--   23/01/2024 Lyaaaaa
 #--     - Removed TOKENIZERS_PATH.
+#--
+#--   08/05/2024 Lyaaaaa
+#--     - Added TORCH_DTYPE_SAFETY.
 #---------------------------------------------------------------------------
 
 import logging
@@ -79,3 +82,5 @@
 DEVICE_MAP        = None # None/see documentation
 TORCH_DTYPE       = None # "Auto"/None/torch.dtype/See torch_dtype.py for more info.
 
+# Safeguards
+TORCH_DTYPE_SAFETY = True # True/False. If CUDA isn't available, will enforce Torch_Dtype to float32 to avoir error. See issue https://github.com/LyaaaaaGames/AIdventure_Server/issues/31

server/model.py

@@ -203,9 +203,16 @@
 #--      - 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.
+#--        Because otherwise, it will lead to an error during generation.
+#--        See https://github.com/LyaaaaaGames/AIdventure_Server/issues/31
 #------------------------------------------------------------------------------
 
 from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer
@@ -327,6 +334,11 @@ def _load_model(self):
       logger.log.debug("Model settings:")
       logger.log.debug(args)
 
+      if not self.is_cuda_available and config.TORCH_DTYPE_SAFETY:
+        logger.log.warn("Cuda isn't available.")
+        logger.log.warn("Setting torch_dtype to float 32 to avoid error.")
+        args["torch_dtype"] = Torch_Dtypes.dtypes.value[Torch_Dtypes.FLOAT_32.value]
+
       self._Model = AutoModelForCausalLM.from_pretrained(self._model_path,
                                                          **args)
     except Exception as e: