From cee4294ecfb9594e384f38f9c7d34919fe8bbab9 Mon Sep 17 00:00:00 2001
From: BuildTools <chenxuzhang45@gmail.com>
Date: Mon, 9 Sep 2024 19:26:56 -0700
Subject: [PATCH] feat(parallel): add support for iMatrix generation tracking

- add support for iMatrix generation tracking
- don't adjust progress bar when indeterminate
---
 src/AutoGGUF.py           |  9 +++++++--
 src/QuantizationThread.py | 19 ++++++++++++++++++-
 src/TaskListItem.py       |  3 +--
 src/main.py               |  2 +-
 4 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/src/AutoGGUF.py b/src/AutoGGUF.py
index 202f349..e2ad663 100644
--- a/src/AutoGGUF.py
+++ b/src/AutoGGUF.py
@@ -1,6 +1,5 @@
 import importlib
 import json
-import re
 import shutil
 import urllib.request
 import urllib.error
@@ -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()))
@@ -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,
             )
@@ -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)
             )
diff --git a/src/QuantizationThread.py b/src/QuantizationThread.py
index d196936..1c07638 100644
--- a/src/QuantizationThread.py
+++ b/src/QuantizationThread.py
@@ -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
diff --git a/src/TaskListItem.py b/src/TaskListItem.py
index dce73fc..293a7cd 100644
--- a/src/TaskListItem.py
+++ b/src/TaskListItem.py
@@ -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))
diff --git a/src/main.py b/src/main.py
index 2b8ff1b..5ff4c92 100644
--- a/src/main.py
+++ b/src/main.py
@@ -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")