From a97a545a287eb5722a6eefea8c02e7aa497a39dd Mon Sep 17 00:00:00 2001 From: BuildTools Date: Thu, 22 Aug 2024 17:08:45 -0700 Subject: [PATCH] refactor: move get helper functions to utils.py - move get_models_data and get_tasks_data to utils.py from AutoGGUF.py --- src/AutoGGUF.py | 34 ++-------------------------------- src/utils.py | 35 +++++++++++++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 32 deletions(-) diff --git a/src/AutoGGUF.py b/src/AutoGGUF.py index aa67371..b39ca92 100644 --- a/src/AutoGGUF.py +++ b/src/AutoGGUF.py @@ -80,6 +80,8 @@ def __init__(self, args): self.browse_output = utils.browse_output.__get__(self) self.browse_logs = utils.browse_logs.__get__(self) self.browse_imatrix = utils.browse_imatrix.__get__(self) + self.get_models_data = utils.get_models_data.__get__(self) + self.get_tasks_data = utils.get_tasks_data.__get__(self) self.update_threads_spinbox = partial(ui_update.update_threads_spinbox, self) self.update_threads_slider = partial(ui_update.update_threads_slider, self) self.update_gpu_offload_spinbox = partial( @@ -1549,38 +1551,6 @@ def browse_imatrix_output(self): if output_file: self.imatrix_output.setText(os.path.abspath(output_file)) - def get_models_data(self): - models = [] - root = self.model_tree.invisibleRootItem() - child_count = root.childCount() - for i in range(child_count): - item = root.child(i) - model_name = item.text(0) - model_type = "sharded" if "sharded" in model_name.lower() else "single" - model_path = item.data(0, Qt.ItemDataRole.UserRole) - models.append({"name": model_name, "type": model_type, "path": model_path}) - return models - - def get_tasks_data(self): - tasks = [] - for i in range(self.task_list.count()): - item = self.task_list.item(i) - task_widget = self.task_list.itemWidget(item) - if task_widget: - tasks.append( - { - "name": task_widget.task_name, - "status": task_widget.status, - "progress": ( - task_widget.progress_bar.value() - if hasattr(task_widget, "progress_bar") - else 0 - ), - "log_file": task_widget.log_file, - } - ) - return tasks - def generate_imatrix(self): self.logger.info(STARTING_IMATRIX_GENERATION) try: diff --git a/src/utils.py b/src/utils.py index bca9a6d..ae36822 100644 --- a/src/utils.py +++ b/src/utils.py @@ -1,3 +1,4 @@ +from PySide6.QtCore import Qt from PySide6.QtWidgets import QFileDialog from error_handling import show_error @@ -8,6 +9,40 @@ from imports_and_globals import ensure_directory +def get_models_data(self): + models = [] + root = self.model_tree.invisibleRootItem() + child_count = root.childCount() + for i in range(child_count): + item = root.child(i) + model_name = item.text(0) + model_type = "sharded" if "sharded" in model_name.lower() else "single" + model_path = item.data(0, Qt.ItemDataRole.UserRole) + models.append({"name": model_name, "type": model_type, "path": model_path}) + return models + + +def get_tasks_data(self): + tasks = [] + for i in range(self.task_list.count()): + item = self.task_list.item(i) + task_widget = self.task_list.itemWidget(item) + if task_widget: + tasks.append( + { + "name": task_widget.task_name, + "status": task_widget.status, + "progress": ( + task_widget.progress_bar.value() + if hasattr(task_widget, "progress_bar") + else 0 + ), + "log_file": task_widget.log_file, + } + ) + return tasks + + def browse_models(self): self.logger.info(BROWSING_FOR_MODELS_DIRECTORY) models_path = QFileDialog.getExistingDirectory(self, SELECT_MODELS_DIRECTORY)