Updated config.py, model.py, server.py and translator.py

config.py:
- Removed TOKENIZERS_PATH.

model.py:
- __init__ now expect a fourth parameter, model_path.
- _model_path is now set to p_model_path.
- Removed _tokenizers_path. A model's files are now all in the same
folder.
- Extracted _load_translator to the Translator class to use
polymorphism.
- Updated _load to always call _load_translator no matter the type of
the model. The specific behavior is implemented in children if needed
(See Translator.py)
- _load_tokens and _save_tokens now use _model_path.

server.py:
- Updated handle_request to receive the models path in p_data and to use
it in the model constructor.

translator.py:
- Added _load_model method (extracted from model.py. Old name
"_load_translator")

server/config.py

@@ -44,6 +44,9 @@
 #--   18/09/2023 Lyaaaaa
 #--     - LOG_FILEMODE default value is now "a" again. The log file is now
 #--         manually deleted to avoid losing logs.
+#--
+#--   23/01/2024 Lyaaaaa
+#--     - Removed TOKENIZERS_PATH.
 #---------------------------------------------------------------------------
 
 import logging
@@ -61,7 +64,6 @@
 # Models.
 #See possible values here: https://huggingface.co/docs/transformers/main_classes/model#transformers.PreTrainedModel.from_pretrained
 
-TOKENIZERS_PATH    = "models/"
 MODELS_PATH        = "models/"
 DEFAULT_MODEL      = "EleutherAI/gpt-neo-125M"
 ALLOW_DOWNLOAD     = None # True/False/None. If True, the server will download AI's files.

server/model.py

@@ -174,6 +174,15 @@
 #--    - Fixed an error in _download_model. It used model_name which doesn't
 #--        exist. Now it uses self._model_name.
 #--    - Updated the logs in  _load_model and _load_tokens.
+#--
+#--  - 23/01/2024 Lyaaaaa
+#--    - __init__ now expect a fourth parameter, model_path.
+#--    - _model_path is now set to p_model_path.
+#--    - Removed _tokenizers_path. A model's files are now all in the same folder.
+#--    - Extracted _load_translator to the Translator class to use polymorphism.
+#--    - Updated _load to always call _load_translator no matter the type of the model.
+#--        The specific behavior is implemented in children if needed (See Translator.py)
+#--    - _load_tokens and _save_tokens now use _model_path.
 #------------------------------------------------------------------------------
 
 from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer
@@ -196,7 +205,6 @@ class Model():
   _model_type : Model_Type
   _Model = None
 
-  _tokenizers_path = config.TOKENIZERS_PATH
   _model_path      = config.MODELS_PATH
   _allow_offload   = config.ALLOW_OFFLOAD
   _limit_memory    = config.LIMIT_MEMORY
@@ -215,10 +223,10 @@ class Model():
   def __init__(self,
                p_model_name = config.DEFAULT_MODEL,
                p_model_type = Model_Type.GENERATION.value,
-               p_parameters = {}):
+               p_parameters = {},
+               p_model_path = config.MODELS_PATH):
 
-    self._tokenizers_path  += p_model_name
-    self._model_path       += p_model_name
+    self._model_path       = p_model_path
     self._model_name       = p_model_name
     self.is_cuda_available = torch.cuda.is_available()
     self._model_type       = p_model_type
@@ -278,12 +286,7 @@ def _load(self):
     else:
       logger.log.info("Tokens successfully loaded from local files")
 
-
-    if self._model_type == Model_Type.GENERATION.value:
-      model_loaded = self._load_model()
-    elif self._model_type == Model_Type.TRANSLATION.value:
-      model_loaded = self._load_translator()
-
+    model_loaded = self._load_model()
 
     if model_loaded == False:
       if self._allow_download == True:
@@ -301,9 +304,9 @@ def _load(self):
 #------------------------------------------------------------------------------
   def _load_tokens(self):
     try:
-      self._Tokenizer = AutoTokenizer.from_pretrained(self._tokenizers_path)
+      self._Tokenizer = AutoTokenizer.from_pretrained(self._model_path)
     except Exception as e:
-      logger.log.error("Error loading tokens in " + self._tokenizers_path)
+      logger.log.error("Error loading tokens in " + self._model_path)
       logger.log.error(e)
       return False
 
@@ -335,26 +338,11 @@ def _load_model(self):
     return True
 
 
-#------------------------------------------------------------------------------
-#--
-#------------------------------------------------------------------------------
-  def _load_translator(self):
-    try:
-      self._Model = AutoModelForSeq2SeqLM.from_pretrained(self._model_path)
-
-    except Exception as e:
-      logger.log.error("An unexpected error happened while loading the translator")
-      logger.log.error(e)
-      return False
-
-    return True
-
-
 #------------------------------------------------------------------------------
 #--
 #------------------------------------------------------------------------------
   def _save_tokens(self):
-    self._Tokenizer.save_pretrained(self._tokenizers_path)
+    self._Tokenizer.save_pretrained(self._model_path)
 
 
 #------------------------------------------------------------------------------

server/server.py

@@ -104,6 +104,10 @@
 #--
 #--  - 25/10/2023 Lyaaaaa
 #--    - Fixed syntax error in handler. See https://github.com/LyaaaaaGames/AIdventure_Server/issues/25
+#--
+#--  - 23/01/2024 Lyaaaaa
+#--    - Updated handle_request to receive the models path in p_data and to use
+#--        it in the model constructor.
 #------------------------------------------------------------------------------
 
 import asyncio
@@ -194,25 +198,31 @@ def handle_request(p_websocket, p_data : dict):
                     "offload_dict"    : p_data['offload_dict'],}
       del generator
       logger.log.debug("loading generator")
-      model_name = p_data['model_name']
+      model_name = p_data["model_name"]
+      model_path = p_data["model_path"]
+
 
       generator = Generator(model_name,
-                             Model_Type.GENERATION.value,
-                             parameters)
+                            Model_Type.GENERATION.value,
+                            parameters,
+                            model_path)
       logger.log.info("Is CUDA available: " + format(generator.is_cuda_available))
 
     elif p_data["model_type"] == Model_Type.TRANSLATION.value:
       parameters = {"low_memory_mode" : p_data['low_memory_mode']}
       logger.log.debug("loading translator")
       model_name = p_data["to_eng_model"]
+      model_path = p_data["model_path"]
       to_eng_translator = Translator(model_name,
                                      Model_Type.TRANSLATION.value,
-                                     parameters)
+                                     parameters,
+                                     model_path)
 
       model_name = p_data["from_eng_model"]
       from_eng_translator = Translator(model_name,
                                        Model_Type.TRANSLATION.value,
-                                       parameters)
+                                       parameters,
+                                       model_path)
 
     p_data['request'] = Request.LOADED_MODEL.value
     p_data            = Json_Utils.json_to_string(p_data)

server/translator.py

@@ -30,6 +30,9 @@
 #--    - Removed inputs.to("cuda"). The translators won't use the gpu anymore (for now).
 #--        The new system isn't supported by the MarianMT models and the generator
 #--        is more in need of the GPU.
+#--
+#--  - 23/01/2024 Lyaaaaa
+#--    - Added _load_model method (extracted from model.py. Old name "_load_translator")
 #------------------------------------------------------------------------------
 
 from model import Model
@@ -48,3 +51,18 @@ def translate_text(self, p_input : str, p_parameters = {}):
     generated_text = self._Tokenizer.batch_decode(outputs, skip_special_tokens=True, use_cache=True)
 
     return generated_text[0]
+
+
+#------------------------------------------------------------------------------
+#--
+#------------------------------------------------------------------------------
+  def _load_model(self):
+    try:
+      self._Model = AutoModelForSeq2SeqLM.from_pretrained(self._model_path)
+
+    except Exception as e:
+      logger.log.error("An unexpected error happened while loading the translator")
+      logger.log.error(e)
+      return False
+
+    return True