char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Aleksei Nikiforovgguf-py : fix passing non-native endian tensors (editor-gui and new-metadata) (#17553)d82b7a7c1

char/aliceai-t5-moe
62.9 KiB1621 linesraw
1#!/usr/bin/env python3
2from __future__ import annotations
3
4import logging
5import argparse
6import os
7import sys
8import numpy
9import enum
10from pathlib import Path
11from typing import Any, Optional, Tuple, Type
12import warnings
13
14import numpy as np
15from PySide6.QtWidgets import (
16    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
17    QPushButton, QLabel, QLineEdit, QFileDialog, QTableWidget,
18    QTableWidgetItem, QComboBox, QMessageBox, QTabWidget,
19    QTextEdit, QFormLayout,
20    QHeaderView, QDialog, QDialogButtonBox
21)
22from PySide6.QtCore import Qt
23
24# Necessary to load the local gguf package
25if "NO_LOCAL_GGUF" not in os.environ and (Path(__file__).parent.parent.parent.parent / 'gguf-py').exists():
26    sys.path.insert(0, str(Path(__file__).parent.parent.parent))
27
28import gguf
29from gguf import GGUFReader, GGUFWriter, GGUFValueType, ReaderField
30from gguf.constants import TokenType, RopeScalingType, PoolingType, GGMLQuantizationType
31
32logger = logging.getLogger("gguf-editor-gui")
33
34# Map of key names to enum types for automatic enum interpretation
35KEY_TO_ENUM_TYPE = {
36    gguf.Keys.Tokenizer.TOKEN_TYPE: TokenType,
37    gguf.Keys.Rope.SCALING_TYPE: RopeScalingType,
38    gguf.Keys.LLM.POOLING_TYPE: PoolingType,
39    gguf.Keys.General.FILE_TYPE: GGMLQuantizationType,
40}
41
42# Define the tokenizer keys that should be edited together
43TOKENIZER_LINKED_KEYS = [
44    gguf.Keys.Tokenizer.LIST,
45    gguf.Keys.Tokenizer.TOKEN_TYPE,
46    gguf.Keys.Tokenizer.SCORES
47]
48
49
50class TokenizerEditorDialog(QDialog):
51    def __init__(self, tokens, token_types, scores, parent=None):
52        super().__init__(parent)
53        self.setWindowTitle("Edit Tokenizer Data")
54        self.resize(900, 600)
55
56        self.tokens = tokens.copy() if tokens else []
57        self.token_types = token_types.copy() if token_types else []
58        self.scores = scores.copy() if scores else []
59
60        # Ensure all arrays have the same length
61        max_len = max(len(self.tokens), len(self.token_types), len(self.scores))
62        if len(self.tokens) < max_len:
63            self.tokens.extend([""] * (max_len - len(self.tokens)))
64        if len(self.token_types) < max_len:
65            self.token_types.extend([0] * (max_len - len(self.token_types)))
66        if len(self.scores) < max_len:
67            self.scores.extend([0.0] * (max_len - len(self.scores)))
68
69        layout = QVBoxLayout(self)
70
71        # Add filter controls
72        filter_layout = QHBoxLayout()
73        filter_layout.addWidget(QLabel("Filter:"))
74        self.filter_edit = QLineEdit()
75        self.filter_edit.setPlaceholderText("Type to filter tokens...")
76        self.filter_edit.textChanged.connect(self.apply_filter)
77        filter_layout.addWidget(self.filter_edit)
78
79        # Add page controls
80        self.page_size = 100  # Show 100 items per page
81        self.current_page = 0
82        self.total_pages = max(1, (len(self.tokens) + self.page_size - 1) // self.page_size)
83
84        self.page_label = QLabel(f"Page 1 of {self.total_pages}")
85        filter_layout.addWidget(self.page_label)
86
87        prev_page = QPushButton("Previous")
88        prev_page.clicked.connect(self.previous_page)
89        filter_layout.addWidget(prev_page)
90
91        next_page = QPushButton("Next")
92        next_page.clicked.connect(self.next_page)
93        filter_layout.addWidget(next_page)
94
95        layout.addLayout(filter_layout)
96
97        # Tokenizer data table
98        self.tokens_table = QTableWidget()
99        self.tokens_table.setColumnCount(4)
100        self.tokens_table.setHorizontalHeaderLabels(["Index", "Token", "Type", "Score"])
101        self.tokens_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.ResizeMode.ResizeToContents)
102        self.tokens_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeMode.Stretch)
103        self.tokens_table.horizontalHeader().setSectionResizeMode(2, QHeaderView.ResizeMode.ResizeToContents)
104        self.tokens_table.horizontalHeader().setSectionResizeMode(3, QHeaderView.ResizeMode.ResizeToContents)
105
106        layout.addWidget(self.tokens_table)
107
108        # Controls
109        controls_layout = QHBoxLayout()
110
111        add_button = QPushButton("Add Token")
112        add_button.clicked.connect(self.add_token)
113        controls_layout.addWidget(add_button)
114
115        remove_button = QPushButton("Remove Selected")
116        remove_button.clicked.connect(self.remove_selected)
117        controls_layout.addWidget(remove_button)
118
119        controls_layout.addStretch()
120
121        layout.addLayout(controls_layout)
122
123        # Buttons
124        buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel)
125        buttons.accepted.connect(self.accept)
126        buttons.rejected.connect(self.reject)
127        layout.addWidget(buttons)
128
129        # Initialize the filtered values
130        self.filtered_indices = list(range(len(self.tokens)))
131
132        # Load data for the first page
133        self.load_page()
134
135    def apply_filter(self):
136        """Filter the tokens based on the search text."""
137        filter_text = self.filter_edit.text().lower()
138
139        if not filter_text:
140            # No filter, show all values
141            self.filtered_indices = list(range(len(self.tokens)))
142        else:
143            # Apply filter
144            self.filtered_indices = []
145            for i, token in enumerate(self.tokens):
146                if filter_text in str(token).lower():
147                    self.filtered_indices.append(i)
148
149        # Reset to first page and reload
150        self.total_pages = max(1, (len(self.filtered_indices) + self.page_size - 1) // self.page_size)
151        self.current_page = 0
152        self.page_label.setText(f"Page 1 of {self.total_pages}")
153        self.load_page()
154
155    def previous_page(self):
156        """Go to the previous page of results."""
157        if self.current_page > 0:
158            self.current_page -= 1
159            self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}")
160            self.load_page()
161
162    def next_page(self):
163        """Go to the next page of results."""
164        if self.current_page < self.total_pages - 1:
165            self.current_page += 1
166            self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}")
167            self.load_page()
168
169    def load_page(self):
170        """Load the current page of tokenizer data."""
171        self.tokens_table.setRowCount(0)  # Clear the table
172
173        # Calculate start and end indices for the current page
174        start_idx = self.current_page * self.page_size
175        end_idx = min(start_idx + self.page_size, len(self.filtered_indices))
176
177        # Pre-allocate rows for better performance
178        self.tokens_table.setRowCount(end_idx - start_idx)
179
180        for row, i in enumerate(range(start_idx, end_idx)):
181            orig_idx = self.filtered_indices[i]
182
183            # Index
184            index_item = QTableWidgetItem(str(orig_idx))
185            index_item.setData(Qt.ItemDataRole.UserRole, orig_idx)  # Store original index
186            index_item.setFlags(index_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
187            self.tokens_table.setItem(row, 0, index_item)
188
189            # Token
190            token_item = QTableWidgetItem(str(self.tokens[orig_idx]))
191            self.tokens_table.setItem(row, 1, token_item)
192
193            # Token Type
194            token_type = self.token_types[orig_idx] if orig_idx < len(self.token_types) else 0
195            try:
196                enum_val = TokenType(token_type)
197                display_text = f"{enum_val.name} ({token_type})"
198            except (ValueError, KeyError):
199                display_text = f"Unknown ({token_type})"
200
201            type_item = QTableWidgetItem(display_text)
202            type_item.setData(Qt.ItemDataRole.UserRole, token_type)
203
204            # Make type cell editable with a double-click handler
205            type_item.setFlags(type_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
206            self.tokens_table.setItem(row, 2, type_item)
207
208            # Score
209            score = self.scores[orig_idx] if orig_idx < len(self.scores) else 0.0
210            score_item = QTableWidgetItem(str(score))
211            self.tokens_table.setItem(row, 3, score_item)
212
213        # Connect double-click handler for token type cells
214        self.tokens_table.cellDoubleClicked.connect(self.handle_cell_double_click)
215
216    def handle_cell_double_click(self, row, column):
217        """Handle double-click on a cell, specifically for token type editing."""
218        if column == 2:  # Token Type column
219            orig_item = self.tokens_table.item(row, 0)
220            if orig_item:
221                orig_idx = orig_item.data(Qt.ItemDataRole.UserRole)
222                self.edit_token_type(row, orig_idx)
223
224    def edit_token_type(self, row, orig_idx):
225        """Edit a token type using a dialog with a dropdown of all enum options."""
226        current_value = self.token_types[orig_idx] if orig_idx < len(self.token_types) else 0
227
228        # Create a dialog with enum options
229        dialog = QDialog(self)
230        dialog.setWindowTitle("Select Token Type")
231        layout = QVBoxLayout(dialog)
232
233        combo = QComboBox()
234        for enum_val in TokenType:
235            combo.addItem(f"{enum_val.name} ({enum_val.value})", enum_val.value)
236
237        # Set current value
238        try:
239            if isinstance(current_value, int):
240                enum_val = TokenType(current_value)
241                combo.setCurrentText(f"{enum_val.name} ({current_value})")
242        except (ValueError, KeyError):
243            pass
244
245        layout.addWidget(combo)
246
247        buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel)
248        buttons.accepted.connect(dialog.accept)
249        buttons.rejected.connect(dialog.reject)
250        layout.addWidget(buttons)
251
252        if dialog.exec() == QDialog.DialogCode.Accepted:
253            # Get the selected value
254            new_value = combo.currentData()
255            enum_val = TokenType(new_value)
256            display_text = f"{enum_val.name} ({new_value})"
257
258            # Update the display
259            type_item = self.tokens_table.item(row, 2)
260            if type_item:
261                type_item.setText(display_text)
262                type_item.setData(Qt.ItemDataRole.UserRole, new_value)
263
264            # Update the actual value
265            self.token_types[orig_idx] = new_value
266
267    def add_token(self):
268        """Add a new token to the end of the list."""
269        # Add to the end of the arrays
270        self.tokens.append("")
271        self.token_types.append(0)  # Default to normal token
272        self.scores.append(0.0)
273
274        orig_idx = len(self.tokens) - 1
275
276        # Add to filtered indices if it matches the current filter
277        filter_text = self.filter_edit.text().lower()
278        if not filter_text or filter_text in "":
279            self.filtered_indices.append(orig_idx)
280
281        # Update pagination
282        self.total_pages = max(1, (len(self.filtered_indices) + self.page_size - 1) // self.page_size)
283
284        # Go to the last page to show the new item
285        self.current_page = self.total_pages - 1
286        self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}")
287
288        # Reload the page
289        self.load_page()
290
291    def remove_selected(self):
292        """Remove selected tokens from all arrays."""
293        selected_rows = []
294        for item in self.tokens_table.selectedItems():
295            row = item.row()
296            if row not in selected_rows:
297                selected_rows.append(row)
298
299        if not selected_rows:
300            return
301
302        # Get original indices in descending order to avoid index shifting
303        orig_indices = []
304        for row in selected_rows:
305            orig_item = self.tokens_table.item(row, 0)
306            if orig_item:
307                orig_indices.append(orig_item.data(Qt.ItemDataRole.UserRole))
308        orig_indices.sort(reverse=True)
309
310        # Remove from all arrays
311        for idx in orig_indices:
312            if idx < len(self.tokens):
313                del self.tokens[idx]
314            if idx < len(self.token_types):
315                del self.token_types[idx]
316            if idx < len(self.scores):
317                del self.scores[idx]
318
319        # Rebuild filtered_indices
320        self.filtered_indices = []
321        filter_text = self.filter_edit.text().lower()
322
323        for i, token in enumerate(self.tokens):
324            if not filter_text or filter_text in str(token).lower():
325                self.filtered_indices.append(i)
326
327        # Update pagination
328        self.total_pages = max(1, (len(self.filtered_indices) + self.page_size - 1) // self.page_size)
329        self.current_page = min(self.current_page, self.total_pages - 1)
330        self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}")
331
332        # Reload the page
333        self.load_page()
334
335    def get_data(self):
336        """Return the edited tokenizer data."""
337        return self.tokens, self.token_types, self.scores
338
339
340class ArrayEditorDialog(QDialog):
341    def __init__(self, array_values, element_type, key=None, parent=None):
342        super().__init__(parent)
343        self.setWindowTitle("Edit Array Values")
344        self.resize(700, 500)
345
346        self.array_values = array_values
347        self.element_type = element_type
348        self.key = key
349
350        # Get enum type for this array if applicable
351        self.enum_type = None
352        if key in KEY_TO_ENUM_TYPE and element_type == GGUFValueType.INT32:
353            self.enum_type = KEY_TO_ENUM_TYPE[key]
354
355        layout = QVBoxLayout(self)
356
357        # Add enum type information if applicable
358        if self.enum_type is not None:
359            enum_info_layout = QHBoxLayout()
360            enum_label = QLabel(f"Editing {self.enum_type.__name__} values:")
361            enum_info_layout.addWidget(enum_label)
362
363            # Add a legend for the enum values
364            enum_values = ", ".join([f"{e.name}={e.value}" for e in self.enum_type])
365            enum_values_label = QLabel(f"Available values: {enum_values}")
366            enum_values_label.setWordWrap(True)
367            enum_info_layout.addWidget(enum_values_label, 1)
368
369            layout.addLayout(enum_info_layout)
370
371        # Add search/filter controls
372        filter_layout = QHBoxLayout()
373        filter_layout.addWidget(QLabel("Filter:"))
374        self.filter_edit = QLineEdit()
375        self.filter_edit.setPlaceholderText("Type to filter values...")
376        self.filter_edit.textChanged.connect(self.apply_filter)
377        filter_layout.addWidget(self.filter_edit)
378
379        # Add page controls for large arrays
380        self.page_size = 100  # Show 100 items per page
381        self.current_page = 0
382        self.total_pages = max(1, (len(array_values) + self.page_size - 1) // self.page_size)
383
384        self.page_label = QLabel(f"Page 1 of {self.total_pages}")
385        filter_layout.addWidget(self.page_label)
386
387        prev_page = QPushButton("Previous")
388        prev_page.clicked.connect(self.previous_page)
389        filter_layout.addWidget(prev_page)
390
391        next_page = QPushButton("Next")
392        next_page.clicked.connect(self.next_page)
393        filter_layout.addWidget(next_page)
394
395        layout.addLayout(filter_layout)
396
397        # Array items table
398        self.items_table = QTableWidget()
399
400        # Set up columns based on whether we have an enum type
401        if self.enum_type is not None:
402            self.items_table.setColumnCount(3)
403            self.items_table.setHorizontalHeaderLabels(["Index", "Value", "Actions"])
404            self.items_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.ResizeMode.ResizeToContents)
405            self.items_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeMode.Stretch)
406            self.items_table.horizontalHeader().setSectionResizeMode(2, QHeaderView.ResizeMode.ResizeToContents)
407        else:
408            self.items_table.setColumnCount(2)
409            self.items_table.setHorizontalHeaderLabels(["Index", "Value"])
410            self.items_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.ResizeMode.ResizeToContents)
411            self.items_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeMode.Stretch)
412
413        layout.addWidget(self.items_table)
414
415        # Controls
416        controls_layout = QHBoxLayout()
417
418        add_button = QPushButton("Add Item")
419        add_button.clicked.connect(self.add_item)
420        controls_layout.addWidget(add_button)
421
422        remove_button = QPushButton("Remove Selected")
423        remove_button.clicked.connect(self.remove_selected)
424        controls_layout.addWidget(remove_button)
425
426        # Add bulk edit button for enum arrays
427        if self.enum_type is not None:
428            bulk_edit_button = QPushButton("Bulk Edit Selected")
429            bulk_edit_button.clicked.connect(self.bulk_edit_selected)
430            controls_layout.addWidget(bulk_edit_button)
431
432        controls_layout.addStretch()
433
434        layout.addLayout(controls_layout)
435
436        # Buttons
437        buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel)
438        buttons.accepted.connect(self.accept)
439        buttons.rejected.connect(self.reject)
440        layout.addWidget(buttons)
441
442        # Initialize the filtered values
443        self.filtered_indices = list(range(len(self.array_values)))
444
445        # Load array values for the first page
446        self.load_page()
447
448    def apply_filter(self):
449        """Filter the array values based on the search text."""
450        filter_text = self.filter_edit.text().lower()
451
452        if not filter_text:
453            # No filter, show all values
454            self.filtered_indices = list(range(len(self.array_values)))
455        else:
456            # Apply filter
457            self.filtered_indices = []
458            for i, value in enumerate(self.array_values):
459                # For enum values, search in both name and value
460                if self.enum_type is not None and isinstance(value, int):
461                    try:
462                        enum_val = self.enum_type(value)
463                        display_text = f"{enum_val.name} ({value})".lower()
464                        if filter_text in display_text:
465                            self.filtered_indices.append(i)
466                    except (ValueError, KeyError):
467                        # If not a valid enum value, just check the raw value
468                        if filter_text in str(value).lower():
469                            self.filtered_indices.append(i)
470                else:
471                    # For non-enum values, just check the string representation
472                    if filter_text in str(value).lower():
473                        self.filtered_indices.append(i)
474
475        # Reset to first page and reload
476        self.total_pages = max(1, (len(self.filtered_indices) + self.page_size - 1) // self.page_size)
477        self.current_page = 0
478        self.page_label.setText(f"Page 1 of {self.total_pages}")
479        self.load_page()
480
481    def previous_page(self):
482        """Go to the previous page of results."""
483        if self.current_page > 0:
484            self.current_page -= 1
485            self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}")
486            self.load_page()
487
488    def next_page(self):
489        """Go to the next page of results."""
490        if self.current_page < self.total_pages - 1:
491            self.current_page += 1
492            self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}")
493            self.load_page()
494
495    def load_page(self):
496        """Load the current page of array values."""
497        self.items_table.setRowCount(0)  # Clear the table
498
499        # Calculate start and end indices for the current page
500        start_idx = self.current_page * self.page_size
501        end_idx = min(start_idx + self.page_size, len(self.filtered_indices))
502
503        # Pre-allocate rows for better performance
504        self.items_table.setRowCount(end_idx - start_idx)
505
506        for row, i in enumerate(range(start_idx, end_idx)):
507            orig_idx = self.filtered_indices[i]
508            value = self.array_values[orig_idx]
509
510            # Index
511            index_item = QTableWidgetItem(str(orig_idx))
512            index_item.setData(Qt.ItemDataRole.UserRole, orig_idx)  # Store original index
513            index_item.setFlags(index_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
514            self.items_table.setItem(row, 0, index_item)
515
516            # Value
517            if self.enum_type is not None:
518                # Display enum value and name
519                try:
520                    if isinstance(value, (int, numpy.signedinteger)):
521                        enum_val = self.enum_type(value)
522                        display_text = f"{enum_val.name} ({value})"
523                    else:
524                        display_text = str(value)
525                except (ValueError, KeyError):
526                    display_text = f"Unknown ({value})"
527
528                # Store the enum value in the item
529                value_item = QTableWidgetItem(display_text)
530                value_item.setData(Qt.ItemDataRole.UserRole, value)
531                value_item.setFlags(value_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
532                self.items_table.setItem(row, 1, value_item)
533
534                # Add an edit button in a separate column
535                edit_button = QPushButton("Edit")
536                edit_button.setProperty("row", row)
537                edit_button.clicked.connect(self.edit_array_enum_value)
538
539                # Create a widget to hold the button
540                button_widget = QWidget()
541                button_layout = QHBoxLayout(button_widget)
542                button_layout.setContentsMargins(2, 2, 2, 2)
543                button_layout.addWidget(edit_button)
544                button_layout.addStretch()
545
546                self.items_table.setCellWidget(row, 2, button_widget)
547            else:
548                value_item = QTableWidgetItem(str(value))
549                self.items_table.setItem(row, 1, value_item)
550
551    def edit_array_enum_value(self):
552        """Handle editing an enum value in the array editor."""
553        button = self.sender()
554        row = button.property("row")
555
556        # Get the original index from the table item
557        orig_item = self.items_table.item(row, 0)
558        new_item = self.items_table.item(row, 1)
559        if orig_item and new_item and self.enum_type and self.edit_enum_value(row, self.enum_type):
560            orig_idx = orig_item.data(Qt.ItemDataRole.UserRole)
561            new_value = new_item.data(Qt.ItemDataRole.UserRole)
562            # Update the stored value in the array
563            if isinstance(new_value, (int, float, str, bool)):
564                self.array_values[orig_idx] = new_value
565
566    def bulk_edit_selected(self):
567        """Edit multiple enum values at once."""
568        if not self.enum_type:
569            return
570
571        selected_rows = set()
572        for item in self.items_table.selectedItems():
573            selected_rows.add(item.row())
574
575        if not selected_rows:
576            QMessageBox.information(self, "No Selection", "Please select at least one row to edit.")
577            return
578
579        # Create a dialog with enum options
580        dialog = QDialog(self)
581        dialog.setWindowTitle(f"Bulk Edit {self.enum_type.__name__} Values")
582        layout = QVBoxLayout(dialog)
583
584        layout.addWidget(QLabel(f"Set {len(selected_rows)} selected items to:"))
585
586        combo = QComboBox()
587        for enum_val in self.enum_type:
588            combo.addItem(f"{enum_val.name} ({enum_val.value})", enum_val.value)
589
590        layout.addWidget(combo)
591
592        buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel)
593        buttons.accepted.connect(dialog.accept)
594        buttons.rejected.connect(dialog.reject)
595        layout.addWidget(buttons)
596
597        if dialog.exec() == QDialog.DialogCode.Accepted:
598            # Get the selected value
599            new_value = combo.currentData()
600            enum_val = self.enum_type(new_value)
601            display_text = f"{enum_val.name} ({new_value})"
602
603            # Update all selected rows
604            for row in selected_rows:
605                orig_item = self.items_table.item(row, 0)
606                new_item = self.items_table.item(row, 1)
607                if orig_item and new_item:
608                    orig_idx = orig_item.data(Qt.ItemDataRole.UserRole)
609                    self.array_values[orig_idx] = new_value
610
611                    # Update the display
612                    new_item.setText(display_text)
613                    new_item.setData(Qt.ItemDataRole.UserRole, new_value)
614
615    def add_item(self):
616        # Add to the end of the array
617        orig_idx = len(self.array_values)
618
619        # Add default value based on type
620        if self.enum_type is not None:
621            # Default to first enum value
622            default_value = list(self.enum_type)[0].value
623            self.array_values.append(default_value)
624        else:
625            if self.element_type == GGUFValueType.STRING:
626                self.array_values.append("")
627            else:
628                self.array_values.append(0)
629
630        # Add to filtered indices if it matches the current filter
631        self.filtered_indices.append(orig_idx)
632
633        # Update pagination
634        self.total_pages = max(1, (len(self.filtered_indices) + self.page_size - 1) // self.page_size)
635
636        # Go to the last page to show the new item
637        self.current_page = self.total_pages - 1
638        self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}")
639
640        # Reload the page
641        self.load_page()
642
643    def remove_selected(self):
644        selected_rows = []
645        for item in self.items_table.selectedItems():
646            row = item.row()
647            if row not in selected_rows:
648                selected_rows.append(row)
649
650        if not selected_rows:
651            return
652
653        # Get original indices in descending order to avoid index shifting
654        orig_indices = list()
655        for row in selected_rows:
656            orig_item = self.items_table.item(row, 0)
657            if orig_item:
658                orig_indices.append(orig_item.data(Qt.ItemDataRole.UserRole))
659        orig_indices.sort(reverse=True)
660
661        # Remove from array_values
662        for idx in orig_indices:
663            del self.array_values[idx]
664
665        # Rebuild filtered_indices
666        self.filtered_indices = []
667        filter_text = self.filter_edit.text().lower()
668
669        for i, value in enumerate(self.array_values):
670            if not filter_text:
671                self.filtered_indices.append(i)
672            else:
673                # Apply filter
674                if self.enum_type is not None and isinstance(value, int):
675                    try:
676                        enum_val = self.enum_type(value)
677                        display_text = f"{enum_val.name} ({value})".lower()
678                        if filter_text in display_text:
679                            self.filtered_indices.append(i)
680                    except (ValueError, KeyError):
681                        if filter_text in str(value).lower():
682                            self.filtered_indices.append(i)
683                else:
684                    if filter_text in str(value).lower():
685                        self.filtered_indices.append(i)
686
687        # Update pagination
688        self.total_pages = max(1, (len(self.filtered_indices) + self.page_size - 1) // self.page_size)
689        self.current_page = min(self.current_page, self.total_pages - 1)
690        self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}")
691
692        # Reload the page
693        self.load_page()
694
695    def edit_enum_value(self, row: int, enum_type: Type[enum.Enum]):
696        """Edit an enum value using a dialog with a dropdown of all enum options."""
697        # Get the original index from the table item
698        orig_item = self.items_table.item(row, 0)
699        if orig_item:
700            orig_idx = orig_item.data(Qt.ItemDataRole.UserRole)
701        else:
702            return
703        current_value = self.array_values[orig_idx]
704
705        # Create a dialog with enum options
706        dialog = QDialog(self)
707        dialog.setWindowTitle(f"Select {enum_type.__name__} Value")
708        layout = QVBoxLayout(dialog)
709
710        # Add description
711        description = QLabel(f"Select a {enum_type.__name__} value:")
712        layout.addWidget(description)
713
714        # Use a combo box for quick selection
715        combo = QComboBox()
716        for enum_val in enum_type:
717            combo.addItem(f"{enum_val.name} ({enum_val.value})", enum_val.value)
718
719        # Set current value
720        try:
721            if isinstance(current_value, int):
722                enum_val = enum_type(current_value)
723                combo.setCurrentText(f"{enum_val.name} ({current_value})")
724        except (ValueError, KeyError):
725            pass
726
727        layout.addWidget(combo)
728
729        buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel)
730        buttons.accepted.connect(dialog.accept)
731        buttons.rejected.connect(dialog.reject)
732        layout.addWidget(buttons)
733
734        if dialog.exec() == QDialog.DialogCode.Accepted:
735            # Update the value display and stored data
736            new_value = combo.currentData()
737            enum_val = enum_type(new_value)
738            display_text = f"{enum_val.name} ({new_value})"
739
740            new_item = self.items_table.item(row, 1)
741            if new_item:
742                new_item.setText(display_text)
743                new_item.setData(Qt.ItemDataRole.UserRole, new_value)
744
745            # Update the actual array value
746            self.array_values[orig_idx] = new_value
747            return True
748        return False
749
750    def get_array_values(self):
751        # The array_values list is kept up-to-date as edits are made
752        return self.array_values
753
754
755class AddMetadataDialog(QDialog):
756    def __init__(self, parent=None):
757        super().__init__(parent)
758        self.setWindowTitle("Add Metadata")
759        self.resize(400, 200)
760
761        layout = QVBoxLayout(self)
762
763        form_layout = QFormLayout()
764
765        self.key_edit = QLineEdit()
766        form_layout.addRow("Key:", self.key_edit)
767
768        self.type_combo = QComboBox()
769        for value_type in GGUFValueType:
770            if value_type != GGUFValueType.ARRAY:  # Skip array type for simplicity
771                self.type_combo.addItem(value_type.name, value_type)
772        form_layout.addRow("Type:", self.type_combo)
773
774        self.value_edit = QTextEdit()
775        form_layout.addRow("Value:", self.value_edit)
776
777        layout.addLayout(form_layout)
778
779        buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel)
780        buttons.accepted.connect(self.accept)
781        buttons.rejected.connect(self.reject)
782        layout.addWidget(buttons)
783
784    def get_data(self) -> Tuple[str, GGUFValueType, Any]:
785        key = self.key_edit.text()
786        value_type = self.type_combo.currentData()
787        value_text = self.value_edit.toPlainText()
788
789        # Convert value based on type
790        if value_type == GGUFValueType.UINT8:
791            value = np.uint8(int(value_text))
792        elif value_type == GGUFValueType.INT8:
793            value = np.int8(int(value_text))
794        elif value_type == GGUFValueType.UINT16:
795            value = np.uint16(int(value_text))
796        elif value_type == GGUFValueType.INT16:
797            value = np.int16(int(value_text))
798        elif value_type == GGUFValueType.UINT32:
799            value = np.uint32(int(value_text))
800        elif value_type == GGUFValueType.INT32:
801            value = np.int32(int(value_text))
802        elif value_type == GGUFValueType.FLOAT32:
803            value = np.float32(float(value_text))
804        elif value_type == GGUFValueType.BOOL:
805            value = value_text.lower() in ('true', 'yes', '1')
806        elif value_type == GGUFValueType.STRING:
807            value = value_text
808        else:
809            value = value_text
810
811        return key, value_type, value
812
813
814class GGUFEditorWindow(QMainWindow):
815    def __init__(self):
816        super().__init__()
817
818        self.setWindowTitle("GGUF Editor")
819        self.resize(1000, 800)
820
821        self.current_file = None
822        self.reader = None
823        self.modified = False
824        self.metadata_changes = {}  # Store changes to apply when saving
825        self.metadata_to_remove = set()  # Store keys to remove when saving
826        self.on_metadata_changed_is_connected = False
827
828        self.setup_ui()
829
830    def setup_ui(self):
831        central_widget = QWidget()
832        self.setCentralWidget(central_widget)
833
834        main_layout = QVBoxLayout(central_widget)
835
836        # File controls
837        file_layout = QHBoxLayout()
838
839        self.file_path_edit = QLineEdit()
840        self.file_path_edit.setReadOnly(True)
841        file_layout.addWidget(self.file_path_edit)
842
843        open_button = QPushButton("Open GGUF")
844        open_button.clicked.connect(self.open_file)
845        file_layout.addWidget(open_button)
846
847        save_button = QPushButton("Save As...")
848        save_button.clicked.connect(self.save_file)
849        file_layout.addWidget(save_button)
850
851        main_layout.addLayout(file_layout)
852
853        # Tabs for different views
854        self.tabs = QTabWidget()
855
856        # Metadata tab
857        self.metadata_tab = QWidget()
858        metadata_layout = QVBoxLayout(self.metadata_tab)
859
860        # Metadata table
861        self.metadata_table = QTableWidget()
862        self.metadata_table.setColumnCount(4)
863        self.metadata_table.setHorizontalHeaderLabels(["Key", "Type", "Value", "Actions"])
864        self.metadata_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.ResizeMode.Stretch)
865        self.metadata_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeMode.ResizeToContents)
866        self.metadata_table.horizontalHeader().setSectionResizeMode(2, QHeaderView.ResizeMode.Stretch)
867        self.metadata_table.horizontalHeader().setSectionResizeMode(3, QHeaderView.ResizeMode.ResizeToContents)
868        metadata_layout.addWidget(self.metadata_table)
869
870        # Metadata controls
871        metadata_controls = QHBoxLayout()
872
873        add_metadata_button = QPushButton("Add Metadata")
874        add_metadata_button.clicked.connect(self.add_metadata)
875        metadata_controls.addWidget(add_metadata_button)
876
877        metadata_controls.addStretch()
878
879        metadata_layout.addLayout(metadata_controls)
880
881        # Tensors tab
882        self.tensors_tab = QWidget()
883        tensors_layout = QVBoxLayout(self.tensors_tab)
884
885        self.tensors_table = QTableWidget()
886        self.tensors_table.setColumnCount(5)
887        self.tensors_table.setHorizontalHeaderLabels(["Name", "Type", "Shape", "Elements", "Size (bytes)"])
888        self.tensors_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.ResizeMode.Stretch)
889        self.tensors_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeMode.ResizeToContents)
890        self.tensors_table.horizontalHeader().setSectionResizeMode(2, QHeaderView.ResizeMode.ResizeToContents)
891        self.tensors_table.horizontalHeader().setSectionResizeMode(3, QHeaderView.ResizeMode.ResizeToContents)
892        self.tensors_table.horizontalHeader().setSectionResizeMode(4, QHeaderView.ResizeMode.ResizeToContents)
893        tensors_layout.addWidget(self.tensors_table)
894
895        # Add tabs to tab widget
896        self.tabs.addTab(self.metadata_tab, "Metadata")
897        self.tabs.addTab(self.tensors_tab, "Tensors")
898
899        main_layout.addWidget(self.tabs)
900
901        # Status bar
902        self.statusBar().showMessage("Ready")
903
904    def load_file(self, file_path):
905        """Load a GGUF file by path"""
906        try:
907            self.statusBar().showMessage(f"Loading {file_path}...")
908            QApplication.processEvents()
909
910            self.reader = GGUFReader(file_path, 'r')
911            self.current_file = file_path
912            self.file_path_edit.setText(file_path)
913
914            self.load_metadata()
915            self.load_tensors()
916
917            self.metadata_changes = {}
918            self.metadata_to_remove = set()
919            self.modified = False
920
921            self.statusBar().showMessage(f"Loaded {file_path}")
922            return True
923        except Exception as e:
924            QMessageBox.critical(self, "Error", f"Failed to open file: {str(e)}")
925            self.statusBar().showMessage("Error loading file")
926            return False
927
928    def open_file(self):
929        file_path, _ = QFileDialog.getOpenFileName(
930            self, "Open GGUF File", "", "GGUF Files (*.gguf);;All Files (*)"
931        )
932
933        if not file_path:
934            return
935
936        self.load_file(file_path)
937
938    def load_metadata(self):
939        self.metadata_table.setRowCount(0)
940
941        if not self.reader:
942            return
943
944        # Disconnect to prevent triggering during loading
945        if self.on_metadata_changed_is_connected:
946            with warnings.catch_warnings():
947                warnings.filterwarnings('ignore')
948                self.metadata_table.itemChanged.disconnect(self.on_metadata_changed)
949            self.on_metadata_changed_is_connected = False
950
951        for i, (key, field) in enumerate(self.reader.fields.items()):
952            self.metadata_table.insertRow(i)
953
954            # Key
955            key_item = QTableWidgetItem(key)
956            key_item.setFlags(key_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
957            self.metadata_table.setItem(i, 0, key_item)
958
959            # Type
960            if not field.types:
961                type_str = "N/A"
962            elif field.types[0] == GGUFValueType.ARRAY:
963                nest_count = len(field.types) - 1
964                element_type = field.types[-1].name
965                # Check if this is an enum array
966                enum_type = self.get_enum_for_key(key)
967                if enum_type is not None and field.types[-1] == GGUFValueType.INT32:
968                    element_type = enum_type.__name__
969                type_str = '[' * nest_count + element_type + ']' * nest_count
970            else:
971                type_str = str(field.types[0].name)
972                # Check if this is an enum field
973                enum_type = self.get_enum_for_key(key)
974                if enum_type is not None and field.types[0] == GGUFValueType.INT32:
975                    type_str = enum_type.__name__
976
977            type_item = QTableWidgetItem(type_str)
978            type_item.setFlags(type_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
979            self.metadata_table.setItem(i, 1, type_item)
980
981            # Value
982            value_str = self.format_field_value(field)
983            value_item = QTableWidgetItem(value_str)
984
985            # Make only simple values editable
986            if len(field.types) == 1 and field.types[0] != GGUFValueType.ARRAY:
987                value_item.setFlags(value_item.flags() | Qt.ItemFlag.ItemIsEditable)
988            else:
989                value_item.setFlags(value_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
990
991            self.metadata_table.setItem(i, 2, value_item)
992
993            # Actions
994            actions_widget = QWidget()
995            actions_layout = QHBoxLayout(actions_widget)
996            actions_layout.setContentsMargins(2, 2, 2, 2)
997
998            # Add Edit button for arrays and enum fields
999            if field.types and field.types[0] == GGUFValueType.ARRAY:
1000                edit_button = QPushButton("Edit")
1001                edit_button.setProperty("row", i)
1002                edit_button.setProperty("key", key)
1003                edit_button.clicked.connect(self.edit_array_metadata)
1004                actions_layout.addWidget(edit_button)
1005
1006                # Add special label for tokenizer linked fields
1007                if key in TOKENIZER_LINKED_KEYS:
1008                    edit_button.setText("Edit Tokenizer")
1009                    edit_button.setToolTip("Edit all tokenizer data together")
1010            elif len(field.types) == 1 and self.get_enum_for_key(key) is not None:
1011                edit_button = QPushButton("Edit")
1012                edit_button.setProperty("row", i)
1013                edit_button.setProperty("key", key)
1014                edit_button.clicked.connect(self.edit_metadata_enum)
1015                actions_layout.addWidget(edit_button)
1016
1017            remove_button = QPushButton("Remove")
1018            remove_button.setProperty("row", i)
1019            remove_button.setProperty("key", key)
1020            remove_button.clicked.connect(self.remove_metadata)
1021            actions_layout.addWidget(remove_button)
1022
1023            self.metadata_table.setCellWidget(i, 3, actions_widget)
1024
1025        # Reconnect after loading
1026        self.metadata_table.itemChanged.connect(self.on_metadata_changed)
1027        self.on_metadata_changed_is_connected = True
1028
1029    def extract_array_values(self, field: ReaderField) -> list:
1030        """Extract all values from an array field."""
1031        if not field.types or field.types[0] != GGUFValueType.ARRAY:
1032            return []
1033
1034        curr_type = field.types[1]
1035        array_values = []
1036        total_elements = len(field.data)
1037
1038        if curr_type == GGUFValueType.STRING:
1039            for element_pos in range(total_elements):
1040                value_string = str(bytes(field.parts[-1 - (total_elements - element_pos - 1) * 2]), encoding='utf-8')
1041                array_values.append(value_string)
1042        elif self.reader and curr_type in self.reader.gguf_scalar_to_np:
1043            for element_pos in range(total_elements):
1044                array_values.append(field.parts[-1 - (total_elements - element_pos - 1)][0])
1045
1046        return array_values
1047
1048    def get_enum_for_key(self, key: str) -> Optional[Type[enum.Enum]]:
1049        """Get the enum type for a given key if it exists."""
1050        return KEY_TO_ENUM_TYPE.get(key)
1051
1052    def format_enum_value(self, value: Any, enum_type: Type[enum.Enum]) -> str:
1053        """Format a value as an enum if possible."""
1054        try:
1055            if isinstance(value, (int, str)):
1056                enum_value = enum_type(value)
1057                return f"{enum_value.name} ({value})"
1058        except (ValueError, KeyError):
1059            pass
1060        return str(value)
1061
1062    def format_field_value(self, field: ReaderField) -> str:
1063        if not field.types:
1064            return "N/A"
1065
1066        if len(field.types) == 1:
1067            curr_type = field.types[0]
1068            if curr_type == GGUFValueType.STRING:
1069                return str(bytes(field.parts[-1]), encoding='utf-8')
1070            elif self.reader and curr_type in self.reader.gguf_scalar_to_np:
1071                value = field.parts[-1][0]
1072                # Check if this field has an enum type
1073                enum_type = self.get_enum_for_key(field.name)
1074                if enum_type is not None:
1075                    return self.format_enum_value(value, enum_type)
1076                return str(value)
1077
1078        if field.types[0] == GGUFValueType.ARRAY:
1079            array_values = self.extract_array_values(field)
1080            render_element = min(5, len(array_values))
1081
1082            # Get enum type for this array if applicable
1083            enum_type = self.get_enum_for_key(field.name)
1084
1085            if enum_type is not None:
1086                array_elements = []
1087                for i in range(render_element):
1088                    array_elements.append(self.format_enum_value(array_values[i], enum_type))
1089            else:
1090                array_elements = [str(array_values[i]) for i in range(render_element)]
1091
1092            return f"[ {', '.join(array_elements).strip()}{', ...' if len(array_values) > len(array_elements) else ''} ]"
1093
1094        return "Complex value"
1095
1096    def load_tensors(self):
1097        self.tensors_table.setRowCount(0)
1098
1099        if not self.reader:
1100            return
1101
1102        for i, tensor in enumerate(self.reader.tensors):
1103            self.tensors_table.insertRow(i)
1104
1105            # Name
1106            name_item = QTableWidgetItem(tensor.name)
1107            name_item.setFlags(name_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
1108            self.tensors_table.setItem(i, 0, name_item)
1109
1110            # Type
1111            type_item = QTableWidgetItem(tensor.tensor_type.name)
1112            type_item.setFlags(type_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
1113            self.tensors_table.setItem(i, 1, type_item)
1114
1115            # Shape
1116            shape_str = " × ".join(str(d) for d in tensor.shape)
1117            shape_item = QTableWidgetItem(shape_str)
1118            shape_item.setFlags(shape_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
1119            self.tensors_table.setItem(i, 2, shape_item)
1120
1121            # Elements
1122            elements_item = QTableWidgetItem(str(tensor.n_elements))
1123            elements_item.setFlags(elements_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
1124            self.tensors_table.setItem(i, 3, elements_item)
1125
1126            # Size
1127            size_item = QTableWidgetItem(f"{tensor.n_bytes:,}")
1128            size_item.setFlags(size_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
1129            self.tensors_table.setItem(i, 4, size_item)
1130
1131    def on_metadata_changed(self, item):
1132        if item.column() != 2:  # Only handle value column changes
1133            return
1134
1135        row = item.row()
1136        orig_item = self.metadata_table.item(row, 0)
1137        key = None
1138        if orig_item:
1139            key = orig_item.text()
1140        new_value = item.text()
1141
1142        field = None
1143        if self.reader and key:
1144            field = self.reader.get_field(key)
1145        if not field or not field.types or not key:
1146            return
1147
1148        value_type = field.types[0]
1149
1150        # Check if this is an enum field
1151        enum_type = self.get_enum_for_key(key)
1152        if enum_type is not None and value_type == GGUFValueType.INT32:
1153            # Try to parse the enum value from the text
1154            try:
1155                # Check if it's a name
1156                try:
1157                    enum_val = enum_type[new_value]
1158                    converted_value = enum_val.value
1159                except (KeyError, AttributeError):
1160                    # Check if it's a number or "NAME (value)" format
1161                    if '(' in new_value and ')' in new_value:
1162                        # Extract the value from "NAME (value)" format
1163                        value_part = new_value.split('(')[1].split(')')[0].strip()
1164                        converted_value = int(value_part)
1165                    else:
1166                        # Try to convert directly to int
1167                        converted_value = int(new_value)
1168
1169                # Validate that it's a valid enum value
1170                enum_type(converted_value)
1171
1172                # Store the change
1173                self.metadata_changes[key] = (value_type, converted_value)
1174                self.modified = True
1175
1176                # Update display with formatted enum value
1177                formatted_value = self.format_enum_value(converted_value, enum_type)
1178                item.setText(formatted_value)
1179
1180                self.statusBar().showMessage(f"Changed {key} to {formatted_value}")
1181                return
1182            except (ValueError, KeyError) as e:
1183                QMessageBox.warning(
1184                    self,
1185                    f"Invalid Enum Value ({e})",
1186                    f"'{new_value}' is not a valid {enum_type.__name__} value.\n"
1187                    f"Valid values are: {', '.join(v.name for v in enum_type)}")
1188
1189                # Revert to original value
1190                original_value = self.format_field_value(field)
1191                item.setText(original_value)
1192                return
1193
1194        try:
1195            # Convert the string value to the appropriate type
1196            if value_type == GGUFValueType.UINT8:
1197                converted_value = np.uint8(int(new_value))
1198            elif value_type == GGUFValueType.INT8:
1199                converted_value = np.int8(int(new_value))
1200            elif value_type == GGUFValueType.UINT16:
1201                converted_value = np.uint16(int(new_value))
1202            elif value_type == GGUFValueType.INT16:
1203                converted_value = np.int16(int(new_value))
1204            elif value_type == GGUFValueType.UINT32:
1205                converted_value = np.uint32(int(new_value))
1206            elif value_type == GGUFValueType.INT32:
1207                converted_value = np.int32(int(new_value))
1208            elif value_type == GGUFValueType.FLOAT32:
1209                converted_value = np.float32(float(new_value))
1210            elif value_type == GGUFValueType.BOOL:
1211                converted_value = new_value.lower() in ('true', 'yes', '1')
1212            elif value_type == GGUFValueType.STRING:
1213                converted_value = new_value
1214            else:
1215                # Unsupported type for editing
1216                return
1217
1218            # Store the change
1219            self.metadata_changes[key] = (value_type, converted_value)
1220            self.modified = True
1221
1222            self.statusBar().showMessage(f"Changed {key} to {new_value}")
1223        except ValueError:
1224            QMessageBox.warning(self, "Invalid Value", f"The value '{new_value}' is not valid for type {value_type.name}")
1225
1226            # Revert to original value
1227            original_value = self.format_field_value(field)
1228            item.setText(original_value)
1229
1230    def remove_metadata(self):
1231        button = self.sender()
1232        key = button.property("key")
1233        row = button.property("row")
1234
1235        reply = QMessageBox.question(
1236            self, "Confirm Removal",
1237            f"Are you sure you want to remove the metadata key '{key}'?",
1238            QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, QMessageBox.StandardButton.No
1239        )
1240
1241        if reply == QMessageBox.StandardButton.Yes:
1242            self.metadata_table.removeRow(row)
1243            self.metadata_to_remove.add(key)
1244
1245            # If we previously had changes for this key, remove them
1246            if key in self.metadata_changes:
1247                del self.metadata_changes[key]
1248
1249            self.modified = True
1250            self.statusBar().showMessage(f"Marked {key} for removal")
1251
1252    def edit_metadata_enum(self):
1253        """Edit an enum metadata field."""
1254        button = self.sender()
1255        key = button.property("key")
1256        row = button.property("row")
1257
1258        field = None
1259        if self.reader:
1260            field = self.reader.get_field(key)
1261        if not field or not field.types:
1262            return
1263
1264        enum_type = self.get_enum_for_key(key)
1265        if enum_type is None:
1266            return
1267
1268        # Get current value
1269        current_value = field.contents()
1270
1271        # Create a dialog with enum options
1272        dialog = QDialog(self)
1273        dialog.setWindowTitle(f"Select {enum_type.__name__} Value")
1274        layout = QVBoxLayout(dialog)
1275
1276        combo = QComboBox()
1277        for enum_val in enum_type:
1278            combo.addItem(f"{enum_val.name} ({enum_val.value})", enum_val.value)
1279
1280        # Set current value
1281        try:
1282            if isinstance(current_value, (int, str)):
1283                enum_val = enum_type(current_value)
1284                combo.setCurrentText(f"{enum_val.name} ({current_value})")
1285        except (ValueError, KeyError):
1286            pass
1287
1288        layout.addWidget(combo)
1289
1290        buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel)
1291        buttons.accepted.connect(dialog.accept)
1292        buttons.rejected.connect(dialog.reject)
1293        layout.addWidget(buttons)
1294
1295        if dialog.exec() == QDialog.DialogCode.Accepted:
1296            # Get the selected value
1297            new_value = combo.currentData()
1298            enum_val = enum_type(new_value)
1299
1300            # Store the change
1301            self.metadata_changes[key] = (field.types[0], new_value)
1302            self.modified = True
1303
1304            # Update display
1305            display_text = f"{enum_val.name} ({new_value})"
1306            target_item = self.metadata_table.item(row, 2)
1307            if target_item:
1308                target_item.setText(display_text)
1309
1310            self.statusBar().showMessage(f"Changed {key} to {display_text}")
1311
1312    def edit_array_metadata(self):
1313        button = self.sender()
1314        key = button.property("key")
1315        row = button.property("row")
1316
1317        # Check if this is one of the linked tokenizer keys
1318        if key in TOKENIZER_LINKED_KEYS:
1319            self.edit_tokenizer_metadata(key)
1320            return
1321
1322        field = None
1323        if self.reader:
1324            field = self.reader.get_field(key)
1325        if not field or not field.types or field.types[0] != GGUFValueType.ARRAY:
1326            return
1327
1328        # Get array element type
1329        element_type = field.types[1]
1330
1331        # Extract array values
1332        array_values = self.extract_array_values(field)
1333
1334        # Open array editor dialog
1335        dialog = ArrayEditorDialog(array_values, element_type, key, self)
1336        if dialog.exec() == QDialog.DialogCode.Accepted:
1337            new_values = dialog.get_array_values()
1338
1339            # Store the change
1340            self.metadata_changes[key] = (GGUFValueType.ARRAY, (element_type, new_values))
1341            self.modified = True
1342
1343            # Update display
1344            enum_type = self.get_enum_for_key(key)
1345            if enum_type is not None and element_type == GGUFValueType.INT32:
1346                value_str = f"[ {', '.join(self.format_enum_value(v, enum_type) for v in new_values[:5])}{', ...' if len(new_values) > 5 else ''} ]"
1347            else:
1348                value_str = f"[ {', '.join(str(v) for v in new_values[:5])}{', ...' if len(new_values) > 5 else ''} ]"
1349            target_item = self.metadata_table.item(row, 2)
1350            if target_item:
1351                target_item.setText(value_str)
1352
1353            self.statusBar().showMessage(f"Updated array values for {key}")
1354
1355    def edit_tokenizer_metadata(self, trigger_key):
1356        """Edit the linked tokenizer metadata arrays together."""
1357        if not self.reader:
1358            return
1359
1360        # Get all three fields
1361        tokens_field = self.reader.get_field(gguf.Keys.Tokenizer.LIST)
1362        token_types_field = self.reader.get_field(gguf.Keys.Tokenizer.TOKEN_TYPE)
1363        scores_field = self.reader.get_field(gguf.Keys.Tokenizer.SCORES)
1364
1365        # Extract values from each field
1366        tokens = self.extract_array_values(tokens_field) if tokens_field else []
1367        token_types = self.extract_array_values(token_types_field) if token_types_field else []
1368        scores = self.extract_array_values(scores_field) if scores_field else []
1369
1370        # Apply any pending changes
1371        if gguf.Keys.Tokenizer.LIST in self.metadata_changes:
1372            _, (_, tokens) = self.metadata_changes[gguf.Keys.Tokenizer.LIST]
1373        if gguf.Keys.Tokenizer.TOKEN_TYPE in self.metadata_changes:
1374            _, (_, token_types) = self.metadata_changes[gguf.Keys.Tokenizer.TOKEN_TYPE]
1375        if gguf.Keys.Tokenizer.SCORES in self.metadata_changes:
1376            _, (_, scores) = self.metadata_changes[gguf.Keys.Tokenizer.SCORES]
1377
1378        # Open the tokenizer editor dialog
1379        dialog = TokenizerEditorDialog(tokens, token_types, scores, self)
1380        if dialog.exec() == QDialog.DialogCode.Accepted:
1381            new_tokens, new_token_types, new_scores = dialog.get_data()
1382
1383            # Store changes for all three arrays
1384            if tokens_field:
1385                self.metadata_changes[gguf.Keys.Tokenizer.LIST] = (
1386                    GGUFValueType.ARRAY,
1387                    (tokens_field.types[1], new_tokens)
1388                )
1389
1390            if token_types_field:
1391                self.metadata_changes[gguf.Keys.Tokenizer.TOKEN_TYPE] = (
1392                    GGUFValueType.ARRAY,
1393                    (token_types_field.types[1], new_token_types)
1394                )
1395
1396            if scores_field:
1397                self.metadata_changes[gguf.Keys.Tokenizer.SCORES] = (
1398                    GGUFValueType.ARRAY,
1399                    (scores_field.types[1], new_scores)
1400                )
1401
1402            self.modified = True
1403
1404            # Update display for all three fields
1405            self.update_tokenizer_display(gguf.Keys.Tokenizer.LIST, new_tokens)
1406            self.update_tokenizer_display(gguf.Keys.Tokenizer.TOKEN_TYPE, new_token_types)
1407            self.update_tokenizer_display(gguf.Keys.Tokenizer.SCORES, new_scores)
1408
1409            self.statusBar().showMessage("Updated tokenizer data")
1410
1411    def update_tokenizer_display(self, key, values):
1412        """Update the display of a tokenizer field in the metadata table."""
1413        for row in range(self.metadata_table.rowCount()):
1414            key_item = self.metadata_table.item(row, 0)
1415            if key_item and key_item.text() == key:
1416                value_str = f"[ {', '.join(str(v) for v in values[:5])}{', ...' if len(values) > 5 else ''} ]"
1417                value_item = self.metadata_table.item(row, 2)
1418                if value_item:
1419                    value_item.setText(value_str)
1420                break
1421
1422    def add_metadata(self):
1423        dialog = AddMetadataDialog(self)
1424        if dialog.exec() == QDialog.DialogCode.Accepted:
1425            key, value_type, value = dialog.get_data()
1426
1427            if not key:
1428                QMessageBox.warning(self, "Invalid Key", "Key cannot be empty")
1429                return
1430
1431            # Check if key already exists
1432            for row in range(self.metadata_table.rowCount()):
1433                orig_item = self.metadata_table.item(row, 0)
1434                if orig_item and orig_item.text() == key:
1435                    QMessageBox.warning(self, "Duplicate Key", f"Key '{key}' already exists")
1436                    return
1437
1438            # Add to table
1439            row = self.metadata_table.rowCount()
1440            self.metadata_table.insertRow(row)
1441
1442            # Key
1443            key_item = QTableWidgetItem(key)
1444            key_item.setFlags(key_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
1445            self.metadata_table.setItem(row, 0, key_item)
1446
1447            # Type
1448            type_item = QTableWidgetItem(value_type.name)
1449            type_item.setFlags(type_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
1450            self.metadata_table.setItem(row, 1, type_item)
1451
1452            # Value
1453            value_item = QTableWidgetItem(str(value))
1454            value_item.setFlags(value_item.flags() | Qt.ItemFlag.ItemIsEditable)
1455            self.metadata_table.setItem(row, 2, value_item)
1456
1457            # Actions
1458            actions_widget = QWidget()
1459            actions_layout = QHBoxLayout(actions_widget)
1460            actions_layout.setContentsMargins(2, 2, 2, 2)
1461
1462            remove_button = QPushButton("Remove")
1463            remove_button.setProperty("row", row)
1464            remove_button.setProperty("key", key)
1465            remove_button.clicked.connect(self.remove_metadata)
1466            actions_layout.addWidget(remove_button)
1467
1468            self.metadata_table.setCellWidget(row, 3, actions_widget)
1469
1470            # Store the change
1471            self.metadata_changes[key] = (value_type, value)
1472            self.modified = True
1473
1474            self.statusBar().showMessage(f"Added new metadata key {key}")
1475
1476    def save_file(self):
1477        if not self.reader:
1478            QMessageBox.warning(self, "No File Open", "Please open a GGUF file first")
1479            return
1480
1481        if not self.modified and not self.metadata_changes and not self.metadata_to_remove:
1482            QMessageBox.information(self, "No Changes", "No changes to save")
1483            return
1484
1485        file_path, _ = QFileDialog.getSaveFileName(
1486            self, "Save GGUF File As", "", "GGUF Files (*.gguf);;All Files (*)"
1487        )
1488
1489        if not file_path:
1490            return
1491
1492        try:
1493            self.statusBar().showMessage(f"Saving to {file_path}...")
1494            QApplication.processEvents()
1495
1496            # Get architecture and endianness from the original file
1497            arch = 'unknown'
1498            field = self.reader.get_field(gguf.Keys.General.ARCHITECTURE)
1499            if field:
1500                arch = field.contents()
1501
1502            # Create writer
1503            writer = GGUFWriter(file_path, arch=arch, endianess=self.reader.endianess)
1504
1505            # Get alignment if present
1506            alignment = None
1507            field = self.reader.get_field(gguf.Keys.General.ALIGNMENT)
1508            if field:
1509                alignment = field.contents()
1510                if alignment is not None:
1511                    writer.data_alignment = alignment
1512
1513            # Copy metadata with changes
1514            for field in self.reader.fields.values():
1515                # Skip virtual fields and fields written by GGUFWriter
1516                if field.name == gguf.Keys.General.ARCHITECTURE or field.name.startswith('GGUF.'):
1517                    continue
1518
1519                # Skip fields marked for removal
1520                if field.name in self.metadata_to_remove:
1521                    continue
1522
1523                # Apply changes if any
1524                sub_type = None
1525                if field.name in self.metadata_changes:
1526                    value_type, value = self.metadata_changes[field.name]
1527                    if value_type == GGUFValueType.ARRAY:
1528                        # Handle array values
1529                        sub_type, value = value
1530                else:
1531                    # Copy original value
1532                    value = field.contents()
1533                    value_type = field.types[0]
1534                    if value_type == GGUFValueType.ARRAY:
1535                        sub_type = field.types[-1]
1536
1537                if value is not None:
1538                    writer.add_key_value(field.name, value, value_type, sub_type=sub_type)
1539
1540            # Add new metadata
1541            for key, (value_type, value) in self.metadata_changes.items():
1542                # Skip if the key already existed (we handled it above)
1543                if self.reader.get_field(key) is not None:
1544                    continue
1545
1546                sub_type = None
1547                if value_type == GGUFValueType.ARRAY:
1548                    # Handle array values
1549                    sub_type, value = value
1550
1551                writer.add_key_value(key, value, value_type, sub_type=sub_type)
1552
1553            # Add tensors (including data)
1554            for tensor in self.reader.tensors:
1555                writer.add_tensor(tensor.name, tensor.data, raw_shape=tensor.data.shape, raw_dtype=tensor.tensor_type, tensor_endianess=self.reader.endianess)
1556
1557            # Write header and metadata
1558            writer.open_output_file(Path(file_path))
1559            writer.write_header_to_file()
1560            writer.write_kv_data_to_file()
1561
1562            # Write tensor data using the optimized method
1563            writer.write_tensors_to_file(progress=False)
1564
1565            writer.close()
1566
1567            self.statusBar().showMessage(f"Saved to {file_path}")
1568
1569            # Ask if user wants to open the new file
1570            reply = QMessageBox.question(
1571                self, "Open Saved File",
1572                "Would you like to open the newly saved file?",
1573                QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, QMessageBox.StandardButton.Yes
1574            )
1575
1576            if reply == QMessageBox.StandardButton.Yes:
1577                self.reader = GGUFReader(file_path, 'r')
1578                self.current_file = file_path
1579                self.file_path_edit.setText(file_path)
1580
1581                self.load_metadata()
1582                self.load_tensors()
1583
1584                self.metadata_changes = {}
1585                self.metadata_to_remove = set()
1586                self.modified = False
1587
1588        except Exception as e:
1589            QMessageBox.critical(self, "Error", f"Failed to save file: {str(e)}")
1590            self.statusBar().showMessage("Error saving file")
1591
1592
1593def main() -> None:
1594    parser = argparse.ArgumentParser(description="GUI GGUF Editor")
1595    parser.add_argument("model_path", nargs="?", help="path to GGUF model file to load at startup")
1596    parser.add_argument("--verbose", action="store_true", help="increase output verbosity")
1597
1598    args = parser.parse_args()
1599
1600    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO)
1601
1602    app = QApplication(sys.argv)
1603    window = GGUFEditorWindow()
1604    window.show()
1605
1606    # Load model if specified
1607    if args.model_path:
1608        if os.path.isfile(args.model_path) and args.model_path.endswith('.gguf'):
1609            window.load_file(args.model_path)
1610        else:
1611            logger.error(f"Invalid model path: {args.model_path}")
1612            QMessageBox.warning(
1613                window,
1614                "Invalid Model Path",
1615                f"The specified file does not exist or is not a GGUF file: {args.model_path}")
1616
1617    sys.exit(app.exec())
1618
1619
1620if __name__ == '__main__':
1621    main()