Spaces:
Running
Running
Update hf_model.py
Browse files- 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
|
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
|