mirror of https://github.com/leafspark/AutoGGUF
80 lines
2.9 KiB
Python
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)
|
|
|