Update my_model/LLAMA2/LLAMA2_model.py

my_model/LLAMA2/LLAMA2_model.py

@@ -3,7 +3,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from typing import Optional
 import bitsandbytes  # only for using on GPU
 import accelerate  # only for using on GPU
-from my_model.config import LLAMA2_config as config  # Importing LLAMA2 configuration file
+from my_model.config import LLAMA2_config as config
 import warnings
 
 # Suppress only FutureWarning from transformers
@@ -32,6 +32,7 @@ class Llama2ModelManager:
         """
         Initializes the Llama2ModelManager class with configuration settings.
         """
+        
         self.device: str = config.DEVICE
         self.model_name: str = config.MODEL_NAME
         self.tokenizer_name: str = config.TOKENIZER_NAME
@@ -51,6 +52,7 @@ class Llama2ModelManager:
         Returns:
             BitsAndBytesConfig: Configuration for BitsAndBytes optimized model.
         """
+        
         if self.quantization == '4bit':
             return BitsAndBytesConfig(
                 load_in_4bit=True,
@@ -68,11 +70,13 @@ class Llama2ModelManager:
 
     def load_model(self) -> AutoModelForCausalLM:
         """
-        Loads the LLaMA-2 model based on the specified configuration. If the model is already loaded, returns the existing model.
+        Loads the LLaMA-2 model based on the specified configuration. 
+        If the model is already loaded, returns the existing model.
 
         Returns:
             AutoModelForCausalLM: Loaded LLaMA-2 model.
         """
+        
         if self.model is not None:
             print("Model is already loaded.")
             return self.model
@@ -99,6 +103,7 @@ class Llama2ModelManager:
         Returns:
             AutoTokenizer: Loaded tokenizer for LLaMA-2 model.
         """
+        
         self.tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name, use_fast=self.use_fast,
                                                        token=self.access_token,
                                                        trust_remote_code=self.trust_remote,
@@ -111,12 +116,17 @@ class Llama2ModelManager:
 
         return self.tokenizer
 
-    def load_model_and_tokenizer(self, for_fine_tuning):
+    def load_model_and_tokenizer(self, for_fine_tuning: bool) -> Tuple[AutoModelForCausalLM, AutoTokenizer]:
         """
-        Loads LLAMa2 model and tokenizer in one method and adds special tokens if the purpose if fine tuning.
-        :param for_fine_tuning: YES(True) / NO (False)
-        :return: LLAMA2 Model and Tokenizer
+        Loads the LLaMA-2 model and tokenizer, and optionally adds special tokens for fine-tuning.
+
+        Args:
+            for_fine_tuning (bool): Whether to prepare the model and tokenizer for fine-tuning.
+
+        Returns:
+            Tuple[AutoModelForCausalLM, AutoTokenizer]: The loaded model and tokenizer.
         """
+        
         if for_fine_tuning:
             self.tokenizer = self.load_tokenizer()
             self.model = self.load_model()
@@ -128,17 +138,17 @@ class Llama2ModelManager:
         return self.model, self.tokenizer
 
 
-    def add_special_tokens(self, tokens: Optional[list[str]] = None) -> None:
+    def add_special_tokens(self, tokens: Optional[List[str]] = None) -> None:
         """
-        Adds special tokens to the tokenizer and updates the model's token embeddings if the model is loaded,
-        only if the tokenizer is loaded.
+        Adds special tokens to the tokenizer and updates the model's token embeddings if the model is loaded.
 
         Args:
-            tokens (list of str, optional): Special tokens to add. Defaults to a predefined set.
+            tokens (Optional[List[str]]): Special tokens to add. Defaults to a predefined set.
 
         Returns:
             None
         """
+        
         if self.tokenizer is None:
             print("Tokenizer is not loaded. Cannot add special tokens.")
             return
@@ -166,7 +176,8 @@ class Llama2ModelManager:
 
 
 if __name__ == "__main__":
-    pass
+    pass # uncomment to to load the mode and tokenizer and add the designed special tokens.
+    
     LLAMA2_manager = Llama2ModelManager()
     LLAMA2_model = LLAMA2_manager.load_model()  # First time loading the model
     LLAMA2_tokenizer = LLAMA2_manager.load_tokenizer()