AutoGGUF/QuantizationThread.py

80 lines
2.9 KiB
Python

from PyQt6.QtWidgets import *
from PyQt6.QtCore import *
from PyQt6.QtGui import *
import os
import sys
import psutil
import subprocess
import time
import signal
import json
import platform
import requests
import zipfile
import traceback
from datetime import datetime
from imports_and_globals import open_file_safe
class QuantizationThread(QThread):
output_signal = pyqtSignal(str)
status_signal = pyqtSignal(str)
finished_signal = pyqtSignal()
error_signal = pyqtSignal(str)
model_info_signal = pyqtSignal(dict)
def __init__(self, command, cwd, log_file):
super().__init__()
self.command = command
self.cwd = cwd
self.log_file = log_file
self.process = None
self.model_info = {}
def run(self):
try:
self.process = subprocess.Popen(self.command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
text=True, cwd=self.cwd)
with open_file_safe(self.log_file, 'w') as log:
for line in self.process.stdout:
line = line.strip()
self.output_signal.emit(line)
log.write(line + '\n')
log.flush()
self.status_signal.emit("In Progress")
self.parse_model_info(line)
self.process.wait()
if self.process.returncode == 0:
self.status_signal.emit("Completed")
self.model_info_signal.emit(self.model_info)
else:
self.error_signal.emit(f"Process exited with code {self.process.returncode}")
self.finished_signal.emit()
except Exception as e:
self.error_signal.emit(str(e))
def parse_model_info(self, line):
if "llama_model_loader: loaded meta data with" in line:
parts = line.split()
self.model_info['kv_pairs'] = parts[6]
self.model_info['tensors'] = parts[9]
elif "general.architecture" in line:
self.model_info['architecture'] = line.split('=')[-1].strip()
elif line.startswith("llama_model_loader: - kv"):
key = line.split(':')[2].strip()
value = line.split('=')[-1].strip()
self.model_info.setdefault('kv_data', {})[key] = value
elif line.startswith("llama_model_loader: - type"):
parts = line.split(':')
if len(parts) > 1:
quant_type = parts[1].strip()
tensors = parts[2].strip().split()[0]
self.model_info.setdefault('quantization_type', []).append(f"{quant_type}: {tensors} tensors")
def terminate(self):
if self.process:
os.kill(self.process.pid, signal.SIGTERM)
self.process.wait(timeout=5)
if self.process.poll() is None:
os.kill(self.process.pid, signal.SIGKILL)