Update llm/huggingfacehub/hf_model.py

llm/huggingfacehub/hf_model.py

@@ -9,14 +9,14 @@ from llm.config import config
 
 from langchain.prompts import PromptTemplate
 from langchain.chains import LLMChain
-from langchain.llms import HuggingFaceEndpoint
+from langchain.llms import HuggingFaceHub
 
 logger = logging.getLogger(__name__)
 
-logger.setLevel(logging.CRITICAL)  # because if something went wrong in execution application can't be work anymore
+logger.setLevel(logging.CRITICAL)  # because if something went wrong in execution, application can't be work anyway
 
 file_handler = logging.FileHandler(
-    "logs/chelsea_llm_huggingfacehub.log")  # for all modules here template for logs file is "../logs/chelsea_{module_name}_{dir_name}.log"
+    "logs/chelsea_llm_huggingfacehub.log")  # for all modules here template for logs file is "llm/logs/chelsea_{module_name}_{dir_name}.log"
 logger.setLevel(logging.INFO)  # informed
 
 formatted = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
@@ -26,16 +26,11 @@ logger.addHandler(file_handler)
 
 logger.info("Getting information from hf_model module")
 
-# 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() 
-
 llm_dir = '/home/user/app/llm/'
 
-print("Path to prompts : ", os.path.join(os.getcwd(), "prompts.yaml"))
+path_to_yaml = os.path.join(os.getcwd(), "llm/prompts.yaml")
+
+print("Path to prompts : ", path_to_yaml)
 
 
 class HF_Mistaril(HFInterface, ABC):
@@ -46,12 +41,16 @@ class HF_Mistaril(HFInterface, ABC):
         self.model_config = config["HF_Mistrail"]
 
         # Додати repetition_penalty, task?, top_p, stop_sequences
-        self.llm = HuggingFaceEndpoint(
+        self.llm = HuggingFaceHub(
             repo_id=self.model_config["model"],
-            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"]},
+            # 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"], 
+                          "temperature": self.model_config["temperature"],
+                          "max_new_tokens": self.model_config["max_new_tokens"],
+                          "top_k": self.model_config["top_k"],
+                         },
             huggingfacehub_api_token=os.environ.get("HUGGINGFACEHUB_API_TOKEN")
         )
 
@@ -95,12 +94,16 @@ class HF_TinyLlama(HFInterface, ABC):
 
         self.model_config = config["HF_TinyLlama"]
 
-        self.llm = HuggingFaceEndpoint(
+        self.llm = HuggingFaceHub(
             repo_id=self.model_config["model"],
-            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"]},
+            # 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"], 
+                          "temperature": self.model_config["temperature"],
+                          "max_new_tokens": self.model_config["max_new_tokens"],
+                          "top_k": self.model_config["top_k"],
+                         },
             huggingfacehub_api_token=os.environ.get("HUGGINGFACEHUB_API_TOKEN")
         )