Update llm/huggingfacehub/hf_model.py

llm/huggingfacehub/hf_model.py

@@ -1,6 +1,6 @@
-import logging
 import os
 import yaml
+import logging
 
 from abc import ABC
 
@@ -26,8 +26,12 @@ logger.addHandler(file_handler)
 
 logger.info("Getting information from hf_model module")
 
-path_to_prompts = os.path.join(os.getcwd(), "/llm/")
-print(path_to_prompts)
+try:
+   os.chdir('/home/user/app/llm/')
+except FileNotFoundError:
+   print("Error: Could not move up. You might be at the root directory.")
+
+work_dir = os.getcwd() 
 
 
 class HF_Mistaril(HFInterface, ABC):
@@ -37,18 +41,21 @@ class HF_Mistaril(HFInterface, ABC):
 
         self.model_config = config["HF_Mistrail"]
 
+        # Додати repetition_penalty, task?, top_p, stop_sequences
         self.llm = HuggingFaceEndpoint(
             repo_id=self.model_config["model"],
-            model_kwargs={"temperature": self.model_config["temperature"],
-                          "max_new_tokens": self.model_config["max_new_tokens"],
-                          "top_k": self.model_config["top_k"], "load_in_8bit": self.model_config["load_in_8bit"]})
+            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"]})
 
     @staticmethod
     def __read_yaml():
         try:
-            yaml_file = os.path.join(path_to_prompts, 'prompts.yaml')
-            with open(yaml_file, 'r') as file:
-                data = yaml.safe_load(file)
+            yaml_file = os.path.join(work_dir, 'prompts.yaml')
+            with open(yaml_file, 'r') as f:
+                data = yaml.safe_load(f)
+            f.close()
             return data
         except Exception as e:
             print(f"Execution filed : {e}")
@@ -84,16 +91,18 @@ class HF_TinyLlama(HFInterface, ABC):
 
         self.llm = HuggingFaceEndpoint(
             repo_id=self.model_config["model"],
-            model_kwargs={"temperature": self.model_config["temperature"],
-                          "max_new_tokens": self.model_config["max_new_tokens"],
-                          "top_k": self.model_config["top_k"], "load_in_8bit": self.model_config["load_in_8bit"]})
+            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"]})
 
     @staticmethod
     def __read_yaml():
         try:
-            yaml_file = os.path.join(path_to_prompts, 'prompts.yaml')
-            with open(yaml_file, 'r') as file:
-                data = yaml.safe_load(file)
+            yaml_file = os.path.join(work_dir, 'prompts.yaml')
+            with open(yaml_file, 'r') as f:
+                data = yaml.safe_load(f)
+            f.close()
             return data
         except Exception as e:
             print(f"Execution filed : {e}")