mew77 commited on
Commit
5d07573
·
verified ·
1 Parent(s): a71142d

Update hf_model.py

Browse files
Files changed (1) hide show
  1. hf_model.py +8 -8
hf_model.py CHANGED
@@ -2,7 +2,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
2
  import torch
3
  from datetime import datetime
4
  import os
5
- from .mew_log import log_info, log_error # Import your custom logging methods
6
 
7
  class HFModel:
8
  def __init__(self, model_name):
@@ -17,7 +17,7 @@ class HFModel:
17
  self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
18
  log_info(f"=== Model Loaded Successfully: {self.friendly_name} ===")
19
  except Exception as e:
20
- log_error(f"ERROR Loading Model: {e}")
21
  raise
22
 
23
  def generate_response(self, input_text, max_length=100, skip_special_tokens=True):
@@ -28,7 +28,7 @@ class HFModel:
28
  log_info(f"Generated Response: {response}")
29
  return response
30
  except Exception as e:
31
- log_error(f"ERROR Generating Response: {e}")
32
  raise
33
 
34
  def stream_response(self, input_text, max_length=100, skip_special_tokens=True):
@@ -38,7 +38,7 @@ class HFModel:
38
  response = self.tokenizer.decode(output, skip_special_tokens=skip_special_tokens).strip()
39
  yield response
40
  except Exception as e:
41
- log_error(f"ERROR Streaming Response: {e}")
42
  raise
43
 
44
  def chat(self, user_input, max_length=100, skip_special_tokens=True):
@@ -59,7 +59,7 @@ class HFModel:
59
 
60
  return model_response
61
  except Exception as e:
62
- log_error(f"ERROR in Chat: {e}")
63
  raise
64
 
65
  def save_chat_log(self):
@@ -71,7 +71,7 @@ class HFModel:
71
  f.write(f"**{role.capitalize()}:**\n\n{content}\n\n---\n\n")
72
  log_info(f"Chat log saved to {self.log_file}")
73
  except Exception as e:
74
- log_error(f"ERROR Saving Chat Log: {e}")
75
  raise
76
 
77
  def clear_chat_history(self):
@@ -79,7 +79,7 @@ class HFModel:
79
  self.chat_history = []
80
  log_info("Chat history cleared.")
81
  except Exception as e:
82
- log_error(f"ERROR Clearing Chat History: {e}")
83
  raise
84
 
85
  def print_chat_history(self):
@@ -90,5 +90,5 @@ class HFModel:
90
  print(f"{role.capitalize()}: {content}\n")
91
  log_info("Printed chat history to console.")
92
  except Exception as e:
93
- log_error(f"ERROR Printing Chat History: {e}")
94
  raise
 
2
  import torch
3
  from datetime import datetime
4
  import os
5
+ from mew_log import log_info, log_error # Import your custom logging methods
6
 
7
  class HFModel:
8
  def __init__(self, model_name):
 
17
  self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
18
  log_info(f"=== Model Loaded Successfully: {self.friendly_name} ===")
19
  except Exception as e:
20
+ log_error(f"ERROR Loading Model: {e}", e)
21
  raise
22
 
23
  def generate_response(self, input_text, max_length=100, skip_special_tokens=True):
 
28
  log_info(f"Generated Response: {response}")
29
  return response
30
  except Exception as e:
31
+ log_error(f"ERROR Generating Response: {e}", e)
32
  raise
33
 
34
  def stream_response(self, input_text, max_length=100, skip_special_tokens=True):
 
38
  response = self.tokenizer.decode(output, skip_special_tokens=skip_special_tokens).strip()
39
  yield response
40
  except Exception as e:
41
+ log_error(f"ERROR Streaming Response: {e}", e)
42
  raise
43
 
44
  def chat(self, user_input, max_length=100, skip_special_tokens=True):
 
59
 
60
  return model_response
61
  except Exception as e:
62
+ log_error(f"ERROR in Chat: {e}", e)
63
  raise
64
 
65
  def save_chat_log(self):
 
71
  f.write(f"**{role.capitalize()}:**\n\n{content}\n\n---\n\n")
72
  log_info(f"Chat log saved to {self.log_file}")
73
  except Exception as e:
74
+ log_error(f"ERROR Saving Chat Log: {e}", e)
75
  raise
76
 
77
  def clear_chat_history(self):
 
79
  self.chat_history = []
80
  log_info("Chat history cleared.")
81
  except Exception as e:
82
+ log_error(f"ERROR Clearing Chat History: {e}", e)
83
  raise
84
 
85
  def print_chat_history(self):
 
90
  print(f"{role.capitalize()}: {content}\n")
91
  log_info("Printed chat history to console.")
92
  except Exception as e:
93
+ log_error(f"ERROR Printing Chat History: {e}", e)
94
  raise