Update services/llama_generator.py

services/llama_generator.py

@@ -45,19 +45,9 @@ class LlamaGenerator(BaseGenerator):
     ):
         print(llama_model_name)
         print(prm_model_path)
-        @observe()
-        def load_model(self, model_name: str):
-            # Code to load your model, e.g., Hugging Face's transformers library
-            from transformers import AutoModelForCausalLM
-            return AutoModelForCausalLM.from_pretrained(model_name)
 
-        @observe()
-        def load_tokenizer(self, model_name: str):
-            # Load the tokenizer associated with the model
-            from transformers import AutoTokenizer
-            return AutoTokenizer.from_pretrained(model_name)
             
-        self.tokenizer = load_tokenizer(llama_model_name)  # Add this line to initialize the tokenizer
+        self.tokenizer = model_manager.load_tokenizer(llama_model_name)  # Add this line to initialize the tokenizer
  
         super().__init__(
             llama_model_name,