Update llm/huggingfacehub/hf_model.py

llm/huggingfacehub/hf_model.py

@@ -33,10 +33,10 @@ logger.info("Getting information from hf_model module")
 
 # work_dir = os.getcwd() 
 
-print("CWD : ", os.getcwd())
-
 llm_dir = '/home/user/app/llm/'
 
+print("Path to prompts : ", os.path.join(os.getcwd(), "prompts.yaml"))
+
 
 class HF_Mistaril(HFInterface, ABC):
     def __init__(self, prompt_entity: str, prompt_id: int = 0):
@@ -51,7 +51,8 @@ class HF_Mistaril(HFInterface, ABC):
             temperature=self.model_config["temperature"],
             max_new_tokens=self.model_config["max_new_tokens"],
             top_k=self.model_config["top_k"],
-            model_kwargs={"load_in_8bit": self.model_config["load_in_8bit"]}
+            model_kwargs={"load_in_8bit": self.model_config["load_in_8bit"]},
+            huggingfacehub_api_token=os.environ.get("HUGGINGFACEHUB_API_TOKEN")
         )
 
     @staticmethod
@@ -99,7 +100,8 @@ class HF_TinyLlama(HFInterface, ABC):
             temperature=self.model_config["temperature"],
             max_new_tokens=self.model_config["max_new_tokens"],
             top_k=self.model_config["top_k"],
-            model_kwargs={"load_in_8bit": self.model_config["load_in_8bit"]}
+            model_kwargs={"load_in_8bit": self.model_config["load_in_8bit"]},
+            huggingfacehub_api_token=os.environ.get("HUGGINGFACEHUB_API_TOKEN")
         )
 
     @staticmethod