Add global config and model factory

config.ini

@@ -0,0 +1,34 @@
+[server]
+port = int:42069
+
+[models.names]
+gpt2 = str:openai-community/gpt2
+gpt2_medium = str:openai-community/gpt2-medium
+gpt2_large = str:openai-community/gpt2-large
+gpt2_xl = str:openai-community/gpt2-xl
+
+llama_3_8b_intruct = str:meta-llama/Meta-Llama-3-8B-Instruct
+llama_3_70b_instruct = str:meta-llama/Meta-Llama-3-70B-Instruct
+
+[models.params]
+dtype = str:bfloat16
+run_device = str:cpu
+load_device = str:cuda
+
+[encrypt.default]
+gen_model = str:gpt2
+start_pos = int:0
+gamma = float:10.0
+msg_base = int:2
+seed_scheme = str:sha_left_hash
+window_length = int:1
+private_key = int:0
+max_new_tokens_ratio = float:2.0
+num_beams = int:4
+
+[decrypt.default]
+gen_model = str:gpt2
+msg_base = int:2
+seed_scheme = str:sha_left_hash
+window_length = int:1
+private_key = int:0

global_config.py

@@ -0,0 +1,36 @@
+import configparser
+from utils import static_init
+
+
+@static_init
+class GlobalConfig:
+    default_file_name = "config.ini"
+    config = configparser.ConfigParser()
+
+    @classmethod
+    def get_section(cls, section_name):
+        if section_name in cls.config :
+            return cls.config[section_name].keys()
+        else:
+            return None
+
+    @classmethod
+    def get(cls, section_name, attr_name):
+        if section_name in cls.config and attr_name in cls.config[section_name]:
+            value = cls.config.get(section_name, attr_name)
+            value = value.split(":")
+            type_name = value[0]
+            value = ":".join(value[1:])
+            if type_name == "str":
+                value = str(value)
+            elif type_name == "float":
+                value = float(value)
+            elif type_name == "int":
+                value = int(value)
+            return value
+        else:
+            return None
+
+    @classmethod
+    def __static_init__(cls):
+        cls.config.read(cls.default_file_name)

model_factory.py

@@ -0,0 +1,72 @@
+import torch
+import transformers
+from utils import load_model, static_init
+from global_config import GlobalConfig
+
+
+@static_init
+class ModelFactory:
+    models_names = {}
+    models = {}
+    tokenizers = {}
+    run_model = None
+    dtype = torch.bfloat16
+    load_device = torch.device("cpu")
+    run_device = torch.device("cpu")
+
+    @classmethod
+    def __static_init__(cls):
+        names_sec = GlobalConfig.get_section("models.names")
+        if names_sec is not None:
+            for name in names_sec:
+                cls.models_names[name] = GlobalConfig.get("models.names", name)
+
+        if GlobalConfig.get_section("models.params") is not None:
+            dtype = GlobalConfig.get("models.params", "dtype")
+            if dtype == "bfloat16":
+                cls.dtype = torch.bfloat16
+            elif dtype == "float16":
+                cls.dtype = torch.float16
+            elif dtype == "float32":
+                cls.dtype = torch.float32
+
+            load_device = GlobalConfig.get("models.params", "load_device")
+            run_device = GlobalConfig.get("models.params", "run_device")
+            if load_device is not None:
+                cls.load_device = torch.device(str(load_device))
+            if run_device is not None:
+                cls.run_device = torch.device(str(run_device))
+
+    @classmethod
+    def __load_model(cls, name):
+        if name not in cls.models_names:
+            print(f"{name} is not a valid model name")
+            return None
+
+        if name not in cls.models:
+            model, tokenizer = load_model(
+                cls.models_names[name], cls.load_device
+            )
+            cls.models[name] = model
+            cls.tokenizers[name] = tokenizer
+        else:
+            model, tokenizer = cls.models[name], cls.tokenizers[name]
+
+        return model, tokenizer
+
+    @classmethod
+    def load_model(cls, name):
+        if name not in cls.models:
+            cls.__load_model(name)
+
+        if name != cls.run_model and cls.run_model is not None:
+            cls.models[cls.run_model].to(cls.load_device)
+
+        cls.models[name].to(cls.run_device)
+        cls.run_model = name
+
+        return cls.models[name], cls.tokenizers[name]
+
+    @classmethod
+    def get_models_names(cls):
+        return list(cls.models_names.keys())