Update hf_model.py

hf_model.py

@@ -1,58 +1,94 @@
-
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 from datetime import datetime
 import os
+from MewUtilities.mew_log import log_info, log_error  # Import your custom logging methods
 
 class HFModel:
     def __init__(self, model_name):
         parts = model_name.split("/")
         self.friendly_name = parts[1]
-        self.model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.bfloat16).cuda()
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
         self.chat_history = []
         self.log_file = f"chat_log_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md"
 
+        try:
+            log_info(f"=== Loading Model: {self.friendly_name} ===")
+            self.model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.bfloat16).cuda()
+            self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+            log_info(f"=== Model Loaded Successfully: {self.friendly_name} ===")
+        except Exception as e:
+            log_error(f"ERROR Loading Model: {e}")
+            raise
+
     def generate_response(self, input_text, max_length=100, skip_special_tokens=True):
-        inputs = self.tokenizer(input_text, return_tensors="pt").to(self.model.device)
-        outputs = self.model.generate(**inputs, max_length=max_length)
-        response = self.tokenizer.decode(outputs[0], skip_special_tokens=skip_special_tokens).strip()
-        return response
+        try:
+            inputs = self.tokenizer(input_text, return_tensors="pt").to(self.model.device)
+            outputs = self.model.generate(**inputs, max_length=max_length)
+            response = self.tokenizer.decode(outputs[0], skip_special_tokens=skip_special_tokens).strip()
+            log_info(f"Generated Response: {response}")
+            return response
+        except Exception as e:
+            log_error(f"ERROR Generating Response: {e}")
+            raise
 
     def stream_response(self, input_text, max_length=100, skip_special_tokens=True):
-        inputs = self.tokenizer(input_text, return_tensors="pt").to(self.model.device)
-        for output in self.model.generate(**inputs, max_length=max_length, do_stream=True):
-            response = self.tokenizer.decode(output, skip_special_tokens=skip_special_tokens).strip()
-            yield response
+        try:
+            inputs = self.tokenizer(input_text, return_tensors="pt").to(self.model.device)
+            for output in self.model.generate(**inputs, max_length=max_length, do_stream=True):
+                response = self.tokenizer.decode(output, skip_special_tokens=skip_special_tokens).strip()
+                yield response
+        except Exception as e:
+            log_error(f"ERROR Streaming Response: {e}")
+            raise
 
     def chat(self, user_input, max_length=100, skip_special_tokens=True):
-        # Add user input to chat history
-        self.chat_history.append({"role": "user", "content": user_input})
+        try:
+            # Add user input to chat history
+            self.chat_history.append({"role": "user", "content": user_input})
+            log_info(f"User Input: {user_input}")
 
-        # Generate model response
-        model_response = self.generate_response(user_input, max_length=max_length, skip_special_tokens=skip_special_tokens)
+            # Generate model response
+            model_response = self.generate_response(user_input, max_length=max_length, skip_special_tokens=skip_special_tokens)
 
-        # Add model response to chat history
-        self.chat_history.append({"role": "assistant", "content": model_response})
+            # Add model response to chat history
+            self.chat_history.append({"role": "assistant", "content": model_response})
+            log_info(f"Assistant Response: {model_response}")
 
-        # Save chat log
-        self.save_chat_log()
+            # Save chat log
+            self.save_chat_log()
 
-        return model_response
+            return model_response
+        except Exception as e:
+            log_error(f"ERROR in Chat: {e}")
+            raise
 
     def save_chat_log(self):
-        with open(self.log_file, "a", encoding="utf-8") as f:
-            for entry in self.chat_history[-2:]:  # Save only the latest interaction
-                role = entry["role"]
-                content = entry["content"]
-                f.write(f"**{role.capitalize()}:**\n\n{content}\n\n---\n\n")
+        try:
+            with open(self.log_file, "a", encoding="utf-8") as f:
+                for entry in self.chat_history[-2:]:  # Save only the latest interaction
+                    role = entry["role"]
+                    content = entry["content"]
+                    f.write(f"**{role.capitalize()}:**\n\n{content}\n\n---\n\n")
+            log_info(f"Chat log saved to {self.log_file}")
+        except Exception as e:
+            log_error(f"ERROR Saving Chat Log: {e}")
+            raise
 
     def clear_chat_history(self):
-        self.chat_history = []
-        print("Chat history cleared.")
+        try:
+            self.chat_history = []
+            log_info("Chat history cleared.")
+        except Exception as e:
+            log_error(f"ERROR Clearing Chat History: {e}")
+            raise
 
     def print_chat_history(self):
-        for entry in self.chat_history:
-            role = entry["role"]
-            content = entry["content"]
-            print(f"{role.capitalize()}: {content}\n")
+        try:
+            for entry in self.chat_history:
+                role = entry["role"]
+                content = entry["content"]
+                print(f"{role.capitalize()}: {content}\n")
+            log_info("Printed chat history to console.")
+        except Exception as e:
+            log_error(f"ERROR Printing Chat History: {e}")
+            raise