Skip to content

Commit

Permalink
More Logger fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
@w4ffl35
w4ffl35 committed Jan 16, 2024
1 parent d57ed5f commit a455da5
Show file tree
Hide file tree
Showing 9 changed files with 78 additions and 75 deletions.
3 changes: 2 additions & 1 deletion src/airunner/aihandler/download_civitai.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@


class DownloadCivitAI:
logger = Logger(prefix="DownloadCivitAI")
cancel_download = False

@staticmethod
Expand All @@ -18,7 +19,7 @@ def get_json(model_id):
try:
json = response.json()
except JSONDecodeError:
Logger.error(f"Failed to decode JSON from {url}")
self.logger.error(f"Failed to decode JSON from {url}")
print(response)
return json

Expand Down
51 changes: 25 additions & 26 deletions src/airunner/aihandler/llm.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,8 @@
from airunner.aihandler.logger import Logger




class LLM(QObject):
logger = Logger(prefix="LLM")
dtype = ""
local_files_only = True
set_attention_mask = False
Expand Down Expand Up @@ -89,7 +88,7 @@ def __init__(self, *args, **kwargs):

def move_to_cpu(self):
if self.model:
Logger.info("Moving model to CPU")
self.logger.info("Moving model to CPU")
self.model.to("cpu")
self.tokenizer = None

Expand All @@ -102,16 +101,16 @@ def move_to_device(self, device=None):
device_name = "cuda"
else:
device_name = "cpu"
Logger.info("Moving model to device {device_name}")
self.logger.info("Moving model to device {device_name}")
self.model.to(device_name)

def load_tokenizer(self, local_files_only = None):
Logger.info(f"Loading tokenizer for {self.requested_generator_name}")
self.logger.info(f"Loading tokenizer for {self.requested_generator_name}")
if self.requested_generator_name == "casuallm":
local_files_only = self.local_files_only if local_files_only is None else local_files_only
if not self.tokenizer is None:
return
Logger.info("Load tokenizer")
self.logger.info("Load tokenizer")
try:
self.tokenizer = AutoTokenizer.from_pretrained(
self.current_model_path,
Expand All @@ -122,19 +121,19 @@ def load_tokenizer(self, local_files_only = None):
except OSError as e:
if "Checkout your internet connection" in str(e):
if local_files_only:
Logger.info("Unable to load tokenizer, model does not exist locally, trying to load from remote")
self.logger.info("Unable to load tokenizer, model does not exist locally, trying to load from remote")
return self.load_tokenizer(local_files_only=False)
else:
Logger.error(e)
self.logger.error(e)
if self.tokenizer:
self.tokenizer.use_default_system_prompt = False

def load_streamer(self):
Logger.info("Loading LLM text streamer")
self.logger.info("Loading LLM text streamer")
self.streamer = TextIteratorStreamer(self.tokenizer)

def load_model(self, local_files_only = None):
Logger.info("Loading model")
self.logger.info("Loading model")
if not self.do_load_model:
return

Expand Down Expand Up @@ -166,20 +165,20 @@ def load_model(self, local_files_only = None):
auto_class_ = InstructBlipForConditionalGeneration
params.pop("use_cache", None)

Logger.info(f"Loading model from {path}")
self.logger.info(f"Loading model from {path}")
try:
self.model = auto_class_.from_pretrained(path, **params)
except OSError as e:
if "Checkout your internet connection" in str(e):
if local_files_only:
return self.load_model(local_files_only=False)
else:
Logger.error(e)
self.logger.error(e)

# print the type of class that self.model is

def load_processor(self, local_files_only = None):
Logger.info(f"Loading processor {self.model_path}")
self.logger.info(f"Loading processor {self.model_path}")
local_files_only = self.local_files_only if local_files_only is None else local_files_only
kwargs = {}
config = self.quantization_config()
Expand All @@ -192,15 +191,15 @@ def load_processor(self, local_files_only = None):
**kwargs
)
if self.processor:
Logger.info("Processor loaded")
self.logger.info("Processor loaded")
else:
Logger.error("Failed to load processor")
self.logger.error("Failed to load processor")

def clear_history(self):
self.history = []

def unload_tokenizer(self):
Logger.info("Unloading tokenizer")
self.logger.info("Unloading tokenizer")
self.tokenizer = None
self.engine.clear_memory()

Expand All @@ -210,14 +209,14 @@ def unload_model(self):
self.engine.clear_memory()

def unload_processor(self):
Logger.info("Unloading processor")
self.logger.info("Unloading processor")
self.processor = None
self.engine.clear_memory()

def quantization_config(self):
config = None
if self.dtype == "8bit":
Logger.info("Loading 8bit model")
self.logger.info("Loading 8bit model")
config = BitsAndBytesConfig(
load_in_4bit=False,
load_in_8bit=True,
Expand All @@ -228,7 +227,7 @@ def quantization_config(self):
bnb_4bit_quant_type='nf4',
)
elif self.dtype == "4bit":
Logger.info("Loading 4bit model")
self.logger.info("Loading 4bit model")
config = BitsAndBytesConfig(
load_in_4bit=True,
load_in_8bit=False,
Expand All @@ -239,7 +238,7 @@ def quantization_config(self):
bnb_4bit_quant_type='nf4',
)
elif self.dtype == "2bit":
Logger.info("Loading 2bit model")
self.logger.info("Loading 2bit model")
config = GPTQConfig(bits=2, dataset = "c4", tokenizer=self.tokenizer)
return config

Expand Down Expand Up @@ -272,7 +271,7 @@ def handle_request(self):
}.get(self.request_type, default_)()

def prepare_input_args(self):
Logger.info("Preparing input args")
self.logger.info("Preparing input args")
self.system_instructions = self.request_data.get("system_instructions", "")
top_k = self.parameters.get("top_k", self.top_k)
eta_cutoff = self.parameters.get("eta_cutoff", self.eta_cutoff)
Expand Down Expand Up @@ -344,7 +343,7 @@ def do_set_seed(self, seed=None):
self.tokenizer.seed = self.seed

def handle_generate_request(self):
Logger.info("Handling generate request")
self.logger.info("Handling generate request")
self._processing_request = True
kwargs = self.prepare_input_args()
self.do_set_seed(kwargs.get("seed"))
Expand All @@ -361,14 +360,14 @@ def clear_conversation(self):
self.chain.clear()

def do_generate(self, data):
Logger.info("Generating with LLM")
self.logger.info("Generating with LLM")
self.process_data(data)
self.handle_request()
self.requested_generator_name = data["request_data"]["generator_name"]
return self.generate()

def generate(self):
Logger.info("Generating with LLM " + self.requested_generator_name)
self.logger.info("Generating with LLM " + self.requested_generator_name)
# Create a FileSystemLoader object with the directory of the template
HERE = os.path.dirname(os.path.abspath(__file__))
file_loader = FileSystemLoader(os.path.join(HERE, "chat_templates"))
Expand Down Expand Up @@ -428,7 +427,7 @@ def generate(self):
model_inputs = encoded.to("cuda" if torch.cuda.is_available() else "cpu")

# Generate the response
Logger.info("Generating...")
self.logger.info("Generating...")
import threading
thread = threading.Thread(target=self.model.generate, kwargs=dict(
model_inputs,
Expand Down Expand Up @@ -507,4 +506,4 @@ def generate(self):
answers.append(answer.strip().lower())
return answers
else:
Logger.error(f"Failed to call generator for {self.requested_generator_name}")
self.logger.error(f"Failed to call generator for {self.requested_generator_name}")
40 changes: 20 additions & 20 deletions src/airunner/aihandler/transformer_runner.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,12 +7,12 @@
from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer
from transformers import InstructBlipForConditionalGeneration
from transformers import InstructBlipProcessor
from airunner.aihandler.enums import MessageCode
from airunner.aihandler.logger import Logger
from transformers import GPTQConfig


class TransformerRunner(QObject):
logger = Logger(prefix="TransformerRunner")
dtype = ""
local_files_only = True
set_attention_mask = False
Expand Down Expand Up @@ -83,7 +83,7 @@ def __init__(self, *args, **kwargs):

def move_to_cpu(self):
if self.model:
Logger.info("Moving model to CPU")
self.logger.info("Moving model to CPU")
self.model.to("cpu")
self.tokenizer = None

Expand All @@ -96,11 +96,11 @@ def move_to_device(self, device=None):
device_name = "cuda"
else:
device_name = "cpu"
Logger.info("Moving model to device {device_name}")
self.logger.info("Moving model to device {device_name}")
self.model.to(device_name)

def unload_tokenizer(self):
Logger.info("Unloading tokenizer")
self.logger.info("Unloading tokenizer")
self.tokenizer = None
self.engine.clear_memory()

Expand All @@ -112,7 +112,7 @@ def unload_model(self):
def quantization_config(self):
config = None
if self.dtype == "8bit":
Logger.info("Loading 8bit model")
self.logger.info("Loading 8bit model")
config = BitsAndBytesConfig(
load_in_4bit=False,
load_in_8bit=True,
Expand All @@ -123,7 +123,7 @@ def quantization_config(self):
bnb_4bit_quant_type='nf4',
)
elif self.dtype == "4bit":
Logger.info("Loading 4bit model")
self.logger.info("Loading 4bit model")
config = BitsAndBytesConfig(
load_in_4bit=True,
load_in_8bit=False,
Expand All @@ -134,17 +134,17 @@ def quantization_config(self):
bnb_4bit_quant_type='nf4',
)
elif self.dtype == "2bit":
Logger.info("Loading 2bit model")
self.logger.info("Loading 2bit model")
config = GPTQConfig(bits=2, dataset = "c4", tokenizer=self.tokenizer)
return config

def unload_processor(self):
Logger.info("Unloading processor")
self.logger.info("Unloading processor")
self.processor = None
self.engine.clear_memory()

def load_processor(self, local_files_only = None):
Logger.info(f"Loading processor {self.model_path}")
self.logger.info(f"Loading processor {self.model_path}")
local_files_only = self.local_files_only if local_files_only is None else local_files_only
kwargs = {}
config = self.quantization_config()
Expand All @@ -157,12 +157,12 @@ def load_processor(self, local_files_only = None):
**kwargs
)
if self.processor:
Logger.info("Processor loaded")
self.logger.info("Processor loaded")
else:
Logger.error("Failed to load processor")
self.logger.error("Failed to load processor")

def load_model(self, local_files_only = None):
Logger.info("Loading model")
self.logger.info("Loading model")
if not self.do_load_model:
return

Expand Down Expand Up @@ -194,15 +194,15 @@ def load_model(self, local_files_only = None):
auto_class_ = InstructBlipForConditionalGeneration
params.pop("use_cache", None)

Logger.info(f"Loading model from {path}")
self.logger.info(f"Loading model from {path}")
try:
self.model = auto_class_.from_pretrained(path, **params)
except OSError as e:
if "Checkout your internet connection" in str(e):
if local_files_only:
return self.load_model(local_files_only=False)
else:
Logger.error(e)
self.logger.error(e)


def process_data(self, data):
Expand Down Expand Up @@ -241,7 +241,7 @@ def clear_conversation(self):
pass

def prepare_input_args(self):
Logger.info("Preparing input args")
self.logger.info("Preparing input args")
self.system_instructions = self.request_data.get("system_instructions", "")
top_k = self.parameters.get("top_k", self.top_k)
eta_cutoff = self.parameters.get("eta_cutoff", self.eta_cutoff)
Expand Down Expand Up @@ -299,12 +299,12 @@ def prepare_input_args(self):
return kwargs

def load_tokenizer(self, local_files_only = None):
Logger.info(f"Loading tokenizer for {self.requested_generator_name}")
self.logger.info(f"Loading tokenizer for {self.requested_generator_name}")
if self.requested_generator_name == "casuallm":
local_files_only = self.local_files_only if local_files_only is None else local_files_only
if not self.tokenizer is None:
return
Logger.info("Load tokenizer")
self.logger.info("Load tokenizer")
try:
self.tokenizer = AutoTokenizer.from_pretrained(
self.current_model_path,
Expand All @@ -315,10 +315,10 @@ def load_tokenizer(self, local_files_only = None):
except OSError as e:
if "Checkout your internet connection" in str(e):
if local_files_only:
Logger.info("Unable to load tokenizer, model does not exist locally, trying to load from remote")
self.logger.info("Unable to load tokenizer, model does not exist locally, trying to load from remote")
return self.load_tokenizer(local_files_only=False)
else:
Logger.error(e)
self.logger.error(e)
if self.tokenizer:
self.tokenizer.use_default_system_prompt = False

Expand All @@ -337,7 +337,7 @@ def do_set_seed(self, seed=None):
self.tokenizer.seed = self.seed

def handle_generate_request(self):
Logger.info("Handling generate request")
self.logger.info("Handling generate request")
self.disable_request_processing()
kwargs = self.prepare_input_args()
self.do_set_seed(kwargs.get("seed"))
Expand Down
Loading

0 comments on commit a455da5

Please sign in to comment.