pysfi 0.1.7__py3-none-any.whl → 0.1.11__py3-none-any.whl

sfi/llmquantize/llmquantize.py

@@ -0,0 +1,480 @@
+"""GGUF量化转换GUI工具, 用于将F16格式的GGUF文件转换为其他主流量化格式."""
+
+from __future__ import annotations
+
+import atexit
+import json
+import logging
+import pathlib
+import subprocess
+import sys
+from pathlib import Path
+from types import SimpleNamespace
+
+from PySide2.QtCore import QThread, Signal, Slot
+from PySide2.QtGui import QMoveEvent, QResizeEvent
+from PySide2.QtWidgets import (
+    QApplication,
+    QCheckBox,
+    QFileDialog,
+    QGridLayout,
+    QGroupBox,
+    QHBoxLayout,
+    QLabel,
+    QMainWindow,
+    QProgressBar,
+    QPushButton,
+    QTextEdit,
+    QVBoxLayout,
+    QWidget,
+)
+
+CONFIG_FILE = Path.home() / ".sfi" / "llmquantize.json"
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class QuantizerConfig(SimpleNamespace):
+    """GGUF量化转换工具配置."""
+
+    TITLE: str = "GGUF量化转换工具"
+    WIN_SIZE: list[int] = [600, 500]  # noqa: RUF012
+    WIN_POS: list[int] = [100, 100]  # noqa: RUF012
+    LAST_INPUT_FILE: str = ""
+    SELECTED_QUANTS: list[str] = ["Q4_K_M", "Q5_K_M"]  # noqa: RUF012
+
+    def __init__(self) -> None:
+        if CONFIG_FILE.exists():
+            logger.info("Loading configuration from %s", CONFIG_FILE)
+            try:
+                # 直接更新，忽略无效字段
+                self.__dict__.update(json.loads(CONFIG_FILE.read_text()))
+            except (json.JSONDecodeError, TypeError) as e:
+                logger.warning("Failed to load configuration: %s", e)
+                logger.info("Using default configuration")
+        else:
+            logger.info("Using default configuration")
+
+    def save(self) -> None:
+        """保存配置."""
+        CONFIG_FILE.parent.mkdir(parents=True, exist_ok=True)
+        CONFIG_FILE.write_text(json.dumps(vars(self), indent=4))
+
+
+conf = QuantizerConfig()
+atexit.register(conf.save)
+
+
+def _process_gguf_stem(filename: str) -> str:
+    """处理文件名, 移除可能的F16后缀.
+
+    Returns:
+        str: 处理后的文件名
+    """
+    if filename.upper().endswith("-F16"):
+        filename = filename[:-4]  # 移除-F16后缀
+    return filename
+
+
+class QuantizationWorker(QThread):
+    """量化执行线程Worker."""
+
+    progress_msg_updated = Signal(str)
+    progress_count_updated = Signal(int)
+    is_finished = Signal()
+
+    def __init__(
+        self,
+        input_file: pathlib.Path,
+        quant_types: list[str],
+    ) -> None:
+        super().__init__()
+
+        self.input_file = input_file
+        self.quant_types = quant_types
+        self.input_dir = input_file.parent
+        self.base_name = _process_gguf_stem(input_file.stem)
+        self.total_files = len(quant_types)
+        self.completed_files = 0
+        self.success = True  # 记录整体转换状态
+
+    def run(self) -> None:
+        """执行量化转换任务."""
+        try:
+            for quant_type in self.quant_types:
+                output_file: pathlib.Path = (
+                    self.input_dir / f"{self.base_name}-{quant_type}.gguf"
+                )
+
+                self.progress_msg_updated.emit(
+                    f"正在转换到 {quant_type} 格式...",
+                )
+
+                # 构建命令行参数（确保所有参数都是字符串）
+                cmd = [
+                    "llama-quantize",
+                    str(self.input_file.name),
+                    str(output_file.name),
+                    quant_type,
+                ]
+
+                # 执行转换命令（使用 cwd 参数避免全局目录变更）
+                try:
+                    process = subprocess.Popen(
+                        cmd,
+                        stdout=subprocess.PIPE,
+                        stderr=subprocess.STDOUT,
+                        universal_newlines=True,
+                        bufsize=1,  # 行缓冲以获得实时输出
+                        cwd=str(self.input_file.parent),
+                    )
+                except (FileNotFoundError, PermissionError) as e:
+                    self.progress_msg_updated.emit(f"无法启动量化进程: {e!s}")
+                    self.success = False
+                    continue
+
+                # 实时输出进度
+                if not process.stdout:
+                    logger.error("无法获取进度信息")
+                    process.wait()
+                    self.success = False
+                    continue
+
+                try:
+                    for line in process.stdout:
+                        self.progress_msg_updated.emit(line.strip())
+                finally:
+                    # 等待进程结束，设置超时防止永久挂起
+                    try:
+                        process.wait(timeout=3600)  # 单个文件最多1小时
+                    except subprocess.TimeoutExpired:
+                        self.progress_msg_updated.emit(f"转换 {quant_type} 超时")
+                        process.kill()
+                        process.wait()
+                        self.success = False
+                        continue
+                    finally:
+                        # 确保关闭管道以避免资源泄漏
+                        if process.stdout:
+                            process.stdout.close()
+
+                # 只有在成功时才更新进度
+                if process.returncode == 0:
+                    self.completed_files += 1
+                    progress = int((self.completed_files / self.total_files) * 100)
+                    self.progress_count_updated.emit(progress)
+                    self.progress_msg_updated.emit(
+                        f"成功生成: {output_file!s}",
+                    )
+                else:
+                    self.progress_msg_updated.emit(f"转换 {quant_type} 失败")
+                    self.success = False
+
+            self.is_finished.emit()
+        except Exception as e:
+            logger.exception("量化转换过程中发生异常")
+            self.progress_msg_updated.emit(f"发生错误: {e!s}")
+            self.success = False
+            self.is_finished.emit()
+
+
+class GGUFQuantizerGUI(QMainWindow):
+    """GGUF量化转换工具GUI界面."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.setWindowTitle(conf.TITLE)
+        self.setGeometry(*conf.WIN_POS, *conf.WIN_SIZE)
+
+        self.input_file: pathlib.Path = pathlib.Path(conf.LAST_INPUT_FILE)
+        self.worker: QuantizationWorker | None = None
+        self.quant_types = {
+            "Q2_K": "Q2_K (极低精度, 最小尺寸)",
+            "Q3_K_S": "Q3_K_S (低精度, 小尺寸)",
+            "Q3_K_M": "Q3_K_M (低精度, 中等尺寸)",
+            "Q3_K_L": "Q3_K_L (低精度, 大尺寸)",
+            "Q4_0": "Q4_0 (基本4位)",
+            "Q4_K_S": "Q4_K_S (4位, 小尺寸)",
+            "Q4_K_M": "Q4_K_M (4位, 中等尺寸)",
+            "Q5_0": "Q5_0 (基本5位)",
+            "Q5_K_S": "Q5_K_S (5位, 小尺寸)",
+            "Q5_K_M": "Q5_K_M (5位, 中等尺寸)",
+            "Q6_K": "Q6_K (6位, 高质量)",
+            "Q8_0": "Q8_0 (8位, 最高质量)",
+        }
+        self.quant_checks = {}  # 存储量化类型对应的checkbox
+
+        self.init_ui()
+
+        # 恢复上次选择的文件
+        if conf.LAST_INPUT_FILE and pathlib.Path(conf.LAST_INPUT_FILE).exists():
+            self.input_file = pathlib.Path(conf.LAST_INPUT_FILE)
+            self.file_label.setText(self.input_file.name)
+            self.check_existing_quant_files()
+            self.convert_btn.setEnabled(True)
+
+    def init_ui(self) -> None:
+        """初始化界面."""
+        main_widget = QWidget()
+        main_layout = QVBoxLayout()
+
+        # 文件选择部分
+        file_group = QGroupBox("选择F16格式的GGUF文件")
+        file_layout = QVBoxLayout()
+
+        self.file_label = QLabel("未选择文件")
+        file_btn = QPushButton("选择文件")
+        file_btn.clicked.connect(self.select_file)
+
+        file_layout.addWidget(self.file_label)
+        file_layout.addWidget(file_btn)
+        file_group.setLayout(file_layout)
+
+        # 量化选项部分
+        quant_group = QGroupBox("选择量化类型")
+        quant_layout = QGridLayout()
+
+        for i, (quant_type, label) in enumerate(self.quant_types.items()):
+            check = QCheckBox(label)
+            # 连接信号，选择时保存配置
+            check.stateChanged.connect(self.on_quant_type_changed)
+            self.quant_checks[quant_type] = check
+            row = i // 2
+            col = i % 2
+            quant_layout.addWidget(check, row, col)
+
+        # 恢复上次选择的量化类型
+        for quant_type in conf.SELECTED_QUANTS:
+            if quant_type in self.quant_checks:
+                self.quant_checks[quant_type].setChecked(True)
+
+        quant_group.setLayout(quant_layout)
+
+        # 进度显示部分
+        progress_group = QGroupBox("转换进度")
+        progress_layout = QVBoxLayout()
+
+        self.progress_bar = QProgressBar()
+        self.progress_bar.setRange(0, 100)
+        self.progress_bar.setValue(0)
+
+        self.output_text = QTextEdit()
+        self.output_text.setReadOnly(True)
+
+        progress_layout.addWidget(self.progress_bar)
+        progress_layout.addWidget(self.output_text)
+        progress_group.setLayout(progress_layout)
+
+        # 操作按钮
+        btn_layout = QHBoxLayout()
+        self.convert_btn = QPushButton("开始转换")
+        self.convert_btn.clicked.connect(self.start_conversion)
+        self.convert_btn.setEnabled(False)
+
+        btn_layout.addWidget(self.convert_btn)
+
+        # 组装主界面
+        main_layout.addWidget(file_group)
+        main_layout.addWidget(quant_group)
+        main_layout.addWidget(progress_group)
+        main_layout.addLayout(btn_layout)
+
+        main_widget.setLayout(main_layout)
+        self.setCentralWidget(main_widget)
+
+    def select_file(self) -> None:
+        """选择文件."""
+        # 使用上次选择的目录作为初始目录
+        initial_dir = ""
+        if conf.LAST_INPUT_FILE and pathlib.Path(conf.LAST_INPUT_FILE).exists():
+            initial_dir = str(pathlib.Path(conf.LAST_INPUT_FILE).parent)
+        elif self.input_file.exists():
+            initial_dir = str(self.input_file.parent)
+
+        file_path, _ = QFileDialog.getOpenFileName(
+            self,
+            "选择F16格式的GGUF文件",
+            initial_dir,
+            "GGUF Files (*.gguf)",
+        )
+
+        if file_path:
+            self.input_file = pathlib.Path(file_path)
+            filename = self.input_file.name
+            self.file_label.setText(filename)
+
+            # 保存最后选择的文件路径
+            conf.LAST_INPUT_FILE = file_path
+
+            # 检查文件名是否包含F16
+            if "-F16" not in filename.upper():
+                self.output_text.append(
+                    "注意: 输入文件名不包含F16后缀,输出文件名将直接添加量化类型",
+                )
+                self._scroll_to_bottom()
+
+            # 检查已存在的量化文件
+            self.check_existing_quant_files()
+
+            self.convert_btn.setEnabled(True)
+            self.output_text.clear()
+            self.progress_bar.setValue(0)
+            self._scroll_to_bottom()
+
+    def check_existing_quant_files(self) -> None:
+        """检查当前目录下已存在的量化文件, 并更新checkbox状态."""
+        if not self.input_file:
+            return
+
+        dir_path = self.input_file.parent
+
+        for quant_type in self.quant_types:
+            filename = f"{_process_gguf_stem(self.input_file.stem)}-{quant_type}.gguf"
+            expected_file = dir_path / filename
+            if expected_file.exists():
+                # 文件已存在，标记但不禁用，允许用户选择重新生成
+                self.quant_checks[quant_type].setText(
+                    f"{self.quant_types[quant_type]} (已存在)",
+                )
+                self.quant_checks[quant_type].setStyleSheet("color: orange;")
+            else:
+                self.quant_checks[quant_type].setText(
+                    self.quant_types[quant_type],
+                )
+                self.quant_checks[quant_type].setStyleSheet("")
+
+    def _scroll_to_bottom(self) -> None:
+        """滚动输出框到底部."""
+        scrollbar = self.output_text.verticalScrollBar()
+        scrollbar.setValue(scrollbar.maximum())
+
+    def on_quant_type_changed(self, _state: int) -> None:
+        """量化类型变更时保存配置."""
+        selected_quants = [
+            q for q, check in self.quant_checks.items() if check.isChecked()
+        ]
+        conf.SELECTED_QUANTS = selected_quants
+
+    def moveEvent(self, event: QMoveEvent) -> None:  # noqa: N802
+        """处理窗口移动事件，保存窗口位置."""
+        top_left = self.geometry().topLeft()
+        conf.WIN_POS = [top_left.x(), top_left.y()]
+        return super().moveEvent(event)
+
+    def resizeEvent(self, event: QResizeEvent) -> None:  # noqa: N802
+        """处理窗口调整大小事件，保存窗口大小."""
+        geometry = self.geometry()
+        conf.WIN_SIZE = [geometry.width(), geometry.height()]
+        return super().resizeEvent(event)
+
+    def start_conversion(self) -> None:
+        """开始转换."""
+        # 检查是否已有任务在运行
+        if self.worker and self.worker.isRunning():
+            self.output_text.append("已有转换任务正在进行, 请等待完成")
+            self._scroll_to_bottom()
+            return
+
+        selected_quants: list[str] = [
+            q for q, check in self.quant_checks.items() if check.isChecked()
+        ]
+
+        if not selected_quants:
+            self.output_text.append("请至少选择一种量化类型")
+            self._scroll_to_bottom()
+            return
+
+        if not self.input_file:
+            self.output_text.append("请先选择输入文件")
+            self._scroll_to_bottom()
+            return
+
+        # 检查是否有已存在的文件将被覆盖
+        existing_files = []
+        for quant_type in selected_quants:
+            filename = f"{_process_gguf_stem(self.input_file.stem)}-{quant_type}.gguf"
+            expected_file = self.input_file.parent / filename
+            if expected_file.exists():
+                existing_files.append(filename)
+
+        if existing_files:
+            self.output_text.append("警告: 将覆盖以下已存在文件:")
+            for existing_file in existing_files:
+                self.output_text.append(f"  - {existing_file}")
+            self._scroll_to_bottom()
+
+        self.convert_btn.setEnabled(False)
+        self.progress_bar.setValue(0)
+        self.output_text.append(f"开始转换: {self.input_file!s}")
+        self.output_text.append(f"选择的量化类型: {', '.join(selected_quants)}")
+        self.output_text.append(f"将生成 {len(selected_quants)} 个量化文件")
+
+        self.worker = QuantizationWorker(self.input_file, selected_quants)
+        # 连接信号
+        self.worker.progress_msg_updated.connect(self.update_progress_msg)
+        self.worker.is_finished.connect(self.conversion_finished)
+        self.worker.progress_count_updated.connect(self.update_progress_value)
+        self.worker.start()
+
+    @Slot(str)
+    def update_progress_msg(self, message: str) -> None:
+        """更新进度信息."""
+        self.output_text.append(message)
+        self.output_text.ensureCursorVisible()
+        self._scroll_to_bottom()
+
+    @Slot(int)
+    def update_progress_value(self, value: int) -> None:
+        """更新进度条."""
+        self.progress_bar.setValue(value)
+
+    @Slot()
+    def conversion_finished(self) -> None:
+        """转换完成回调函数."""
+        success = self.worker.success if self.worker else False
+
+        # 断开信号连接，允许 worker 被垃圾回收
+        if self.worker:
+            try:
+                self.worker.progress_msg_updated.disconnect(self.update_progress_msg)
+                self.worker.is_finished.disconnect(self.conversion_finished)
+                self.worker.progress_count_updated.disconnect(
+                    self.update_progress_value
+                )
+            except RuntimeError:
+                # Worker 可能已经被销毁
+                pass
+
+        self.convert_btn.setEnabled(True)
+
+        if success:
+            self.output_text.append("所有量化转换完成!")
+            self.progress_bar.setValue(100)
+            # 重新检查已存在的量化文件
+            self.check_existing_quant_files()
+        else:
+            self.output_text.append("量化转换过程中出现错误!")
+            self._scroll_to_bottom()
+
+
+def main() -> None:
+    app = QApplication(sys.argv)
+
+    # 检查是否安装了llama.cpp
+    try:
+        subprocess.run(
+            ["llama-quantize", "--help"],
+            capture_output=True,
+            check=False,
+        )
+    except FileNotFoundError:
+        logger.exception("错误: 未找到llama.cpp/quantize工具")
+        logger.exception(
+            "请确保已编译llama.cpp并将quantize工具放在llama.cpp/目录下",
+        )
+        sys.exit(1)
+
+    window = GGUFQuantizerGUI()
+    window.show()
+    sys.exit(app.exec_())