Skip to content

Commit

Permalink
feat(parallel): add support for iMatrix generation tracking
Browse files Browse the repository at this point in the history
- add support for iMatrix generation tracking
- don't adjust progress bar when indeterminate
  • Loading branch information
leafspark committed Sep 10, 2024
1 parent c96380c commit cee4294
Show file tree
Hide file tree
Showing 4 changed files with 27 additions and 6 deletions.
9 changes: 7 additions & 2 deletions src/AutoGGUF.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
import importlib
import json
import re
import shutil
import urllib.request
import urllib.error
Expand Down Expand Up @@ -1547,6 +1546,7 @@ def task_finished(self, thread, task_item) -> None:
if thread in self.quant_threads:
self.quant_threads.remove(thread)
task_item.update_status(COMPLETED)
self.setAttribute(Qt.WA_WindowModified, True) # Set modified flag

def show_task_details(self, item) -> None:
self.logger.debug(SHOWING_TASK_DETAILS_FOR.format(item.text()))
Expand Down Expand Up @@ -1649,7 +1649,7 @@ def generate_imatrix(self) -> None:
task_item = TaskListItem(
task_name,
log_file,
show_progress_bar=False,
show_progress_bar=True,
logger=self.logger,
quant_threads=self.quant_threads,
)
Expand All @@ -1658,7 +1658,12 @@ def generate_imatrix(self) -> None:
self.task_list.addItem(list_item)
self.task_list.setItemWidget(list_item, task_item)

imatrix_chunks = None

thread.status_signal.connect(task_item.update_status)
thread.output_signal.connect(
lambda line, ti=task_item: self.parse_progress(line, ti, imatrix_chunks)
)
thread.finished_signal.connect(
lambda: self.task_finished(thread, task_item)
)
Expand Down
19 changes: 18 additions & 1 deletion src/QuantizationThread.py
Original file line number Diff line number Diff line change
Expand Up @@ -79,14 +79,31 @@ def parse_model_info(self, line) -> None:
f"{quant_type}: {tensors} tensors"
)

def parse_progress(self, line, task_item) -> None:
def parse_progress(self, line, task_item, imatrix_chunks=None) -> None:
# Parses the output line for progress information and updates the task item.
match = re.search(r"\[\s*(\d+)\s*/\s*(\d+)\s*].*", line)

if match:
current = int(match.group(1))
total = int(match.group(2))
progress = int((current / total) * 100)
task_item.update_progress(progress)
else:
imatrix_match = re.search(
r"compute_imatrix: computing over (\d+) chunks with batch_size \d+",
line,
)
if imatrix_match:
imatrix_chunks = int(imatrix_match.group(1))
elif imatrix_chunks is not None:
save_match = re.search(
r"save_imatrix: stored collected data after (\d+) chunks in .*",
line,
)
if save_match:
saved_chunks = int(save_match.group(1))
progress = int((saved_chunks / self.imatrix_chunks) * 100)
task_item.update_progress(progress)

def terminate(self) -> None:
# Terminate the subprocess if it's still running
Expand Down
3 changes: 1 addition & 2 deletions src/TaskListItem.py
Original file line number Diff line number Diff line change
Expand Up @@ -162,8 +162,7 @@ def update_progress(self, value=None) -> None:
self.progress_value = value
self.progress_bar.setValue(self.progress_value)
else:
# Set progress bar to zero for indeterminate progress
self.progress_bar.setValue(0)
return

def restart_task(self, task_item) -> None:
self.logger.info(RESTARTING_TASK.format(task_item.task_name))
Expand Down
2 changes: 1 addition & 1 deletion src/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ class Config:


class Task(BaseModel):
id: str = Field(..., description="Unique identifier for the task")
# id: str = Field(..., description="Unique identifier for the task")
status: str = Field(..., description="Current status of the task")
progress: float = Field(..., description="Progress of the task as a percentage")

Expand Down

0 comments on commit cee4294

Please sign in to comment.