\n{document.text}\n"
- if document.kind == "conversation":
- return f".*?)```", re.DOTALL)
- last = 0
- formatter = HtmlFormatter(style="monokai", noclasses=True, nowrap=True)
- for match in pattern.finditer(markdown_text):
- prose = markdown_text[last:match.start()]
- if prose.strip():
- body_parts.append(markdown_lib.markdown(prose, extensions=["tables", "nl2br"]))
- code = match.group("code")
- lang = match.group("lang").strip()
- try:
- lexer = get_lexer_by_name(lang) if lang else guess_lexer(code)
- except Exception:
- lexer = TextLexer()
- highlighted = highlight(code, lexer, formatter)
- label = lang.title() if lang else lexer.name
- body_parts.append(code_block_html(label, highlighted, code))
- last = match.end()
- prose = markdown_text[last:]
- if prose.strip():
- body_parts.append(markdown_lib.markdown(prose, extensions=["tables", "nl2br"]))
- body = "\n".join(body_parts) if body_parts else ""
- return (
- f"""
-
-
-
-
- {body}
-
- """
- )
- except Exception:
- return basic_markdown_html(markdown_text)
-
-def basic_markdown_html(markdown_text: str) -> str:
- """Render basic Markdown with simple code coloring.
-
- Args:
- markdown_text: Raw Markdown.
-
- Returns:
- Basic HTML.
- """
-
- text = normalize_code_blocks(markdown_text)
- parts: list[str] = []
- pattern = re.compile(r"```(?:\w+)?\n(.*?)```", re.DOTALL)
- last = 0
- for match in pattern.finditer(text):
- parts.append(render_basic_prose(text[last:match.start()]))
- code = match.group(1)
- parts.append(code_block_html("Code", colorize_code(code), code))
- last = match.end()
- parts.append(render_basic_prose(text[last:]))
- return (
- ""
- ""
- + "".join(parts)
- + ""
- )
-
-def code_block_html(label: str, highlighted_html: str, raw_code: str) -> str:
- """Build a code panel with a copy link.
-
- Args:
- label: Code language label.
- highlighted_html: Highlighted code HTML.
- raw_code: Raw code for clipboard copy.
-
- Returns:
- Code panel HTML.
- """
-
- return (
- ""
- f""
- f"{highlighted_html}
"
- ""
- )
-
-def render_basic_prose(text: str) -> str:
- """Render a small Markdown subset for fallback mode.
-
- Args:
- text: Markdown prose.
-
- Returns:
- HTML fragment.
- """
-
- html_lines: list[str] = []
- in_ordered = False
- in_unordered = False
-
- def close_lists() -> None:
- nonlocal in_ordered, in_unordered
- if in_ordered:
- html_lines.append("")
- in_ordered = False
- if in_unordered:
- html_lines.append("")
- in_unordered = False
-
- for raw_line in text.splitlines():
- line = raw_line.strip()
- if not line:
- close_lists()
- html_lines.append("
")
- continue
- if line.startswith("### "):
- close_lists()
- html_lines.append(f"{inline_basic_markdown(line[4:])}
")
- continue
- if line.startswith("## "):
- close_lists()
- html_lines.append(f"{inline_basic_markdown(line[3:])}
")
- continue
- if line.startswith("# "):
- close_lists()
- html_lines.append(f"{inline_basic_markdown(line[2:])}
")
- continue
- ordered = re.match(r"^\d+\.\s+(.*)$", line)
- if ordered:
- if not in_ordered:
- close_lists()
- html_lines.append("")
- in_ordered = True
- html_lines.append(f"- {inline_basic_markdown(ordered.group(1))}
")
- continue
- unordered = re.match(r"^[-*]\s+(.*)$", line)
- if unordered:
- if not in_unordered:
- close_lists()
- html_lines.append("")
- in_unordered = True
- html_lines.append(f"- {inline_basic_markdown(unordered.group(1))}
")
- continue
- close_lists()
- html_lines.append(f"{inline_basic_markdown(line)}
")
- close_lists()
- return "\n".join(html_lines)
-
-def inline_basic_markdown(text: str) -> str:
- """Render inline Markdown for fallback mode.
-
- Args:
- text: Inline Markdown text.
-
- Returns:
- HTML fragment.
- """
-
- escaped = escape_html(text)
- escaped = re.sub(r"`([^`]+)`", r"\1", escaped)
- escaped = re.sub(r"\*\*(.+?)\*\*", r"\1", escaped)
- escaped = re.sub(r"\*(.+?)\*", r"\1", escaped)
- return escaped
-
-def escape_html(text: str) -> str:
- """Escape text for HTML.
-
- Args:
- text: Raw text.
-
- Returns:
- Escaped text.
- """
-
- return text.replace("&", "&").replace("<", "<").replace(">", ">")
-
-def colorize_code(code: str) -> str:
- """Apply simple inline colors to Python-like code.
-
- Args:
- code: Source code.
-
- Returns:
- HTML code.
- """
-
- escaped = escape_html(code)
- keywords = {
- "def", "class", "import", "from", "for", "while", "if", "else", "elif",
- "try", "except", "return", "print", "with", "as", "in", "function", "const",
- "let", "var", "new", "typeof", "await", "async", "true", "false", "null",
- "True", "False", "None",
- }
- builtins = {"console", "Object", "process", "JSON", "Array", "String", "Number", "Boolean", "Math", "os", "sys"}
- token_pattern = re.compile(
- r"(?P//.*|#.*)"
- r"|(?P`(?:\\.|[^`])*`|'(?:\\.|[^'])*'|\"(?:\\.|[^\"])*\")"
- r"|(?P\b\d+(?:\.\d+)?\b)"
- r"|(?P\b[A-Za-z_][A-Za-z0-9_]*\b)"
- )
- colored_lines: list[str] = []
- for line in escaped.splitlines():
- segments: list[str] = []
- last = 0
- for match in token_pattern.finditer(line):
- segments.append(line[last:match.start()])
- value = match.group(0)
- if match.lastgroup == "comment":
- segments.append(f"{value}")
- elif match.lastgroup == "string":
- segments.append(f"{value}")
- elif match.lastgroup == "number":
- segments.append(f"{value}")
- elif match.lastgroup == "word":
- next_chars = line[match.end(): match.end() + 2]
- previous = line[max(0, match.start() - 1): match.start()]
- if value in keywords:
- segments.append(f"{value}")
- elif value in builtins:
- segments.append(f"{value}")
- elif next_chars.startswith("(") and previous != ".":
- segments.append(f"{value}")
- elif previous == ".":
- segments.append(f"{value}")
- else:
- segments.append(value)
- last = match.end()
- segments.append(line[last:])
- colored_lines.append("".join(segments))
- return "\n".join(colored_lines)
-
-def normalize_code_blocks(markdown_text: str) -> str:
- """Fence obvious loose code blocks so syntax highlighting can run.
-
- Args:
- markdown_text: Raw model Markdown.
-
- Returns:
- Markdown with likely code blocks fenced.
- """
-
- if "```" in markdown_text:
- if markdown_text.count("```") % 2:
- return f"{markdown_text}\n```"
- return markdown_text
- lines = markdown_text.splitlines()
- normalized: list[str] = []
- code_block: list[str] = []
-
- def is_code_line(line: str) -> bool:
- stripped = line.strip()
- if not stripped:
- return bool(code_block)
- if line.startswith((" ", "\t")):
- return True
- if re.match(
- r"^(def|class|import|from|for|while|if|else:?|elif|try:?|except|return|print|with|"
- r"function|const|let|var|console\.|Object\.|process\.)\b",
- stripped,
- ):
- return True
- if stripped in {"{", "}", "};", "})", "});"}:
- return True
- if stripped.startswith(("#", "@")):
- return True
- return sum(stripped.count(symbol) for symbol in "()[]{}:=<>+-*/") >= 3
-
- def flush() -> None:
- nonlocal code_block
- if len([line for line in code_block if line.strip()]) >= 3:
- normalized.append(f"```{guess_code_language(code_block)}")
- normalized.extend(code_block)
- normalized.append("```")
- else:
- normalized.extend(code_block)
- code_block = []
-
- for line in lines:
- if is_code_line(line):
- code_block.append(line)
- else:
- flush()
- normalized.append(line)
- flush()
- return "\n".join(normalized)
-
-def guess_code_language(lines: list[str]) -> str:
- """Guess a fence language for loose code.
-
- Args:
- lines: Code lines.
-
- Returns:
- Markdown fence language.
- """
-
- joined = "\n".join(lines).lower()
- if any(marker in joined for marker in ("console.", "const ", "let ", "function ", "process.env", "object.keys")):
- return "javascript"
- if any(marker in joined for marker in ("#include", "std::", "cout", "cin")):
- return "cpp"
- if any(marker in joined for marker in ("public class", "system.out", "private ", "protected ")):
- return "java"
- if any(marker in joined for marker in ("def ", "import ", "print(", "self.")):
- return "python"
- return "text"
diff --git a/llm_trainer/ui/micro_llm_creator_lightning.ico b/llm_trainer/ui/micro_llm_creator_lightning.ico
deleted file mode 100644
index ee1de17..0000000
Binary files a/llm_trainer/ui/micro_llm_creator_lightning.ico and /dev/null differ
diff --git a/llm_trainer/ui/startup.py b/llm_trainer/ui/startup.py
deleted file mode 100644
index e08fab7..0000000
--- a/llm_trainer/ui/startup.py
+++ /dev/null
@@ -1,461 +0,0 @@
-īģŋfrom __future__ import annotations
-
-"""Startup validation and project selection UI."""
-import ctypes
-from datetime import datetime
-import html
-import json
-import logging
-import os
-from pathlib import Path
-import sys
-from typing import Optional
-
-from PySide6.QtCore import QEvent, Qt
-from PySide6.QtGui import QFont, QFontDatabase
-from PySide6.QtWidgets import (
- QApplication,
- QDialog,
- QHBoxLayout,
- QLabel,
- QListWidget,
- QListWidgetItem,
- QPushButton,
- QProgressBar,
- QTextBrowser,
- QVBoxLayout,
- QWidget,
-)
-
-from .startup_validation import (
- _run_startup_tests,
- _run_startup_validations,
- _validate_writable_directory,
-)
-
-
-APP_NAME = "DrunkenBot LLM-IDE"
-WINDOWS_APP_ID = "DrunkenBot.LLMIDE"
-LOGGER = logging.getLogger(__name__)
-APP_HOME_DIR = Path.home() / ".drunkenbot_ide"
-DEFAULT_CACHE_DIR = APP_HOME_DIR / "cache"
-DEFAULT_PROJECTS_DIR = APP_HOME_DIR / "projects"
-RECENT_PROJECTS_PATH = APP_HOME_DIR / "recent_projects.json"
-_WINDOWS_ICON_HANDLES: list[int] = []
-_LOGO_FONT_FAMILY: Optional[str] = None
-
-
-def _main_window():
- """Load the main window lazily to avoid a startup-module import cycle."""
- from .app import MainWindow
-
- return MainWindow
-
-def _load_recent_projects(limit: int = 12) -> list[Path]:
- """Return recently opened project files that still exist."""
-
- try:
- payload = json.loads(RECENT_PROJECTS_PATH.read_text(encoding="utf-8"))
- except Exception:
- return []
- if not isinstance(payload, list):
- return []
- results: list[Path] = []
- for item in payload:
- if not isinstance(item, dict):
- continue
- path_text = str(item.get("path", "")).strip()
- if not path_text:
- continue
- path = Path(path_text)
- if path.exists() and path.is_file():
- results.append(path)
- if len(results) >= limit:
- break
- return results
-
-
-def _register_recent_project(project_file: Path, limit: int = 12) -> None:
- """Insert/update a project file in recent history."""
-
- APP_HOME_DIR.mkdir(parents=True, exist_ok=True)
- now = datetime.utcnow().isoformat() + "Z"
- try:
- payload = json.loads(RECENT_PROJECTS_PATH.read_text(encoding="utf-8"))
- except Exception:
- payload = []
- rows: list[dict[str, str]] = []
- resolved_new = project_file.resolve()
- for item in payload if isinstance(payload, list) else []:
- if not isinstance(item, dict):
- continue
- path_text = str(item.get("path", "")).strip()
- if not path_text:
- continue
- path = Path(path_text)
- if not path.exists() or not path.is_file():
- continue
- if path.resolve() == resolved_new:
- continue
- rows.append(
- {
- "path": str(path),
- "last_opened": str(item.get("last_opened", now)),
- }
- )
- rows.insert(0, {"path": str(project_file), "last_opened": now})
- RECENT_PROJECTS_PATH.write_text(json.dumps(rows[:limit], indent=2), encoding="utf-8")
-
-
-def _apply_windows_taskbar_icon(widget: QWidget) -> None:
- """Apply the app icon to a Qt widget taskbar entry on Windows."""
-
- if sys.platform != "win32":
- return
- try:
- ctypes.windll.shell32.SetCurrentProcessExplicitAppUserModelID(WINDOWS_APP_ID)
- except Exception:
- LOGGER.exception("Could not set Windows app user model ID for widget")
- icon_path = _main_window()._ensure_windows_icon_file()
- if icon_path is None:
- return
- hwnd = int(widget.winId())
- if not hwnd:
- return
- wm_seticon = 0x0080
- icon_small = 0
- icon_big = 1
- image_icon = 1
- lr_loadfromfile = 0x0010
- user32 = ctypes.windll.user32
- hicon_big = user32.LoadImageW(None, str(icon_path), image_icon, 256, 256, lr_loadfromfile)
- hicon_small = user32.LoadImageW(None, str(icon_path), image_icon, 32, 32, lr_loadfromfile)
- if hicon_big:
- user32.SendMessageW(hwnd, wm_seticon, icon_big, hicon_big)
- _WINDOWS_ICON_HANDLES.append(hicon_big)
- if hicon_small:
- user32.SendMessageW(hwnd, wm_seticon, icon_small, hicon_small)
- _WINDOWS_ICON_HANDLES.append(hicon_small)
-
-
-def _logo_font_family() -> Optional[str]:
- """Load and cache the custom logo font family when available."""
-
- global _LOGO_FONT_FAMILY
- if _LOGO_FONT_FAMILY is not None:
- return _LOGO_FONT_FAMILY
- font_path = Path(__file__).resolve().parents[2] / "fonts" / "Blue-Whale Heavy.otf"
- if not font_path.exists():
- _LOGO_FONT_FAMILY = ""
- return None
- font_id = QFontDatabase.addApplicationFont(str(font_path))
- if font_id < 0:
- _LOGO_FONT_FAMILY = ""
- return None
- families = QFontDatabase.applicationFontFamilies(font_id)
- if not families:
- _LOGO_FONT_FAMILY = ""
- return None
- _LOGO_FONT_FAMILY = families[0]
- return _LOGO_FONT_FAMILY
-
-
-class StartupValidationSplash(QDialog):
- """Modal splash screen that shows startup validation progress."""
-
- def __init__(self) -> None:
- super().__init__()
- self.setWindowTitle(APP_NAME)
- self.setWindowFlags(Qt.Window | Qt.FramelessWindowHint | Qt.WindowStaysOnTopHint)
- self.setModal(True)
- self.setMinimumSize(560, 760)
- self.setFont(QFont("Arial", 10))
- self._checks: dict[str, str] = {}
- self._check_order: list[str] = []
- self._build_ui()
-
- def _build_ui(self) -> None:
- self.setStyleSheet(
- """
- QDialog { background: #111111; color: #d0d0d0; border: 0; border-radius: 0; font-family: Arial, "Segoe UI", sans-serif; }
- QLabel#Title { color: #d0d0d0; font-size: 22px; }
- QLabel#Subtitle { color: #bfbfbf; font-size: 13px; }
- QLabel#Step { color: #c7c7c7; font-size: 13px; }
- QTextBrowser { background: #111111; color: #d0d0d0; border: 0; padding: 10px; }
- QProgressBar { background: #222222; border: 0; border-radius: 2px; }
- QProgressBar::chunk { background: #bcbcbc; border-radius: 2px; }
- """
- )
- root = QVBoxLayout(self)
- root.setContentsMargins(24, 24, 24, 24)
- root.setSpacing(12)
-
- header = QHBoxLayout()
- logo = QLabel()
- logo.setFixedSize(128, 128)
- logo_pixmap = _main_window()._app_logo_pixmap(118)
- if logo_pixmap.isNull():
- logo.setText("DB")
- logo.setAlignment(Qt.AlignCenter)
- logo.setStyleSheet("color:#f5b041;font-size:38px;")
- else:
- logo.setPixmap(logo_pixmap)
- logo.setAlignment(Qt.AlignCenter)
- title_box = QVBoxLayout()
- title = QLabel(APP_NAME)
- title.setObjectName("Title")
- logo_family = _logo_font_family()
- if logo_family:
- title.setFont(QFont(logo_family, 22))
- title_box.addWidget(title)
- title_box.addSpacing(4)
- header.addWidget(logo)
- header.addSpacing(10)
- header.addLayout(title_box, 1)
- root.addLayout(header)
-
- self.step_label = QLabel("Preparing checks...")
- self.step_label.setObjectName("Step")
- root.addWidget(self.step_label)
-
- self.progress = QProgressBar()
- self.progress.setRange(0, 100)
- self.progress.setTextVisible(False)
- self.progress.setFixedHeight(4)
- self.progress.setValue(0)
- root.addWidget(self.progress)
-
- self.checks_view = QTextBrowser()
- self.checks_view.setOpenExternalLinks(False)
- self.checks_view.setReadOnly(True)
- root.addWidget(self.checks_view, 1)
- self.footer_label = QLabel("")
- self.footer_label.setObjectName("Subtitle")
- root.addWidget(self.footer_label)
-
- def update_step(self, text: str, index: int, total: int) -> None:
- self.step_label.setText(text)
- percent = int((max(0, index) / max(1, total)) * 100)
- self.progress.setValue(percent)
- QApplication.processEvents()
-
- def set_checks(self, checks: list[str]) -> None:
- """Initialize the checklist in pending state."""
-
- self._check_order = list(checks)
- self._checks = {label: "pending" for label in checks}
- self._render_checks()
-
- def add_check(self, label: str) -> None:
- """Add a dynamically discovered check to the startup checklist.
-
- Args:
- label: Human-readable check name.
- """
- if label in self._checks:
- self._checks[label] = "running"
- self._render_checks()
- return
- self._check_order.append(label)
- self._checks[label] = "running"
- self._render_checks()
-
- def mark_check_running(self, label: str) -> None:
- self._checks[label] = "running"
- self._render_checks()
-
- def mark_check_done(self, label: str) -> None:
- self._checks[label] = "done"
- self._render_checks()
-
- def mark_check_failed(self, label: str) -> None:
- self._checks[label] = "failed"
- self._render_checks()
-
- def append_log(self, text: str) -> None:
- self.footer_label.setText(text)
- QApplication.processEvents()
-
- def showEvent(self, event: QEvent) -> None:
- super().showEvent(event)
- _apply_windows_taskbar_icon(self)
-
- def _render_checks(self) -> None:
- rows: list[str] = [""]
- for label in self._check_order:
- state = self._checks.get(label, "pending")
- escaped = html.escape(label)
- if state == "done":
- rows.append(f"- [OK] {escaped}
")
- elif state == "running":
- rows.append(f"- [*] {escaped}
")
- elif state == "failed":
- rows.append(f"- [FAIL] {escaped}
")
- else:
- rows.append(f"- - {escaped}
")
- rows.append("
")
- self.checks_view.setHtml("".join(rows))
- QApplication.processEvents()
-
-
-class ProjectChoiceDialog(QDialog):
- """Prompt shown after startup checks to choose project creation/open flow."""
-
- def __init__(self) -> None:
- super().__init__()
- self.choice = ""
- self.selected_project_file: Optional[Path] = None
- self.setWindowTitle(APP_NAME)
- self.setWindowFlags(Qt.Window | Qt.WindowCloseButtonHint)
- self.setModal(True)
- self.setMinimumSize(760, 520)
- self.setFont(QFont("Arial", 10))
- self._build_ui()
-
- def _build_ui(self) -> None:
- self.setStyleSheet(
- """
- QDialog { background: #111111; color: #eeeeee; border: 0; border-radius: 0; font-family: Arial, "Segoe UI", sans-serif; }
- QLabel#Title { color: #f5b041; font-size: 24px; }
- QLabel#Body { color: #dddddd; font-size: 13px; }
- QLabel#CardTitle { color: #f1f1f1; font-size: 16px; }
- QLabel#CardBody { color: #c9c9c9; font-size: 12px; }
- QWidget#ChoiceCard { background: #171717; border: 1px solid #3a3a3a; border-radius: 8px; }
- QListWidget { background: #171717; color: #d8d8d8; border: 1px solid #3a3a3a; border-radius: 8px; padding: 4px; }
- QListWidget::item { padding: 6px 8px; }
- QListWidget::item:selected { background: #2a2a2a; color: #ffffff; }
- QPushButton { background: #242424; color: #eeeeee; border: 0; border-radius: 6px; padding: 8px 12px; }
- QPushButton:hover { background: #f5b041; color: #151515; }
- """
- )
- root = QVBoxLayout(self)
- root.setContentsMargins(28, 24, 28, 24)
- root.setSpacing(16)
-
- logo = QLabel()
- logo_pixmap = _main_window()._app_logo_pixmap(144)
- if logo_pixmap.isNull():
- logo.setText("DB")
- logo.setStyleSheet("color:#f5b041;font-size:56px;")
- logo.setAlignment(Qt.AlignCenter)
- else:
- logo.setPixmap(logo_pixmap)
- logo.setAlignment(Qt.AlignCenter)
- root.addWidget(logo, 0, Qt.AlignHCenter)
-
- title = QLabel("Get started")
- title.setObjectName("Title")
- logo_family = _logo_font_family()
- if logo_family:
- title.setFont(QFont(logo_family, 26))
- title.setAlignment(Qt.AlignLeft)
- root.addWidget(title)
-
- body = QLabel(
- "Startup checks are complete.\n"
- "Choose how you want to begin with DrunkenBot LLM-IDE."
- )
- body.setObjectName("Body")
- body.setAlignment(Qt.AlignLeft)
- root.addWidget(body)
-
- new_card = QWidget()
- new_card.setObjectName("ChoiceCard")
- new_layout = QVBoxLayout(new_card)
- new_layout.setContentsMargins(16, 14, 16, 14)
- new_layout.setSpacing(8)
- new_title = QLabel("Create a new project")
- new_title.setObjectName("CardTitle")
- new_body = QLabel("Start with a clean workspace, default folders, and bundled starter data.")
- new_body.setObjectName("CardBody")
- new_body.setWordWrap(True)
- new_button = QPushButton("Create New Project")
- new_layout.addWidget(new_title)
- new_layout.addWidget(new_body)
- new_layout.addWidget(new_button, 0, Qt.AlignLeft)
- root.addWidget(new_card)
-
- open_card = QWidget()
- open_card.setObjectName("ChoiceCard")
- open_layout = QVBoxLayout(open_card)
- open_layout.setContentsMargins(16, 14, 16, 14)
- open_layout.setSpacing(8)
- open_title = QLabel("Open an existing project")
- open_title.setObjectName("CardTitle")
- open_body = QLabel("Open a saved project.json and continue where you left off.")
- open_body.setObjectName("CardBody")
- open_body.setWordWrap(True)
- open_button = QPushButton("Open Existing Project")
- open_layout.addWidget(open_title)
- open_layout.addWidget(open_body)
- open_layout.addWidget(open_button, 0, Qt.AlignLeft)
- root.addWidget(open_card)
-
- test_chat_card = QWidget()
- test_chat_card.setObjectName("ChoiceCard")
- test_chat_layout = QVBoxLayout(test_chat_card)
- test_chat_layout.setContentsMargins(16, 14, 16, 14)
- test_chat_layout.setSpacing(8)
- test_chat_title = QLabel("Test local LLM")
- test_chat_title.setObjectName("CardTitle")
- test_chat_body = QLabel("Jump directly to the Chat tab to load a local model and start chatting.")
- test_chat_body.setObjectName("CardBody")
- test_chat_body.setWordWrap(True)
- test_chat_button = QPushButton("Test Local LLM")
- test_chat_layout.addWidget(test_chat_title)
- test_chat_layout.addWidget(test_chat_body)
- test_chat_layout.addWidget(test_chat_button, 0, Qt.AlignLeft)
- root.addWidget(test_chat_card)
-
- recent_paths = _load_recent_projects()
- self.recent_list: Optional[QListWidget] = None
- if recent_paths:
- recent_card = QWidget()
- recent_card.setObjectName("ChoiceCard")
- recent_layout = QVBoxLayout(recent_card)
- recent_layout.setContentsMargins(16, 14, 16, 14)
- recent_layout.setSpacing(8)
- recent_title = QLabel("Recent projects")
- recent_title.setObjectName("CardTitle")
- recent_layout.addWidget(recent_title)
- self.recent_list = QListWidget()
- for path in recent_paths:
- item = QListWidgetItem(str(path))
- item.setData(Qt.UserRole, str(path))
- self.recent_list.addItem(item)
- self.recent_list.setCurrentRow(0)
- recent_layout.addWidget(self.recent_list)
- recent_button = QPushButton("Open Selected Recent Project")
- recent_button.clicked.connect(self._open_selected_recent)
- recent_layout.addWidget(recent_button, 0, Qt.AlignLeft)
- root.addWidget(recent_card)
-
- row = QHBoxLayout()
- row.addStretch(1)
- exit_button = QPushButton("Exit")
- new_button.clicked.connect(lambda: self._choose("new"))
- open_button.clicked.connect(lambda: self._choose("open"))
- test_chat_button.clicked.connect(lambda: self._choose("test_local_llm"))
- exit_button.clicked.connect(self.reject)
- row.addWidget(exit_button)
- root.addLayout(row)
-
- def _choose(self, choice: str) -> None:
- self.choice = choice
- self.accept()
-
- def _open_selected_recent(self) -> None:
- if self.recent_list is None:
- return
- item = self.recent_list.currentItem()
- if item is None:
- return
- raw = item.data(Qt.UserRole)
- if not raw:
- return
- self.selected_project_file = Path(str(raw))
- self._choose("recent")
-
- def showEvent(self, event: QEvent) -> None:
- super().showEvent(event)
- _apply_windows_taskbar_icon(self)
diff --git a/llm_trainer/ui/startup_splash.py b/llm_trainer/ui/startup_splash.py
deleted file mode 100644
index 80d1780..0000000
--- a/llm_trainer/ui/startup_splash.py
+++ /dev/null
@@ -1,127 +0,0 @@
-from __future__ import annotations
-
-import html
-import os
-import sys
-from pathlib import Path
-
-from PySide6.QtCore import QEvent, Qt
-from PySide6.QtGui import QFont, QPixmap
-from PySide6.QtWidgets import QDialog, QLabel, QHBoxLayout, QProgressBar, QTextBrowser, QVBoxLayout
-
-
-class StartupSplash(QDialog):
- """Splash screen shown while startup validation is running."""
-
- def __init__(self) -> None:
- super().__init__()
- self.setWindowTitle("DrunkenBot LLM-IDE")
- self.setWindowFlags(Qt.Window | Qt.FramelessWindowHint | Qt.WindowStaysOnTopHint)
- self.setModal(True)
- self.setMinimumSize(560, 760)
- layout = QVBoxLayout(self)
- layout.setContentsMargins(24, 24, 24, 24)
- layout.setSpacing(12)
- header = QHBoxLayout()
- logo = QLabel()
- logo.setFixedSize(128, 128)
- logo_candidates = [
- Path(__file__).resolve().parents[2] / "drunken_bot_logo_small.png",
- Path(__file__).resolve().parents[3] / "drunken_bot_logo_small.png",
- ]
- if hasattr(sys, "_MEIPASS"):
- logo_candidates.insert(0, Path(sys._MEIPASS) / "drunken_bot_logo_small.png")
- app_root = os.environ.get("DRUNKENBOT_APP_ROOT")
- if app_root:
- logo_candidates.insert(0, Path(app_root) / "drunken_bot_logo_small.png")
- pixmap = QPixmap(str(next((path for path in logo_candidates if path.exists()), logo_candidates[0])))
- if pixmap.isNull():
- logo.setText("DB")
- logo.setAlignment(Qt.AlignCenter)
- logo.setStyleSheet("color:#f5b041;font-size:38px;")
- else:
- logo.setPixmap(pixmap.scaled(118, 118, Qt.KeepAspectRatio, Qt.SmoothTransformation))
- logo.setAlignment(Qt.AlignCenter)
- title = QLabel("DrunkenBot LLM-IDE")
- title.setObjectName("Title")
- title.setFont(QFont("Arial", 22))
- header.addWidget(logo)
- header.addSpacing(10)
- header.addWidget(title, 1)
- layout.addLayout(header)
- self.status = QLabel("Preparing checks...")
- self.status.setObjectName("Step")
- layout.addWidget(self.status)
- self.progress = QProgressBar()
- self.progress.setRange(0, 100)
- self.progress.setTextVisible(False)
- self.progress.setFixedHeight(4)
- layout.addWidget(self.progress)
- self.checks_view = QTextBrowser()
- self.checks_view.setOpenExternalLinks(False)
- self.checks_view.setReadOnly(True)
- layout.addWidget(self.checks_view, 1)
- self.footer_label = QLabel("")
- layout.addWidget(self.footer_label)
- self.setStyleSheet(
- "QDialog { background: #111111; color: #d0d0d0; border: 0; border-radius: 0; "
- 'font-family: Arial, "Segoe UI", sans-serif; } QLabel#Title { color: #d0d0d0; '
- "font-size: 22px; } QLabel#Step { color: #c7c7c7; font-size: 13px; } "
- "QTextBrowser { background: #111111; color: #d0d0d0; border: 0; padding: 10px; } "
- "QProgressBar { background: #222222; border: 0; border-radius: 2px; } "
- "QProgressBar::chunk { background: #bcbcbc; border-radius: 2px; }"
- )
- self._checks = {}
- self._check_order = []
-
- def append_log(self, message: str) -> None:
- self.footer_label.setText(message)
-
- def set_checks(self, checks: list[str]) -> None:
- self._check_order = list(checks)
- self._checks = {label: "pending" for label in checks}
- self._render_checks()
-
- def add_check(self, label: str) -> None:
- """Add a dynamically discovered check to the checklist.
-
- Args:
- label: Human-readable check name.
- """
- if label in self._checks:
- self._checks[label] = "running"
- self._render_checks()
- return
- self._check_order.append(label)
- self._checks[label] = "running"
- self._render_checks()
-
- def update_step(self, text: str, index: int, total: int) -> None:
- self.status.setText(text)
- self.progress.setValue(int(index / max(total, 1) * 100))
-
- def mark_check_running(self, label: str) -> None:
- self._checks[label] = "running"
- self._render_checks()
-
- def mark_check_done(self, label: str) -> None:
- self._checks[label] = "done"
- self._render_checks()
-
- def mark_check_failed(self, label: str) -> None:
- self._checks[label] = "failed"
- self._render_checks()
-
- def _render_checks(self) -> None:
- rows = [""]
- for label in self._check_order:
- state = self._checks.get(label, "pending")
- escaped = html.escape(label)
- marker = {"done": "[OK]", "running": "[*]", "failed": "[FAIL]"}.get(state, "-")
- color = {"done": "#ffffff", "running": "#e2cfaa", "failed": "#ff9a9a"}.get(state, "#bdbdbd")
- rows.append(f"- {marker} {escaped}
")
- rows.append("
")
- self.checks_view.setHtml("".join(rows))
-
- def showEvent(self, event: QEvent) -> None:
- super().showEvent(event)
diff --git a/llm_trainer/ui/startup_validation.py b/llm_trainer/ui/startup_validation.py
deleted file mode 100644
index ca36273..0000000
--- a/llm_trainer/ui/startup_validation.py
+++ /dev/null
@@ -1,120 +0,0 @@
-īģŋ"""Startup validation helpers and repository test execution."""
-
-from __future__ import annotations
-
-import ast
-import importlib
-import subprocess
-import sys
-from pathlib import Path
-from typing import Any, Optional
-
-from PySide6.QtWidgets import QApplication
-
-from engine.app_logging import DEFAULT_LOG_DIR
-
-APP_HOME_DIR = Path.home() / ".drunkenbot_ide"
-DEFAULT_CACHE_DIR = APP_HOME_DIR / "cache"
-DEFAULT_PROJECTS_DIR = APP_HOME_DIR / "projects"
-
-
-def _validate_writable_directory(path: Path) -> None:
- """Ensure a directory exists and can be written."""
- path.mkdir(parents=True, exist_ok=True)
- probe = path / ".startup_probe"
- probe.write_text("ok", encoding="utf-8")
- probe.unlink(missing_ok=True)
-
-
-def _test_display_name(label: str) -> str:
- """Extract the concise unittest method name from verbose output."""
- value = label.removeprefix("Test: ").strip()
- return value.split(" ", 1)[0].removesuffix("...")
-
-
-def _discover_test_labels(tests_root: Path) -> list[str]:
- """Discover unittest-style test method labels without executing tests."""
- labels: list[str] = []
- for path in sorted(tests_root.glob("test_*.py")):
- tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
- for node in ast.walk(tree):
- if not isinstance(node, ast.ClassDef):
- continue
- labels.extend(
- method.name
- for method in node.body
- if isinstance(method, (ast.FunctionDef, ast.AsyncFunctionDef))
- and method.name.startswith("test")
- )
- return labels
-
-
-def _run_startup_tests(repo_root: Path, tests_root: Path, on_test: Optional[Any] = None) -> None:
- """Run repository tests and raise on failure."""
- if not tests_root.exists():
- raise RuntimeError(f"Tests folder not found: {tests_root}")
- command = [sys.executable, "-m", "unittest", "discover", "-s", "tests", "-v", "-p", "test_*.py"]
- process = subprocess.Popen(
- command,
- cwd=str(repo_root),
- stdout=subprocess.PIPE,
- stderr=subprocess.STDOUT,
- text=True,
- encoding="utf-8",
- errors="replace",
- )
- output_lines: list[str] = []
- assert process.stdout is not None
- for line in process.stdout:
- clean_line = line.strip()
- if clean_line:
- output_lines.append(clean_line)
- if on_test is not None and clean_line.startswith("test"):
- on_test(f"Test: {clean_line}")
- QApplication.processEvents()
- return_code = process.wait()
- if return_code != 0:
- tail = "\n".join(output_lines[-25:]).strip()
- raise RuntimeError(f"Startup tests failed.\n{tail}")
-
-
-def _run_startup_validations(splash: StartupValidationSplash) -> None:
- """Run startup checks shown on the splash screen."""
- # ``interface`` is a top-level package after the engine/interface split.
- repo_root = Path(__file__).resolve().parents[2]
- tests_root = repo_root / "tests"
- required_modules = [
- "PySide6", "torch", "PyPDF2", "numpy", "tokenizers",
- "llm_trainer.dataset_build", "llm_trainer.training", "llm_trainer.ui.app",
- ]
- steps: list[tuple[str, Any]] = [
- ("Checking log folder", lambda: _validate_writable_directory(DEFAULT_LOG_DIR)),
- ("Checking cache folder", lambda: _validate_writable_directory(DEFAULT_CACHE_DIR)),
- ("Checking projects folder", lambda: _validate_writable_directory(DEFAULT_PROJECTS_DIR)),
- ("Checking required imports", lambda: [importlib.import_module(name) for name in required_modules]),
- ]
- # Populate the checklist before the subprocess starts. Test callbacks
- # update these entries while unittest is streaming verbose output.
- splash.set_checks(_discover_test_labels(tests_root) if tests_root.is_dir() else [])
- if tests_root.is_dir():
- steps.append((
- "Running test suite",
- lambda: _run_startup_tests(
- repo_root,
- tests_root,
- lambda label: (
- splash.add_check(_test_display_name(label)),
- splash.mark_check_done(_test_display_name(label)),
- ),
- ),
- ))
- else:
- splash.append_log("Repository tests are not included in this packaged installation; skipping test suite.")
- splash.append_log(f"Workspace: {repo_root}")
- for index, (label, action) in enumerate(steps, start=1):
- splash.update_step(f"{label}...", index - 1, len(steps))
- action()
- splash.append_log(f"Completed: {label}")
- splash.update_step("Startup checks complete", len(steps), len(steps))
- splash.append_log("All startup validations passed.")
-
diff --git a/llm_trainer/ui/styles.qss b/llm_trainer/ui/styles.qss
deleted file mode 100644
index 5f0e6c5..0000000
--- a/llm_trainer/ui/styles.qss
+++ /dev/null
@@ -1,288 +0,0 @@
-* { font-family: Arial, "Segoe UI", sans-serif; }
-
-QMainWindow, QWidget#AppShell {
- background: #111111;
- color: #eeeeee;
-}
-
-QWidget#TopBar {
- background: #1f1f1f;
- border-bottom: 1px solid #3a3a3a;
-}
-
-QWidget#SideRail {
- background: #171717;
- border-right: 1px solid #3a3a3a;
-}
-
-QWidget#Panel {
- background: #181818;
- color: #eeeeee;
-}
-
-QWidget#Card {
- background: #242424;
- border: 1px solid #3d3d3d;
- border-radius: 8px;
-}
-
-QWidget#LossChart {
- background: #141414;
- border: 1px solid #3d3d3d;
- border-radius: 8px;
-}
-
-QLabel {
- color: #dddddd;
- font-size: 13px;
-}
-
-QLabel#Logo {
- color: #f5b041;
- font-size: 23px;
- font-weight: 900;
-}
-
-QLabel#PageTitle {
- color: #f2f2f2;
- font-size: 24px;
- font-weight: 700;
-}
-
-QLabel#SectionLabel {
- color: #f5b041;
- font-size: 12px;
- font-weight: 800;
- text-transform: uppercase;
-}
-
-QLabel#TopStatus {
- color: #d8eec2;
- background: #20231d;
- border: 1px solid #8fbf5a;
- border-radius: 7px;
- padding: 5px 8px;
- font-size: 12px;
-}
-
-QLabel#Metric {
- color: #b6d77a;
- font-size: 14px;
- font-weight: 700;
-}
-
-QLabel#MetricChip {
- color: #d8eec2;
- background: #1b1f18;
- border: 1px solid #6f8f45;
- border-radius: 7px;
- padding: 7px 10px;
- font-size: 12px;
- font-weight: 800;
-}
-
-QLabel#TrainingSampleLine {
- color: #d8eec2;
- background: #171a14;
- border: 1px solid #6f8f45;
- border-radius: 6px;
- padding: 5px 10px;
- font-family: Arial;
- font-size: 12px;
-}
-
-QLabel#MessageMeta {
- color: #b6d77a;
- font-size: 11px;
- font-weight: 700;
-}
-
-QLineEdit, QSpinBox, QDoubleSpinBox, QComboBox {
- background: #111111;
- color: #f0f0f0;
- border: 1px solid #555555;
- border-radius: 7px;
- padding: 5px 8px;
- min-height: 20px;
-}
-
-QLineEdit:focus, QSpinBox:focus, QDoubleSpinBox:focus, QComboBox:focus {
- border-color: #f5b041;
- background: #181818;
-}
-
-QCheckBox {
- color: #dddddd;
- spacing: 8px;
-}
-
-QPushButton {
- background: #242424;
- color: #eeeeee;
- border: 1px solid #555555;
- border-radius: 7px;
- padding: 7px 12px;
- font-weight: 800;
-}
-
-QPushButton:hover {
- background: #f5b041;
- color: #151515;
- border-color: #ffd27a;
-}
-
-QPushButton:pressed {
- background: #c98216;
- color: #111111;
- border-color: #f5b041;
-}
-
-QPushButton:disabled {
- background: #1a1a1a;
- color: #777777;
- border-color: #333333;
-}
-
-QPushButton#NavButton {
- background: #202020;
- color: #dddddd;
- border: 1px solid #4a4a4a;
- border-radius: 8px;
- min-width: 52px;
- min-height: 52px;
- padding: 0;
- font-size: 12px;
-}
-
-QPushButton#NavButton:checked, QPushButton#NavButton:hover {
- background: #f5b041;
- color: #151515;
- border-color: #ffd27a;
-}
-
-QPushButton#MessageAction {
- background: transparent;
- color: #d4d4d4;
- border: 0;
- border-radius: 5px;
- padding: 2px 6px;
- font-size: 15px;
- font-weight: 700;
-}
-
-QPushButton#MessageAction:hover {
- color: #f5b041;
- background: #242424;
-}
-
-QTextEdit {
- background: #111111;
- color: #eeeeee;
- border: 1px solid #555555;
- border-radius: 8px;
- padding: 8px;
- font-family: Consolas, monospace;
- font-size: 12px;
-}
-
-QTableWidget {
- background: #111111;
- alternate-background-color: #1b1b1b;
- color: #eeeeee;
- border: 1px solid #555555;
- border-radius: 7px;
- gridline-color: #333333;
- selection-background-color: #3a2d19;
- selection-color: #ffffff;
-}
-
-QHeaderView::section {
- background: #202020;
- color: #f5b041;
- border: 0;
- border-right: 1px solid #3d3d3d;
- border-bottom: 1px solid #3d3d3d;
- padding: 6px 8px;
- font-weight: 800;
-}
-
-QTextBrowser {
- background: #111111;
- color: #eeeeee;
- border: 1px solid #555555;
- border-radius: 8px;
- padding: 12px;
- font-family: Arial, "Segoe UI", sans-serif;
- font-size: 14px;
-}
-
-QScrollArea#ChatScroll {
- background: #111111;
- border: 0;
- border-radius: 0;
-}
-
-QScrollArea#PageScroll {
- background: #181818;
- border: 0;
- border-radius: 0;
-}
-
-QWidget#ChatCanvas {
- background: #111111;
-}
-
-QWidget#UserBubble {
- background: #3a3326;
- border: 1px solid #c98f2e;
- border-radius: 8px;
-}
-
-QWidget#AssistantBubble {
- background: transparent;
- border: 0;
- border-radius: 0;
-}
-
-QTextBrowser#BubbleText {
- background: transparent;
- border: 0;
- border-radius: 0;
- padding: 2px;
-}
-
-QTextEdit#ChatInput, QTextEdit#SystemPrompt {
- background: #111111;
- color: #eeeeee;
- border: 1px solid #555555;
- border-radius: 8px;
- padding: 8px;
- font-family: Arial, "Segoe UI", sans-serif;
- font-size: 13px;
-}
-
-QProgressBar {
- background: #2a2a2a;
- border: 0;
- border-radius: 2px;
- min-height: 4px;
- max-height: 4px;
-}
-
-QProgressBar::chunk {
- background: #f5b041;
- border-radius: 2px;
-}
-
-QProgressBar#HardwareMeter {
- background: #1a1a1a;
- border: 0;
- border-radius: 4px;
- min-height: 8px;
- max-height: 8px;
-}
-
-QProgressBar#HardwareMeter::chunk {
- background: #f5b041;
- border-radius: 4px;
-}
diff --git a/llm_trainer/ui/tabs/__init__.py b/llm_trainer/ui/tabs/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/llm_trainer/ui/tabs/benchmark_tab.py b/llm_trainer/ui/tabs/benchmark_tab.py
deleted file mode 100644
index 0eab9a2..0000000
--- a/llm_trainer/ui/tabs/benchmark_tab.py
+++ /dev/null
@@ -1,87 +0,0 @@
-from __future__ import annotations
-
-from PySide6.QtCore import Qt
-from PySide6.QtWidgets import (
- QCheckBox,
- QGridLayout,
- QLabel,
- QPushButton,
- QScrollArea,
- QSizePolicy,
- QTextEdit,
- QVBoxLayout,
- QWidget,
-)
-
-from llm_trainer.evaluation import DEFAULT_BENCHMARK_PROMPTS
-
-
-def build_benchmark_tab(window) -> QWidget:
- """Build the benchmark prompt page.
-
- Returns:
- Benchmark page widget.
- """
-
- page = window._panel()
- outer = QVBoxLayout(page)
- outer.setContentsMargins(0, 0, 0, 0)
- outer.setSpacing(0)
- scroll = QScrollArea()
- scroll.setObjectName("PageScroll")
- scroll.setWidgetResizable(True)
- scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
- content = QWidget()
- content.setObjectName("Panel")
- layout = QVBoxLayout(content)
- layout.setContentsMargins(18, 18, 18, 10)
- layout.setSpacing(10)
- scroll.setWidget(content)
- outer.addWidget(scroll, 1)
- layout.addWidget(window._page_title("Benchmark Console"))
-
- benchmark_grid = QGridLayout()
- benchmark_grid.setHorizontalSpacing(12)
- benchmark_grid.setVerticalSpacing(8)
- window.benchmark_prompts = QTextEdit()
- window.benchmark_prompts.setMinimumHeight(260)
- window.benchmark_prompts.setPlainText("\n\n".join(DEFAULT_BENCHMARK_PROMPTS))
- window._tip(window.benchmark_prompts, "Benchmark prompts separated by blank lines. Run the same prompts after each training run.")
- window.benchmark_tokens = window._spin(16, 1024, 128)
- window._tip(window.benchmark_tokens, "Maximum generated tokens per benchmark prompt.")
- window.benchmark_temperature = window._double_spin(0.0, 2.0, 0.7, 0.05, 2)
- window._tip(window.benchmark_temperature, "Sampling randomness for benchmark generation.")
- window.benchmark_kv_cache = QCheckBox("Use KV cache")
- window.benchmark_kv_cache.setChecked(True)
- window._tip(window.benchmark_kv_cache, "Reuse attention key/value tensors during MicroGPT benchmark generation for faster inference.")
- window.run_benchmark_button = QPushButton("Run Benchmark")
- window.run_benchmark_button.setMaximumWidth(180)
- window.run_benchmark_button.clicked.connect(window.run_benchmark)
- window._tip(window.run_benchmark_button, "Generate benchmark outputs from final_model.pt and save a benchmark JSON file.")
- window.stop_benchmark_button = QPushButton("Stop")
- window.stop_benchmark_button.setMaximumWidth(120)
- window.stop_benchmark_button.setEnabled(False)
- window.stop_benchmark_button.clicked.connect(window.stop_active_task)
- window._tip(window.stop_benchmark_button, "Request a graceful stop for benchmark generation.")
- benchmark_grid.addWidget(window.benchmark_prompts, 0, 0, 5, 1)
- benchmark_grid.addWidget(QLabel("Max tokens"), 0, 1)
- benchmark_grid.addWidget(window.benchmark_tokens, 0, 2)
- benchmark_grid.addWidget(QLabel("Temperature"), 1, 1)
- benchmark_grid.addWidget(window.benchmark_temperature, 1, 2)
- benchmark_grid.addWidget(window.benchmark_kv_cache, 2, 1, 1, 2)
- benchmark_grid.addWidget(window.run_benchmark_button, 3, 1, 1, 2)
- benchmark_grid.addWidget(window.stop_benchmark_button, 4, 1, 1, 2)
- benchmark_grid.setColumnStretch(0, 1)
- layout.addWidget(window._card("BENCHMARK PROMPTS", benchmark_grid), 0)
-
- window.benchmark_log = QTextEdit()
- window.benchmark_log.setReadOnly(True)
- window.benchmark_log.setMinimumHeight(260)
- window.benchmark_log.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
- benchmark_log_layout = QVBoxLayout()
- benchmark_log_layout.addWidget(window.benchmark_log, 1)
- layout.addWidget(window._card("BENCHMARK TELEMETRY", benchmark_log_layout), 1)
-
- window.benchmark_progress = window._thin_progress()
- outer.addWidget(window.benchmark_progress)
- return page
diff --git a/llm_trainer/ui/tabs/chat_tab.py b/llm_trainer/ui/tabs/chat_tab.py
deleted file mode 100644
index 00fb5dc..0000000
--- a/llm_trainer/ui/tabs/chat_tab.py
+++ /dev/null
@@ -1,173 +0,0 @@
-from __future__ import annotations
-
-from PySide6.QtWidgets import (
- QCheckBox,
- QComboBox,
- QFormLayout,
- QHBoxLayout,
- QLabel,
- QLineEdit,
- QPushButton,
- QScrollArea,
- QSizePolicy,
- QTextEdit,
- QVBoxLayout,
- QWidget,
-)
-
-from llm_trainer.ui.chat_widgets import ChatInputEdit
-
-
-def build_chat_tab(window) -> QWidget:
- """Build the GGUF model test chat page.
-
- Returns:
- Chat page widget.
- """
-
- page = window._panel()
- layout = QVBoxLayout(page)
- layout.setContentsMargins(24, 20, 24, 14)
- layout.setSpacing(12)
-
- main = QHBoxLayout()
- main.setSpacing(14)
-
- chat_column = QVBoxLayout()
- chat_column.setSpacing(10)
-
- window.chat_scroll = QScrollArea()
- window.chat_scroll.setObjectName("ChatScroll")
- window.chat_scroll.setWidgetResizable(True)
- window.chat_scroll.setMinimumHeight(420)
- window._tip(window.chat_scroll, "Rendered Markdown conversation view.")
- window.chat_canvas = QWidget()
- window.chat_canvas.setObjectName("ChatCanvas")
- window.chat_messages = QVBoxLayout(window.chat_canvas)
- window.chat_messages.setContentsMargins(14, 14, 14, 14)
- window.chat_messages.setSpacing(12)
- window.chat_messages.addStretch(1)
- window.chat_scroll.setWidget(window.chat_canvas)
- window.chat_event_log = QTextEdit()
- window.chat_event_log.setVisible(False)
- window._add_chat_message("assistant", "Load a GGUF model to start testing.")
- window.chat_stats = QLabel("Idle")
- window.chat_stats.setObjectName("Metric")
- window.chat_stats.setVisible(False)
- window._tip(window.chat_stats, "Generation timing, produced tokens, and approximate token speed.")
- chat_column.addWidget(window.chat_scroll, 1)
-
- prompt_row = QHBoxLayout()
- prompt_row.setSpacing(10)
- window.chat_input = ChatInputEdit()
- window.chat_input.setObjectName("ChatInput")
- window.chat_input.setMaximumHeight(92)
- window.chat_input.setPlaceholderText("Send a message...")
- window._tip(window.chat_input, "Prompt to send to the loaded model.")
- window.chat_input.sendRequested.connect(window.send_chat_message)
- window.send_chat_button = QPushButton("Send")
- window.send_chat_button.setMaximumWidth(120)
- window.send_chat_button.clicked.connect(window.send_chat_message)
- window._tip(window.send_chat_button, "Send the message to the already loaded model.")
- window.stop_chat_button = QPushButton("Stop")
- window.stop_chat_button.setMaximumWidth(120)
- window.stop_chat_button.setEnabled(False)
- window.stop_chat_button.clicked.connect(window.stop_active_task)
- window._tip(window.stop_chat_button, "Stop the current streamed reply.")
- prompt_row.addWidget(window.chat_input, 1)
- prompt_row.addWidget(window.send_chat_button)
- prompt_row.addWidget(window.stop_chat_button)
- chat_column.addLayout(prompt_row)
-
- settings_column = QVBoxLayout()
- settings_column.setSpacing(12)
- settings_panel = QWidget()
- settings_panel.setMaximumWidth(390)
- settings_panel.setMinimumWidth(340)
- settings_panel.setLayout(settings_column)
-
- model_form = QFormLayout()
- window._configure_form(model_form)
- window.chat_model_backend = QComboBox()
- window.chat_model_backend.addItems(["GGUF / llama.cpp", "MicroGPT checkpoint"])
- window.chat_model_backend.setMaximumWidth(260)
- window._tip(window.chat_model_backend, "Choose whether chat loads a GGUF model or a native MicroGPT final_model.pt checkpoint.")
- window.gguf_path = QLineEdit()
- window._tip(window.gguf_path, "Path to a GGUF model file produced by llama.cpp-compatible export tooling.")
- window.microgpt_chat_path = QLineEdit()
- window._tip(window.microgpt_chat_path, "MicroGPT final_model.pt checkpoint. A model folder containing final_model.pt and tokenizer.json also works if typed.")
- window.llama_context = window._spin(256, 131072, 2048)
- window._tip(window.llama_context, "llama.cpp context window. Larger values allow longer chats but use more memory.")
- window.llama_threads = window._spin(1, 128, 4)
- window._tip(window.llama_threads, "CPU threads used by llama.cpp inference.")
- window.llama_gpu_layers = window._spin(-1, 200, -1)
- window._tip(window.llama_gpu_layers, "Number of transformer layers to offload to GPU. Use -1 to offload all possible layers.")
- model_form.addRow("Model type", window.chat_model_backend)
- window.gguf_path_row = window._path_row(window.gguf_path, directory=False, file_filter="GGUF models (*.gguf);;All files (*)")
- window.microgpt_path_row = window._path_row(window.microgpt_chat_path, directory=False, file_filter="Checkpoints (*.pt);;All files (*)")
- model_form.addRow("GGUF model", window.gguf_path_row)
- model_form.addRow("MicroGPT checkpoint", window.microgpt_path_row)
- model_form.addRow("Context", window.llama_context)
- model_form.addRow("CPU threads", window.llama_threads)
- model_form.addRow("GPU layers", window.llama_gpu_layers)
- window.load_llm_button = QPushButton("Load Model")
- window.load_llm_button.setMaximumWidth(180)
- window.load_llm_button.clicked.connect(window.toggle_llm_model)
- window._tip(window.load_llm_button, "Load the selected model into memory once for repeated chat messages.")
- window.reset_chat_button = QPushButton("Reset Chat")
- window.reset_chat_button.setMaximumWidth(180)
- window.reset_chat_button.clicked.connect(window.reset_chat)
- window._tip(window.reset_chat_button, "Clear conversation memory while keeping the model loaded.")
- loader_buttons = QHBoxLayout()
- loader_buttons.addWidget(window.load_llm_button)
- loader_buttons.addWidget(window.reset_chat_button)
- loader_buttons.addStretch(1)
- model_form.addRow("", loader_buttons)
- window.chat_model_backend.currentTextChanged.connect(window._update_chat_backend_controls)
-
- sample_form = QFormLayout()
- window._configure_form(sample_form)
- window.thinking_enabled = QCheckBox("Thinking")
- window.thinking_enabled.setChecked(True)
- window._tip(window.thinking_enabled, "When enabled, the prompt asks the model to reason according to the selected effort level. Turn off for direct answers.")
- window.reasoning_effort = QComboBox()
- window.reasoning_effort.addItems(["Balanced", "Fast", "Deep"])
- window.reasoning_effort.setMaximumWidth(260)
- window._tip(window.reasoning_effort, "Controls the instruction style sent with each prompt. Deep asks for more careful reasoning.")
- window.thinking_enabled.toggled.connect(window.reasoning_effort.setEnabled)
- window.chat_max_tokens = window._spin(16, 8192, 512)
- window._tip(window.chat_max_tokens, "Maximum new tokens for each assistant reply.")
- window.chat_temperature = window._double_spin(0.0, 2.0, 0.7, 0.05, 2)
- window._tip(window.chat_temperature, "Sampling randomness. Lower is more focused; higher is more creative.")
- window.chat_top_p = window._double_spin(0.01, 1.0, 0.9, 0.01, 2)
- window._tip(window.chat_top_p, "Nucleus sampling. Lower values restrict the model to more likely tokens.")
- window.chat_repeat_penalty = window._double_spin(0.8, 2.0, 1.1, 0.01, 2)
- window._tip(window.chat_repeat_penalty, "Penalty for repeated text. Higher can reduce loops.")
- sample_form.addRow("", window.thinking_enabled)
- sample_form.addRow("Reasoning effort", window.reasoning_effort)
- sample_form.addRow("Max tokens", window.chat_max_tokens)
- sample_form.addRow("Temperature", window.chat_temperature)
- sample_form.addRow("Top-p", window.chat_top_p)
- sample_form.addRow("Repeat penalty", window.chat_repeat_penalty)
-
- window.system_prompt = QTextEdit()
- window.system_prompt.setObjectName("SystemPrompt")
- window.system_prompt.setMaximumHeight(120)
- window.system_prompt.setPlaceholderText("Optional system prompt")
- window._tip(window.system_prompt, "Optional behavior instruction sent to the model with each message.")
- system_layout = QVBoxLayout()
- system_layout.addWidget(window.system_prompt)
-
- settings_column.addWidget(window._card("MODEL LOADER", model_form))
- settings_column.addWidget(window._card("RESPONSE TUNING", sample_form))
- settings_column.addWidget(window._card("SYSTEM PROMPT", system_layout))
- settings_column.addStretch(1)
-
- main.addLayout(chat_column, 1)
- main.addWidget(settings_panel)
- layout.addLayout(main, 1)
-
- window.chat_progress = window._thin_progress()
- layout.addWidget(window.chat_progress)
- window._update_chat_backend_controls()
- return page
diff --git a/llm_trainer/ui/tabs/dataset_plan_tab.py b/llm_trainer/ui/tabs/dataset_plan_tab.py
deleted file mode 100644
index 8f5b283..0000000
--- a/llm_trainer/ui/tabs/dataset_plan_tab.py
+++ /dev/null
@@ -1,457 +0,0 @@
-from __future__ import annotations
-
-from pathlib import Path
-import re
-from typing import Any
-
-from PySide6.QtCore import Qt
-from PySide6.QtWidgets import (
- QCheckBox,
- QComboBox,
- QFormLayout,
- QGridLayout,
- QHBoxLayout,
- QLabel,
- QLineEdit,
- QMenu,
- QPushButton,
- QScrollArea,
- QSizePolicy,
- QTreeWidget,
- QTreeWidgetItem,
- QVBoxLayout,
- QWidget,
- QWidgetAction,
-)
-
-from llm_trainer.conversation_datasets import CONVERSATION_DATASET_PRESETS
-
-
-
-CODE_SUFFIXES = {".py", ".js", ".ts", ".tsx", ".jsx", ".java", ".c", ".cpp", ".h", ".hpp", ".cs", ".go", ".rs", ".sh", ".ps1"}
-SUPPORTED_DEFAULT_SUFFIXES = {".txt", ".md", ".text", ".jsonl", ".json", *CODE_SUFFIXES}
-
-
-def default_data_root() -> Path:
- """Return the bundled default data folder.
-
- Returns:
- Absolute path to the packaged ``default_data`` folder.
- """
-
- return Path(__file__).resolve().parents[2] / "default_data"
-
-
-def blueprint_data_root(window: Any | None = None) -> Path:
- """Return the active Dataset Blueprint data root.
-
- Args:
- window: Optional main window carrying a project-local data root.
-
- Returns:
- Project-local training data root when available, otherwise bundled data.
- """
-
- root = getattr(window, "blueprint_data_root", None)
- if root:
- path = Path(root)
- if path.exists():
- return path
- return default_data_root()
-
-
-def _slugify_category(value: str) -> str:
- """Convert folder/file text into a stable category key.
-
- Args:
- value: Folder name, file stem, or user-facing text.
-
- Returns:
- Lowercase underscore category key.
- """
-
- slug = re.sub(r"[^a-z0-9]+", "_", value.lower()).strip("_")
- return slug or "general_prose"
-
-
-def dataset_category_label(key: str) -> str:
- """Return a readable label for a category key.
-
- Args:
- key: Dataset category key.
-
- Returns:
- User-facing label.
- """
-
- return key.replace("_", " ").title()
-
-
-def default_data_category(path: Path, root: Path | None = None) -> str:
- """Infer the Dataset Blueprint category for a bundled file.
-
- Args:
- path: Bundled source file.
-
- Returns:
- Dataset category key used by the sampler.
- """
-
- root = root or default_data_root()
- try:
- relative = path.relative_to(root)
- except ValueError:
- relative = path
- # The first directory below the configured root is the category. Do not
- # interpret names or extensions: adding a folder is the complete
- # configuration needed to add a new category.
- folders = relative.parts[:-1]
- if folders:
- return _slugify_category(folders[0])
- return _slugify_category(path.stem)
-
-
-def default_data_stage(path: Path, root: Path | None = None) -> str:
- """Infer which training stage should use a bundled file.
-
- Args:
- path: Bundled source file.
-
- Returns:
- Stage key: base, instruction, conversation, or code.
- """
-
- root = root or default_data_root()
- try:
- relative = path.relative_to(root)
- except ValueError:
- relative = path
- folders = relative.parts[:-1]
- return _slugify_category(folders[0]) if folders else "base"
-
-
-def iter_default_data_files(root: Path | None = None) -> list[tuple[Path, str]]:
- """List default/project data files with categories.
-
- Args:
- root: Optional source root. Defaults to bundled default data.
-
- Returns:
- Pairs of file path and Dataset Blueprint category.
- """
-
- root = root or default_data_root()
- if not root.exists():
- return []
- return [
- (path, default_data_category(path, root))
- for path in sorted(root.rglob("*"))
- if (
- path.is_file()
- and path.suffix.lower() in SUPPORTED_DEFAULT_SUFFIXES
- and path.stat().st_size > 0
- )
- ]
-
-
-def file_token_vocab_stats(path: Path, sample_bytes: int = 256 * 1024) -> dict[str, int | bool]:
- """Estimate token and vocabulary counts for a data file.
-
- Args:
- path: Source file path.
- sample_bytes: Maximum bytes to read for a fast estimate.
-
- Returns:
- Dictionary containing size, estimated tokens, estimated vocab, and
- whether values were extrapolated from a sample.
- """
-
- size = path.stat().st_size
- if path.suffix.lower() in {".json", ".jsonl", ".txt", ".md", ".text", *CODE_SUFFIXES}:
- with path.open("rb") as handle:
- raw = handle.read(sample_bytes)
- text = raw.decode("utf-8", errors="ignore")
- pieces = re.findall(r"\w+|[^\w\s]", text)
- vocab = {piece.lower() for piece in pieces if piece.strip()}
- multiplier = size / max(len(raw), 1) if raw and size > len(raw) else 1.0
- return {
- "bytes": size,
- "characters": int(round(len(text) * multiplier)),
- "tokens": int(round(len(pieces) * multiplier)),
- "vocab": int(round(len(vocab) * min(multiplier, 3.0))),
- "sampled": size > len(raw),
- }
- return {"bytes": size, "characters": 0, "tokens": 0, "vocab": 0, "sampled": False}
-
-
-def format_estimate(value: int, sampled: bool) -> str:
- """Format a numeric estimate for the tree widget."""
-
- prefix = "~" if sampled else ""
- return f"{prefix}{value:,}"
-
-
-def dataset_plan_defaults(default_files: list[tuple[Path, str]] | None = None) -> dict[str, float]:
- """Return default blueprint weights plus discovered default-data categories.
-
- Args:
- default_files: Optional pre-discovered bundled file/category pairs.
-
- Returns:
- Default category weight mapping.
- """
-
- if default_files is None:
- return {}
- categories: list[str] = []
- seen: set[str] = set()
- for _path, category in default_files:
- if category not in seen:
- categories.append(category)
- seen.add(category)
- if not categories:
- return {}
- return {category: 0.0 for category in categories}
-
-
-def build_dataset_plan_tab(window) -> QWidget:
- """Build the dataset blueprint page.
-
- Args:
- window: Main application window that owns shared helper methods.
-
- Returns:
- Dataset blueprint page widget.
- """
-
- page = QWidget()
- outer = QVBoxLayout(page)
- outer.setContentsMargins(18, 18, 18, 12)
- outer.setSpacing(12)
-
- scroll = QScrollArea()
- scroll.setWidgetResizable(True)
- scroll.setFrameShape(QScrollArea.NoFrame)
- content = QWidget()
- layout = QVBoxLayout(content)
- layout.setContentsMargins(0, 0, 0, 0)
- layout.setSpacing(12)
- layout.setAlignment(Qt.AlignTop)
-
- title_row = QHBoxLayout()
- title = QLabel("Dataset Sources")
- title.setObjectName("PageTitle")
- active_data_root = blueprint_data_root(window)
- default_files = iter_default_data_files(active_data_root)
- window.blueprint_data_root = active_data_root
-
- window.dataset_plan_source_label = QLabel(f"Source: {active_data_root}")
- window.dataset_plan_source_label.setObjectName("Muted")
- window.dataset_plan_refresh_button = QPushButton("Refresh")
- window.dataset_plan_refresh_button.setMaximumWidth(110)
- title_row.addWidget(title)
- title_row.addSpacing(12)
- title_row.addWidget(window.dataset_plan_source_label, 1)
- title_row.addStretch(1)
- layout.addLayout(title_row)
-
- external_form = QFormLayout()
- window._configure_form(external_form)
- window.external_dataset_dir = QLineEdit(str(Path.home() / "drunkenbot_datasets" / "default"))
- window.external_dataset_version = QLabel("Installed version: not installed")
- window.external_dataset_version.setObjectName("Muted")
- window._tip(window.external_dataset_dir, "Folder where downloaded dataset categories are installed and used as the ingestion source.")
- window._tip(window.external_dataset_version, "Version recorded from the installed dataset manifest.")
- window.external_dataset_download_button = QPushButton("Download latest dataset")
- window._tip(window.external_dataset_download_button, "Download, verify, and extract the latest dataset release into the install folder.")
- window.external_dataset_download_button.clicked.connect(window.download_latest_external_dataset)
- external_form.addRow("Install folder", window._path_row(window.external_dataset_dir, directory=True))
- external_form.addRow("Status", window.external_dataset_version)
- external_form.addRow("", window.external_dataset_download_button)
- external_card = window._card("EXTERNAL DATASET", external_form)
- body_grid = QGridLayout()
- body_grid.setHorizontalSpacing(14)
- body_grid.setVerticalSpacing(12)
-
- conversation_form = QFormLayout()
- window._configure_form(conversation_form)
- window.dataset_stage = QComboBox()
- # Workflows are directories in the configured training-data root. This
- # keeps the UI in sync with the corpus instead of requiring code changes
- # for every new training workflow.
- workflow_names = sorted(
- (path.name for path in active_data_root.iterdir() if path.is_dir()),
- key=str.casefold,
- ) if active_data_root.exists() else []
- window.dataset_stage.addItems(workflow_names or ["base"])
- window.dataset_stage.setMaximumWidth(240)
- window.include_conversation_datasets = QCheckBox("Online")
- window.include_conversation_datasets.setChecked(False)
- purpose_row = QWidget()
- purpose_layout = QHBoxLayout(purpose_row)
- purpose_layout.setContentsMargins(0, 0, 0, 0)
- purpose_layout.setSpacing(8)
- purpose_layout.addWidget(window.dataset_stage, 1)
- purpose_layout.addWidget(window.include_conversation_datasets)
- conversation_form.addRow("Purpose", purpose_row)
- window.conversation_datasets_status = QLabel("Base pretraining: choose optional online corpus datasets, or use local files only.")
- window.conversation_datasets_status.setObjectName("Muted")
- conversation_form.addRow("", window.conversation_datasets_status)
- window.conversation_dataset_button = QPushButton("Online datasets off")
- window.conversation_dataset_button.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed)
- window.conversation_dataset_menu = QMenu(window.conversation_dataset_button)
- window.conversation_dataset_button.setMenu(window.conversation_dataset_menu)
- window.conversation_dataset_actions = {}
- window.conversation_dataset_widget_actions = {}
- for dataset_id, preset in CONVERSATION_DATASET_PRESETS.items():
- checkbox = QCheckBox(preset.label)
- checkbox.setEnabled(False)
- checkbox.setToolTip(preset.description)
- checkbox.toggled.connect(lambda _checked=False: window._update_conversation_dataset_button_text())
- widget_action = QWidgetAction(window.conversation_dataset_menu)
- widget_action.setDefaultWidget(checkbox)
- window.conversation_dataset_menu.addAction(widget_action)
- window.conversation_dataset_actions[dataset_id] = checkbox
- window.conversation_dataset_widget_actions[dataset_id] = widget_action
- conversation_form.addRow("Online sets", window.conversation_dataset_button)
- window.custom_huggingface_dataset = QLineEdit()
- window.custom_huggingface_dataset.setPlaceholderText("owner/dataset or Hugging Face URL")
- window._tip(window.custom_huggingface_dataset, "Optional Hugging Face dataset repository to load in addition to the selected presets.")
- conversation_form.addRow("Custom HF dataset", window.custom_huggingface_dataset)
- window.custom_huggingface_download = QPushButton("Download custom dataset")
- window.custom_huggingface_download.clicked.connect(window._download_custom_huggingface_dataset)
- window._tip(window.custom_huggingface_download, "Enable the custom dataset and download it during the next dataset preparation run.")
- conversation_form.addRow("", window.custom_huggingface_download)
- window.conversation_sample_limit = window._spin(0, 2_000_000, 20000)
- window.conversation_sample_limit.setMaximumHeight(30)
- window.conversation_sample_limit.setEnabled(False)
- window.include_conversation_datasets.toggled.connect(window.conversation_sample_limit.setEnabled)
- window.include_conversation_datasets.toggled.connect(window._update_online_dataset_stage_controls)
- window.dataset_stage.currentTextChanged.connect(window._update_online_dataset_stage_controls)
- conversation_form.addRow("Rows / set", window.conversation_sample_limit)
- conversation_card = window._card("OPTIONAL EXTERNAL / STRUCTURED DATA", conversation_form)
- conversation_card.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Maximum)
- conversation_card.setMaximumHeight(250)
-
- window.default_data_tree_updating = False
- window.default_data_tree = QTreeWidget()
- window.default_data_tree.setHeaderLabels(["Category / file", "Characters", "Vocab"])
- window.default_data_tree.setRootIsDecorated(True)
- window.default_data_tree.setAlternatingRowColors(False)
- window.default_data_tree.setMinimumHeight(600)
- window.default_data_tree.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
- window.default_data_tree.setColumnWidth(0, 360)
- window.default_data_tree.setColumnWidth(1, 130)
- window.default_data_tree.setColumnWidth(2, 130)
- window.default_data_actions = {}
- window.default_data_category_items = {}
- window.default_data_tree.clear()
- grouped_files: dict[str, list[Path]] = {}
- for path, category in default_files:
- grouped_files.setdefault(category, []).append(path)
- for category in sorted(grouped_files, key=dataset_category_label):
- total_characters = 0
- total_vocab = 0
- category_sampled = False
- category_item = QTreeWidgetItem([dataset_category_label(category), "0", "0"])
- category_item.setData(0, Qt.UserRole, {"kind": "category", "category": category})
- category_item.setFlags(category_item.flags() | Qt.ItemIsUserCheckable)
- category_item.setCheckState(0, Qt.Checked)
- window.default_data_tree.addTopLevelItem(category_item)
- window.default_data_category_items[category] = category_item
- for path in sorted(grouped_files[category], key=lambda item: item.name.lower()):
- try:
- stats = file_token_vocab_stats(path)
- except OSError:
- stats = {"characters": 0, "vocab": 0, "sampled": False}
- sampled = bool(stats.get("sampled", False))
- characters = int(stats.get("characters", 0))
- vocab = int(stats.get("vocab", 0))
- child = QTreeWidgetItem([path.name, format_estimate(characters, sampled), format_estimate(vocab, sampled)])
- child.setToolTip(0, str(path))
- child.setData(0, Qt.UserRole, {"kind": "file", "path": str(path), "category": category})
- child.setFlags(child.flags() | Qt.ItemIsUserCheckable)
- child.setCheckState(0, Qt.Checked)
- category_item.addChild(child)
- window.default_data_actions[str(path)] = child
- total_characters += characters
- total_vocab += vocab
- category_sampled = category_sampled or sampled
- category_item.setText(1, format_estimate(total_characters, category_sampled))
- category_item.setText(2, format_estimate(total_vocab, category_sampled))
- category_item.setExpanded(False)
- if not window.default_data_actions:
- window.default_data_tree.addTopLevelItem(QTreeWidgetItem(["No project/default data files were found.", "", ""]))
- window.default_data_tree.itemChanged.connect(window._handle_default_data_tree_changed)
- default_layout = QVBoxLayout()
- tree_title_row = QHBoxLayout()
- tree_title_row.addWidget(QLabel("Downloaded and local dataset files"))
- tree_title_row.addStretch(1)
- tree_title_row.addWidget(window.dataset_plan_refresh_button)
- default_layout.addLayout(tree_title_row)
- default_layout.addWidget(window.default_data_tree)
- default_card = window._card("BUNDLED DEFAULT DATA", default_layout)
- body_grid.addWidget(external_card, 0, 0)
- body_grid.addWidget(conversation_card, 1, 0)
- body_grid.addWidget(default_card, 0, 1, 2, 1)
- body_grid.setColumnStretch(0, 1)
- body_grid.setColumnStretch(1, 1)
- layout.addLayout(body_grid)
- window.dataset_plan_refresh_button.clicked.connect(window.refresh_dataset_blueprint_files)
- window._tip(window.dataset_plan_refresh_button, "Reload this tree to include newly copied files and folders.")
- window._update_online_dataset_stage_controls()
-
- scroll.setWidget(content)
- outer.addWidget(scroll, 1)
- window.dataset_plan_progress = window._thin_progress()
- window.dataset_plan_progress.setVisible(False)
- outer.addWidget(window.dataset_plan_progress)
- return page
-
-
-def populate_default_data_tree(window: Any, root: Path) -> None:
- """Reload the existing dataset tree without rebuilding the containing tab."""
- tree = window.default_data_tree
- tree.blockSignals(True)
- try:
- tree.clear()
- window.default_data_actions = {}
- window.default_data_category_items = {}
- grouped_files: dict[str, list[Path]] = {}
- for path, category in iter_default_data_files(root):
- grouped_files.setdefault(category, []).append(path)
- for category in sorted(grouped_files, key=dataset_category_label):
- category_item = QTreeWidgetItem([dataset_category_label(category), "0", "0"])
- category_item.setData(0, Qt.UserRole, {"kind": "category", "category": category})
- category_item.setFlags(category_item.flags() | Qt.ItemIsUserCheckable)
- category_item.setCheckState(0, Qt.Checked)
- tree.addTopLevelItem(category_item)
- window.default_data_category_items[category] = category_item
- total_characters = total_vocab = 0
- sampled_category = False
- for path in sorted(grouped_files[category], key=lambda item: item.name.lower()):
- try:
- stats = file_token_vocab_stats(path)
- except OSError:
- stats = {"characters": 0, "vocab": 0, "sampled": False}
- sampled = bool(stats.get("sampled", False))
- characters = int(stats.get("characters", 0))
- vocab = int(stats.get("vocab", 0))
- child = QTreeWidgetItem(
- [path.name, format_estimate(characters, sampled), format_estimate(vocab, sampled)]
- )
- child.setToolTip(0, str(path))
- child.setData(0, Qt.UserRole, {"kind": "file", "path": str(path), "category": category})
- child.setFlags(child.flags() | Qt.ItemIsUserCheckable)
- child.setCheckState(0, Qt.Checked)
- category_item.addChild(child)
- window.default_data_actions[str(path)] = child
- total_characters += characters
- total_vocab += vocab
- sampled_category = sampled_category or sampled
- category_item.setText(1, format_estimate(total_characters, sampled_category))
- category_item.setText(2, format_estimate(total_vocab, sampled_category))
- finally:
- tree.blockSignals(False)
- if not window.default_data_actions:
- tree.addTopLevelItem(QTreeWidgetItem(["No project/default data files were found.", "", ""]))
diff --git a/llm_trainer/ui/tabs/dataset_tab.py b/llm_trainer/ui/tabs/dataset_tab.py
deleted file mode 100644
index 7f481f5..0000000
--- a/llm_trainer/ui/tabs/dataset_tab.py
+++ /dev/null
@@ -1,300 +0,0 @@
-from __future__ import annotations
-
-from pathlib import Path
-
-from PySide6.QtCore import Qt
-from PySide6.QtWidgets import (
- QCheckBox,
- QComboBox,
- QFormLayout,
- QGridLayout,
- QHBoxLayout,
- QLabel,
- QLineEdit,
- QMenu,
- QPushButton,
- QScrollArea,
- QSizePolicy,
- QTextEdit,
- QVBoxLayout,
- QWidget,
- QWidgetAction,
-)
-
-from llm_trainer.conversation_datasets import CONVERSATION_DATASET_PRESETS
-from llm_trainer.ui.charts import DatasetBarChartWidget
-
-
-def build_dataset_tab(window) -> QWidget:
- """Build the dataset preparation page.
-
- Returns:
- Dataset page widget.
- """
-
- page = window._panel()
- outer = QVBoxLayout(page)
- outer.setContentsMargins(0, 0, 0, 0)
- outer.setSpacing(0)
- scroll = QScrollArea()
- scroll.setObjectName("PageScroll")
- scroll.setWidgetResizable(True)
- scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
- content = QWidget()
- content.setObjectName("Panel")
- layout = QVBoxLayout(content)
- layout.setContentsMargins(18, 18, 18, 10)
- layout.setSpacing(10)
- scroll.setWidget(content)
- outer.addWidget(scroll, 1)
- title_row = QHBoxLayout()
- title_row.setSpacing(10)
- title = window._page_title("Data Ingestion Matrix")
- title_row.addWidget(title, 0)
- window.dataset_quality_samples = window._metric_chip("Documents: -", "Prepared source documents before token sliding windows.")
- window.dataset_quality_tokens = window._metric_chip("Tokens: -", "Total encoded tokens available for training.")
- window.dataset_quality_windows = window._metric_chip("Windows: -", "Sliding context windows the trainer can sample.")
- window.dataset_quality_vocab = window._metric_chip("Vocab: -", "Tokenizer vocabulary size used by the dataset.")
- window.dataset_quality_rating = window._metric_chip("Rating: -", "Five-star dataset quality score based on tokens, windows, vocabulary, diversity, and extraction health.")
- window.dataset_quality_code = window._metric_chip("Code/prose: -", "Code and prose sample split.")
- window.dataset_quality_balance = window._metric_chip("Balance: -", "Code/prose balance detected during preview or preparation.")
- window.dataset_quality_readiness = window._metric_chip("Readiness: -", "Training readiness score based on size, duplicates, extraction quality, and dataset mix.")
- window.dataset_quality_cache = window._metric_chip("Cache: -", "Files reused from cache versus processed this run.")
- window.dataset_quality_duplicates = window._metric_chip("Duplicates: -", "Likely exact or extracted-text duplicate files.")
- window.dataset_quality_extraction = window._metric_chip("Extraction: -", "Files with suspicious text extraction quality.")
- window.dataset_quality_warning = window._metric_chip("Warnings: none", "Dataset quality warnings, if any.")
- header_quality_items = [
- window.dataset_quality_samples,
- window.dataset_quality_tokens,
- window.dataset_quality_windows,
- window.dataset_quality_vocab,
- window.dataset_quality_rating,
- window.dataset_quality_code,
- window.dataset_quality_readiness,
- window.dataset_quality_warning,
- ]
- for item in header_quality_items:
- item.setMaximumWidth(210)
- title_row.addWidget(item, 1)
- layout.addLayout(title_row)
-
- ingestion_body = QHBoxLayout()
- ingestion_body.setSpacing(14)
- left_column = QVBoxLayout()
- left_column.setSpacing(10)
- right_column = QVBoxLayout()
- right_column.setSpacing(10)
- ingestion_body.addLayout(left_column, 1)
- ingestion_body.addLayout(right_column, 1)
-
- source_form = QFormLayout()
- window._configure_form(source_form)
- tokenizer_form = QFormLayout()
- window._configure_form(tokenizer_form)
-
- form = QFormLayout()
- window._configure_form(form)
-
- window.input_dir = QLineEdit()
- window._tip(window.input_dir, "Folder containing PDFs, text, Markdown, or JSONL files. More clean text usually improves the model.")
- window.dataset_dir = QLineEdit(str(Path.cwd() / "runs" / "dataset"))
- window._tip(window.dataset_dir, "Folder where prepared corpus, tokenizer, token files, and dataset summary are saved.")
- window.auto_vocab = QCheckBox("Choose automatically")
- window.auto_vocab.setChecked(True)
- window._tip(window.auto_vocab, "Automatically choose vocabulary size based on corpus size and word variety. Safer for most users.")
- window.manual_vocab_size = window._spin(256, 100000, 8000)
- window.manual_vocab_size.setEnabled(False)
- window._tip(window.manual_vocab_size, "Manual tokenizer vocabulary size. Larger vocab can preserve more words but increases model output size.")
- window.auto_vocab.toggled.connect(lambda checked: window.manual_vocab_size.setEnabled(not checked and not window._tokenizer_strategy_reuses()))
- window.auto_vocab_label = QLabel("Auto after reading files")
- window.auto_vocab_label.setObjectName("Metric")
- window._tip(window.auto_vocab_label, "The actual vocabulary size selected after reading the corpus.")
- window.min_frequency = window._spin(1, 1000, 2)
- window._tip(window.min_frequency, "Minimum token frequency for tokenizer training. Higher values remove rare fragments and can reduce noise.")
- window.context_length = window._spin(16, 4096, 128)
- window._tip(window.context_length, "Number of tokens per training sequence. Longer context lets the model learn longer dependencies but uses more memory.")
- window.validation_split = window._double_spin(0.0, 0.5, 0.1, 0.01, 3)
- window._tip(window.validation_split, "Fraction of tokens held out for validation. Validation helps detect overfitting during training.")
- window.max_workers = window._spin(1, 64, 4)
- window._tip(
- window.max_workers,
- "Number of source files extracted in parallel, each in its own process (capped by your CPU core count). "
- "Faster on multi-core machines, but peak memory scales with this number -- each worker holds one "
- "file's full text in memory while processing it. Lower this if you are extracting many very large "
- "files (e.g. multi-gigabyte dumps) and see high memory use.",
- )
- window.tokenizer_training_max_gb = window._double_spin(0.0, 256.0, 2.0, 0.5, 1)
- window._tip(
- window.tokenizer_training_max_gb,
- "Maximum corpus size (in GiB) shown to the tokenizer trainer when learning vocabulary. The trainer "
- "keeps a frequency table in memory sized to whatever it is shown, so very large corpora are sampled "
- "down to this size by default -- vocabulary quality does not meaningfully improve past a few GiB of "
- "sample text. Raise this if you have more RAM to spare (this is separate from, and much smaller than, "
- "your training data itself -- the full corpus is always encoded into training tokens regardless of "
- "this setting). Set to 0 to disable the cap entirely and train on the full corpus; only do this if "
- "you are confident you have enough RAM to hold a frequency table sized to your whole corpus at once.",
- )
- window.prepare_mode = QComboBox()
- window.prepare_mode.addItems(["Incremental update", "Full rebuild", "Force reprocess"])
- window.prepare_mode.setMaximumWidth(260)
- window._tip(
- window.prepare_mode,
- "Incremental update reuses cached extracted text and the existing tokenizer. Full rebuild rebuilds tokenizer/tokens. Force reprocess ignores cache.",
- )
- window.tokenizer_strategy = QComboBox()
- window.tokenizer_strategy.addItems(["Auto", "Train new tokenizer", "Reuse dataset tokenizer", "Import tokenizer.json"])
- window.tokenizer_strategy.setMaximumWidth(260)
- window._tip(
- window.tokenizer_strategy,
- "Controls tokenizer reuse. Auto reuses the dataset tokenizer during incremental updates; Import lets you use a compatible tokenizer.json.",
- )
- window.tokenizer_path = QLineEdit()
- window.tokenizer_path.setEnabled(False)
- window._tip(window.tokenizer_path, "Existing tokenizer.json to import. Use this when continuing a compatible tokenizer family.")
- window.tokenizer_strategy.currentTextChanged.connect(window._update_tokenizer_strategy_controls)
- window.code_training_mode = QCheckBox("Code-aware processing")
- window.code_training_mode.setChecked(True)
- window._tip(
- window.code_training_mode,
- "Use code-aware cleaning, category tags, and code/prose balancing. Keep this on for programming books, source folders, and technical datasets.",
- )
- window.include_prose = QCheckBox("Include explanations")
- window.include_prose.setChecked(True)
- window._tip(window.include_prose, "Keep prose from PDFs/books. This helps the model learn programming concepts and explanations.")
- window.include_source_code = QCheckBox("Include source files")
- window.include_source_code.setChecked(True)
- window._tip(window.include_source_code, "Include real code files such as .py, .js, .java, .cpp, .cs, .go, .rs, and similar.")
- window.extract_code_blocks = QCheckBox("Extract code blocks")
- window.extract_code_blocks.setChecked(True)
- window._tip(
- window.extract_code_blocks,
- "Detect code snippets inside PDFs and plain text. If you train only from real source files, this can be turned off.",
- )
- window.preserve_indentation = QCheckBox("Preserve indentation")
- window.preserve_indentation.setChecked(True)
- window._tip(window.preserve_indentation, "Keep line breaks and indentation for code. This is important for Python and readable generated code.")
- window.instruction_samples = QCheckBox("Instruction-style samples")
- window.instruction_samples.setChecked(True)
- window._tip(window.instruction_samples, "Wrap code samples with simple instruction tags so the model sees code as task-oriented examples.")
- window.reasoning_sample_mode = QComboBox()
- window.reasoning_sample_mode.addItems(["Reasoning scaffold", "Detailed code reasoning", "No reasoning wrapper"])
- window.reasoning_sample_mode.setMaximumWidth(260)
- window._tip(
- window.reasoning_sample_mode,
- "Shapes code samples as task/reasoning/answer examples. This teaches response structure, not guaranteed deep reasoning by itwindow.",
- )
- window.instruction_samples.toggled.connect(window.reasoning_sample_mode.setEnabled)
-
- source_form.addRow("Source vault", window._path_row(window.input_dir, directory=True))
- source_form.addRow("Dataset core", window._path_row(window.dataset_dir, directory=True))
- source_pipeline_row = QWidget()
- source_pipeline_layout = QHBoxLayout(source_pipeline_row)
- source_pipeline_layout.setContentsMargins(0, 0, 0, 0)
- source_pipeline_layout.setSpacing(8)
- lanes_label = QLabel("Parallel lanes")
- lanes_label.setMinimumWidth(92)
- mode_label = QLabel("Prepare mode")
- mode_label.setMinimumWidth(92)
- source_pipeline_layout.addWidget(lanes_label)
- source_pipeline_layout.addWidget(window.max_workers, 1)
- source_pipeline_layout.addWidget(mode_label)
- source_pipeline_layout.addWidget(window.prepare_mode, 2)
- source_form.addRow("Pipeline", source_pipeline_row)
- source_form.addRow("Tokenizer training cap (GiB)", window.tokenizer_training_max_gb)
-
- source_options_row = QWidget()
- source_options_layout = QHBoxLayout(source_options_row)
- source_options_layout.setContentsMargins(0, 0, 0, 0)
- source_options_layout.setSpacing(14)
- source_options_layout.addWidget(window.code_training_mode)
- source_options_layout.addWidget(window.include_source_code)
- source_options_layout.addStretch(1)
- source_form.addRow("Options", source_options_row)
-
- tokenizer_form.addRow("Auto vocabulary", window.auto_vocab)
- tokenizer_form.addRow("Manual vocabulary", window.manual_vocab_size)
- tokenizer_form.addRow("Selected vocab", window.auto_vocab_label)
- window.tokenizer_path_row = window._path_row(window.tokenizer_path, directory=False, file_filter="Tokenizer JSON (*.json);;All files (*)")
- window.tokenizer_path_row.setEnabled(False)
- tokenizer_form.addRow("Tokenizer policy", window.tokenizer_strategy)
- tokenizer_form.addRow("Import tokenizer", window.tokenizer_path_row)
- tokenizer_form.addRow("Min frequency", window.min_frequency)
- tokenizer_form.addRow("Context window", window.context_length)
- tokenizer_form.addRow("Validation split", window.validation_split)
- tokenizer_form.addRow("", window.include_prose)
- tokenizer_form.addRow("", window.extract_code_blocks)
- tokenizer_form.addRow("", window.preserve_indentation)
- tokenizer_form.addRow("", window.instruction_samples)
- tokenizer_form.addRow("Reasoning samples", window.reasoning_sample_mode)
- source_card = window._card("SOURCE ARRAY", source_form)
- source_card.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Maximum)
- tokenizer_card = window._card("TOKENIZER CORE", tokenizer_form)
- tokenizer_card.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Maximum)
- left_column.addWidget(source_card, 0)
- right_column.addWidget(tokenizer_card, 0)
-
- window.dataset_mix_chart = DatasetBarChartWidget("Dataset Composition", "Percent")
- window.dataset_sequence_chart = DatasetBarChartWidget("Token Distribution", "Tokens")
- stats_grid = QGridLayout()
- stats_grid.setHorizontalSpacing(8)
- stats_grid.setVerticalSpacing(8)
- stats_grid.addWidget(window.dataset_mix_chart, 0, 0)
- stats_grid.addWidget(window.dataset_sequence_chart, 0, 1)
- stats_grid.setColumnStretch(0, 1)
- stats_grid.setColumnStretch(1, 1)
- stats_card = window._card("DATASET STATISTICS", stats_grid)
- stats_card.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Maximum)
- right_column.addWidget(stats_card, 0)
-
- window.dataset_advisor = QTextEdit()
- window.dataset_advisor.setReadOnly(True)
- window.dataset_advisor.setMinimumHeight(210)
- window.dataset_advisor.setPlainText("Run Preview Dataset to get cleanup suggestions.")
- window._tip(window.dataset_advisor, "Actionable dataset cleanup advice from preview quality checks.")
- advisor_layout = QVBoxLayout()
- advisor_layout.addWidget(window.dataset_advisor)
- advisor_card = window._card("DATASET ADVISOR", advisor_layout)
- advisor_card.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
- right_column.addWidget(advisor_card, 1)
-
- window.health_check_button = QPushButton("Check Health")
- window._tip(window.health_check_button, "Validate source, dataset, model, export, GGUF, and hardware readiness before long work.")
- window.health_check_button.clicked.connect(window.check_project_health)
- window.health_check_button.setMaximumWidth(160)
- window.preview_dataset_button = QPushButton("Preview Dataset")
- window._tip(window.preview_dataset_button, "Scan source files and show dataset quality plus sample text/code snippets without preparing tokens.")
- window.preview_dataset_button.clicked.connect(window.preview_dataset)
- window.preview_dataset_button.setMaximumWidth(180)
- window.prepare_button = QPushButton("Prepare Dataset")
- window._tip(window.prepare_button, "Read source files, clean text, train tokenizer, split tokens, and save the dataset project.")
- window.prepare_button.clicked.connect(window.prepare_dataset)
- window.prepare_button.setMaximumWidth(320)
- window.stop_dataset_button = QPushButton("Stop")
- window.stop_dataset_button.setEnabled(False)
- window.stop_dataset_button.setMaximumWidth(120)
- window.stop_dataset_button.clicked.connect(window.stop_active_task)
- window._tip(window.stop_dataset_button, "Request a graceful stop for dataset preparation.")
-
- window.dataset_log = QTextEdit()
- window.dataset_log.setReadOnly(True)
- window.dataset_log.document().setMaximumBlockCount(1200)
- window.dataset_log.setMinimumHeight(260)
- window.dataset_log.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
- log_layout = QVBoxLayout()
- log_layout.addWidget(window.dataset_log, 1)
- left_column.addWidget(window._card("INGEST TELEMETRY", log_layout), 1)
- action_row = QHBoxLayout()
- action_row.setSpacing(10)
- action_row.addWidget(window.health_check_button)
- action_row.addWidget(window.preview_dataset_button)
- action_row.addWidget(window.prepare_button, 1)
- action_row.addWidget(window.stop_dataset_button)
- action_row.addStretch(1)
- right_column.addLayout(action_row)
- right_column.addStretch(1)
- layout.addLayout(ingestion_body, 1)
-
- window.dataset_progress = window._thin_progress()
- outer.addWidget(window.dataset_progress)
- window._update_online_dataset_stage_controls()
- return page
\ No newline at end of file
diff --git a/llm_trainer/ui/tabs/export_tab.py b/llm_trainer/ui/tabs/export_tab.py
deleted file mode 100644
index 6e02eaf..0000000
--- a/llm_trainer/ui/tabs/export_tab.py
+++ /dev/null
@@ -1,118 +0,0 @@
-from __future__ import annotations
-
-from pathlib import Path
-
-from PySide6.QtCore import Qt
-from PySide6.QtWidgets import (
- QComboBox,
- QFormLayout,
- QHBoxLayout,
- QLineEdit,
- QPushButton,
- QScrollArea,
- QSizePolicy,
- QTextEdit,
- QVBoxLayout,
- QWidget,
-)
-
-
-def build_export_tab(window) -> QWidget:
- """Build the export page.
-
- Returns:
- Export page widget.
- """
-
- page = window._panel()
- outer = QVBoxLayout(page)
- outer.setContentsMargins(0, 0, 0, 0)
- outer.setSpacing(0)
- scroll = QScrollArea()
- scroll.setObjectName("PageScroll")
- scroll.setWidgetResizable(True)
- scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
- content = QWidget()
- content.setObjectName("Panel")
- layout = QVBoxLayout(content)
- layout.setContentsMargins(18, 18, 18, 10)
- layout.setSpacing(10)
- scroll.setWidget(content)
- outer.addWidget(scroll, 1)
- layout.addWidget(window._page_title("Export Bay"))
-
- form = QFormLayout()
- window._configure_form(form)
- window.export_model_dir = QLineEdit(str(Path.cwd() / "runs" / "model"))
- window._tip(window.export_model_dir, "Trained model folder containing final_model.pt and tokenizer.json.")
- window.export_dir = QLineEdit(str(Path.cwd() / "runs" / "export"))
- window._tip(window.export_dir, "Folder where export bundles or quantized checkpoints are written.")
- window.quant_mode = QComboBox()
- window.quant_mode.addItems(["FP16 checkpoint", "GGUF Q8_0 (planned)", "GGUF Q4_K_M (planned)", "GGUF Q5_K_M (planned)"])
- window.quant_mode.setMaximumWidth(260)
- window._tip(window.quant_mode, "Quantization target. FP16 reduces checkpoint size now; GGUF modes are planned for llama.cpp export.")
- window.llama_cpp_dir = QLineEdit()
- window._tip(window.llama_cpp_dir, "Local llama.cpp checkout folder containing convert_hf_to_gguf.py. This is not the GGUF output folder.")
- window.gguf_output_path = QLineEdit(str(Path.cwd() / "runs" / "export" / "model.gguf"))
- window._tip(window.gguf_output_path, "Destination GGUF file. Requires an HF-compatible hf_model folder in the model core.")
- window.gguf_outtype = QComboBox()
- window.gguf_outtype.addItems(["f16", "f32", "bf16", "q8_0"])
- window.gguf_outtype.setMaximumWidth(260)
- window._tip(window.gguf_outtype, "llama.cpp converter outtype. f16 is the usual starting point.")
- form.addRow("Model core", window._path_row(window.export_model_dir, directory=True))
- form.addRow("Output bay", window._path_row(window.export_dir, directory=True))
- form.addRow("Quantization", window.quant_mode)
- form.addRow("llama.cpp", window._path_row(window.llama_cpp_dir, directory=True))
- form.addRow("GGUF output", window._path_row(window.gguf_output_path, directory=False, file_filter="GGUF models (*.gguf);;All files (*)"))
- form.addRow("GGUF outtype", window.gguf_outtype)
- layout.addWidget(window._card("ARTIFACT CONFIGURATION", form))
-
- row = QHBoxLayout()
- row.setSpacing(10)
- bundle_button = QPushButton("Create Bundle")
- window._tip(bundle_button, "Copy final model, tokenizer, and summary into a portable export folder.")
- bundle_button.clicked.connect(window.create_bundle)
- quant_button = QPushButton("Quantize Model")
- window._tip(quant_button, "Create a smaller FP16 checkpoint for inference or later conversion workflows.")
- quant_button.clicked.connect(window.quantize_model)
- hf_button = QPushButton("Export HF Package")
- window._tip(hf_button, "Create model_core/hf_model with config, weights, tokenizer, lineage, and README.")
- hf_button.clicked.connect(window.export_hf_package)
- llama_button = QPushButton("Export Llama Adapter")
- window._tip(llama_button, "Export Llama-format weights only when the checkpoint uses RoPE, RMSNorm, SwiGLU, no bias, and full attention.")
- llama_button.clicked.connect(window.export_llama_adapter)
- window.gguf_convert_button = QPushButton("Convert HF to GGUF")
- window._tip(window.gguf_convert_button, "Run llama.cpp convert_hf_to_gguf.py for model_core/hf_model when the architecture is supported by llama.cpp.")
- window.gguf_convert_button.clicked.connect(window.convert_hf_to_gguf)
- bundle_button.setMaximumWidth(220)
- quant_button.setMaximumWidth(220)
- hf_button.setMaximumWidth(220)
- llama_button.setMaximumWidth(220)
- window.gguf_convert_button.setMaximumWidth(220)
- row.addWidget(bundle_button)
- row.addWidget(quant_button)
- row.addWidget(hf_button)
- row.addWidget(llama_button)
- row.addWidget(window.gguf_convert_button)
- row.addStretch(1)
-
- window.export_log = QTextEdit()
- window.export_log.setReadOnly(True)
- window.export_log.setMinimumHeight(320)
- window.export_log.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
- window.export_log.setPlainText(
- "Export options:\n"
- "- Bundle copies final_model.pt, tokenizer.json, and training_summary.json.\n"
- "- HF package writes model_core/hf_model for portable MicroGPT loading.\n"
- "- FP16 checkpoint quantization works now.\n"
- "- GGUF conversion uses llama.cpp when model_core/hf_model exists.\n"
- "- Native MicroGPT checkpoints are not written as fake GGUF files.\n"
- )
- export_log_layout = QVBoxLayout()
- export_log_layout.addWidget(window.export_log, 1)
- layout.addWidget(window._card("EXPORT TELEMETRY", export_log_layout), 1)
- layout.addLayout(row)
-
- window.export_progress = window._thin_progress()
- outer.addWidget(window.export_progress)
- return page
diff --git a/llm_trainer/ui/tabs/fine_tuning_tab.py b/llm_trainer/ui/tabs/fine_tuning_tab.py
deleted file mode 100644
index 755c2e2..0000000
--- a/llm_trainer/ui/tabs/fine_tuning_tab.py
+++ /dev/null
@@ -1,255 +0,0 @@
-from __future__ import annotations
-
-from PySide6.QtCore import Qt, QTimer
-from PySide6.QtWidgets import (
- QComboBox,
- QFormLayout,
- QGridLayout,
- QHBoxLayout,
- QLabel,
- QLineEdit,
- QPushButton,
- QScrollArea,
- QSizePolicy,
- QTextEdit,
- QVBoxLayout,
- QWidget,
-)
-
-
-def build_fine_tuning_tab(window) -> QWidget:
- """Build the dedicated fine-tuning page.
-
- Returns:
- Fine-tuning page widget.
- """
-
- page = window._panel()
- outer = QVBoxLayout(page)
- outer.setContentsMargins(0, 0, 0, 0)
- outer.setSpacing(0)
- scroll = QScrollArea()
- scroll.setObjectName("PageScroll")
- scroll.setWidgetResizable(True)
- scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
- content = QWidget()
- content.setObjectName("Panel")
- layout = QVBoxLayout(content)
- layout.setContentsMargins(18, 18, 18, 10)
- layout.setSpacing(10)
- scroll.setWidget(content)
- outer.addWidget(scroll, 1)
- layout.addWidget(window._page_title("Fine-Tuning Lab"))
-
- body = QHBoxLayout()
- body.setSpacing(12)
- left_zone = QVBoxLayout()
- left_zone.setSpacing(10)
- right_zone = QVBoxLayout()
- right_zone.setSpacing(10)
- body.addLayout(left_zone, 1)
- body.addLayout(right_zone, 1)
- layout.addLayout(body, 1)
-
- mode_form = QFormLayout()
- window._configure_form(mode_form)
- window.training_mode = QComboBox()
- window.training_mode.addItems(["Instruction fine-tune", "Conversation fine-tune", "Code fine-tune", "Fine-tune checkpoint"])
- window.training_mode.setMaximumWidth(300)
- window.training_mode.currentTextChanged.connect(window._update_training_mode_controls)
- window.training_mode.currentTextChanged.connect(window.refresh_fine_tune_workflow)
- window.training_mode.currentTextChanged.connect(window._refresh_fine_tune_default_output)
- window._tip(
- window.training_mode,
- "Instruction tunes request-following, conversation tunes chat behavior, code tunes programming help, generic fine-tune adapts domain data.",
- )
- window.fine_tune_checkpoint = QLineEdit()
- window._tip(window.fine_tune_checkpoint, "Base MicroGPT checkpoint used for fine-tuning. Must match tokenizer and model architecture.")
- # Architecture fields (n_embd, n_layer, etc.) are shared with the AI tab
- # and are not automatically kept in sync with whichever checkpoint is
- # selected here. Without this, a mismatched leftover architecture (e.g.
- # a "Tiny" preset still selected from an earlier session) either fails
- # the resume-compatibility check outright, or -- worse -- if the
- # mismatch happens to pass, training silently starts from a randomly
- # initialized model that only coincidentally matches the same shapes,
- # rather than actually continuing from the selected base checkpoint.
- # Syncing on every change to this field (not just via the separate
- # "Apply Recommended" button) makes the common path safe by default.
- window.fine_tune_checkpoint.textChanged.connect(lambda _text: window._sync_architecture_from_fine_tune_base())
- window.fine_tune_output_dir = QLineEdit()
- window._tip(
- window.fine_tune_output_dir,
- "Separate folder where fine-tuned checkpoints, adapters, telemetry, and final tuned model are saved.",
- )
- window.peft_method = QComboBox()
- window.peft_method.addItems(["Full fine-tune", "LoRA adapters"])
- window.peft_method.setMaximumWidth(300)
- window.peft_method.currentTextChanged.connect(window._update_training_mode_controls)
- window._tip(window.peft_method, "Parameter-efficient fine-tuning method. LoRA trains small adapters while freezing the base model.")
- window.lora_rank = window._spin(1, 256, 8)
- window._tip(window.lora_rank, "LoRA rank. Higher values increase adapter capacity and adapter size.")
- window.lora_alpha = window._double_spin(1.0, 512.0, 16.0, 1.0, 1)
- window._tip(window.lora_alpha, "LoRA alpha scaling. Common default is 2x the rank.")
- window.lora_dropout = window._double_spin(0.0, 0.9, 0.05, 0.01, 3)
- window._tip(window.lora_dropout, "Dropout used only inside LoRA adapters.")
- window.lora_targets = QComboBox()
- window.lora_targets.addItems(["Attention projections", "MLP projections", "Attention + MLP"])
- window.lora_targets.setMaximumWidth(300)
- window._tip(window.lora_targets, "Modules where LoRA adapters are attached.")
- window.fine_tune_check_button = QPushButton("Check Fine-tune")
- window.fine_tune_check_button.setMaximumWidth(180)
- window.fine_tune_check_button.clicked.connect(window.preview_fine_tune_compatibility)
- window._tip(window.fine_tune_check_button, "Inspect whether the base checkpoint can be used for fine-tuning.")
- window.fine_tune_dataset_status = QLabel("Dataset: not checked")
- window.fine_tune_dataset_status.setObjectName("Metric")
- window._tip(
- window.fine_tune_dataset_status,
- "Shows whether the prepared dataset purpose matches this fine-tune workflow.",
- )
- window.fine_tune_refresh_button = QPushButton("Refresh Dataset Fit")
- window.fine_tune_refresh_button.setMaximumWidth(190)
- window.fine_tune_refresh_button.clicked.connect(window.refresh_fine_tune_workflow)
- window._tip(window.fine_tune_refresh_button, "Re-read dataset_summary.json and check whether this dataset fits the fine-tune type.")
- window.apply_lora_preset_button = QPushButton("Apply Recommended LoRA")
- window.apply_lora_preset_button.setMaximumWidth(220)
- window.apply_lora_preset_button.clicked.connect(window.apply_recommended_fine_tune_settings)
- window._tip(window.apply_lora_preset_button, "Apply conservative LoRA, learning-rate, and clipping defaults for the selected fine-tune type.")
-
- mode_form.addRow("Fine-tune type", window.training_mode)
- mode_form.addRow("Base model", window._path_row(window.fine_tune_checkpoint, directory=False))
- mode_form.addRow("Fine-tune output", window._path_row(window.fine_tune_output_dir, directory=True))
- mode_form.addRow("PEFT", window.peft_method)
- mode_form.addRow("LoRA rank", window.lora_rank)
- mode_form.addRow("LoRA alpha", window.lora_alpha)
- mode_form.addRow("LoRA dropout", window.lora_dropout)
- mode_form.addRow("LoRA target", window.lora_targets)
- mode_form.addRow("Dataset fit", window.fine_tune_dataset_status)
- mode_form.addRow("", window.fine_tune_refresh_button)
- mode_form.addRow("", window.apply_lora_preset_button)
- mode_form.addRow("", window.fine_tune_check_button)
- left_zone.addWidget(window._card("ADAPTATION CONTROL", mode_form), 0)
-
- runtime_form = QFormLayout()
- window._configure_form(runtime_form)
- window.fine_tune_launch_target = QComboBox()
- window.fine_tune_launch_target.addItems(["Local machine", "Remote workers", "RunPod cloud"])
- window.fine_tune_launch_target.setMaximumWidth(300)
- window._tip(
- window.fine_tune_launch_target,
- "Local runs fine-tuning on this computer. Remote queues it for workers. RunPod creates a cloud GPU worker automatically.",
- )
- window.fine_tune_runtime_hint = QLabel("Uses AI tab device, precision, resume, and checkpoint settings.")
- window.fine_tune_runtime_hint.setObjectName("Metric")
- window.fine_tune_runtime_hint.setWordWrap(True)
- runtime_form.addRow("Launch", window.fine_tune_launch_target)
- runtime_form.addRow("Settings", window.fine_tune_runtime_hint)
- left_zone.addWidget(window._card("FINE-TUNE RUNTIME", runtime_form), 0)
-
- metrics_grid = QGridLayout()
- metrics_grid.setHorizontalSpacing(8)
- metrics_grid.setVerticalSpacing(8)
- window.fine_tune_eta_metric = window._metric_chip("ETA: -", "Estimated time remaining for the fine-tune run.")
- window.fine_tune_epoch_metric = window._metric_chip("Epoch: -", "Current fine-tune epoch and total epochs.")
- window.fine_tune_step_metric = window._metric_chip("Step: -", "Current optimizer step and total planned steps.")
- window.fine_tune_loss_metric = window._metric_chip("Train loss: -", "Latest fine-tune training loss.")
- window.fine_tune_val_metric = window._metric_chip("Val loss: -", "Latest fine-tune validation loss.")
- window.fine_tune_lr_metric = window._metric_chip("LR: -", "Current fine-tune learning rate.")
- window.fine_tune_speed_metric = window._metric_chip("Speed: -", "Fine-tune token throughput.")
- window.fine_tune_grad_metric = window._metric_chip("Grad: -", "Current fine-tune gradient norm.")
- for index, metric in enumerate((
- window.fine_tune_eta_metric,
- window.fine_tune_epoch_metric,
- window.fine_tune_step_metric,
- window.fine_tune_loss_metric,
- window.fine_tune_val_metric,
- window.fine_tune_lr_metric,
- window.fine_tune_speed_metric,
- window.fine_tune_grad_metric,
- )):
- metrics_grid.addWidget(metric, index // 2, index % 2)
- metrics_grid.setColumnStretch(0, 1)
- metrics_grid.setColumnStretch(1, 1)
- left_zone.addWidget(window._card("FINE-TUNE METRICS", metrics_grid), 0)
-
- builder_form = QFormLayout()
- window._configure_form(builder_form)
- window.fine_tune_dataset_builder_stage = QComboBox()
- window.fine_tune_dataset_builder_stage.addItems(["Instruction fine-tune", "Conversation fine-tune", "Code fine-tune"])
- window.fine_tune_dataset_builder_stage.setMaximumWidth(300)
- window.fine_tune_dataset_builder_stage.currentTextChanged.connect(window._refresh_fine_tune_default_output)
- window._tip(
- window.fine_tune_dataset_builder_stage,
- "Choose the fine-tune dataset type to configure in the Ingest tab.",
- )
- window.configure_fine_tune_dataset_button = QPushButton("Configure Ingest")
- window.configure_fine_tune_dataset_button.setMaximumWidth(180)
- window.configure_fine_tune_dataset_button.clicked.connect(window.configure_fine_tune_dataset_builder)
- window._tip(
- window.configure_fine_tune_dataset_button,
- "Switch to Ingest, set the matching dataset purpose, enable online datasets, and preselect a small starter dataset.",
- )
- builder_form.addRow("Dataset type", window.fine_tune_dataset_builder_stage)
- builder_form.addRow("", window.configure_fine_tune_dataset_button)
- left_zone.addWidget(window._card("FINE-TUNE DATASET BUILDER", builder_form), 0)
-
- guidance = QTextEdit()
- guidance.setReadOnly(True)
- guidance.setMaximumHeight(130)
- guidance.setPlainText(
- "Fine-tuning starts from a compatible checkpoint and uses the prepared dataset selected in AI.\n"
- "- Prepare Instruction fine-tune data for request-following behavior.\n"
- "- Prepare Conversation fine-tune data for chat behavior.\n"
- "- Prepare Code fine-tune data for programming help.\n"
- "- Reuse or import the base tokenizer so token IDs stay compatible.\n"
- "- LoRA is recommended for most experiments.\n"
- "- Stop saves a resumable checkpoint; Resume latest can continue it."
- )
- left_zone.addWidget(window._card("WORKFLOW", _single_widget_layout(guidance)), 0)
-
- window.fine_tune_preview = QTextEdit()
- window.fine_tune_preview.setReadOnly(True)
- window.fine_tune_preview.setMinimumHeight(160)
- window.fine_tune_preview.setText("No compatibility check has been run.")
- window._tip(window.fine_tune_preview, "Compatibility report for the selected base checkpoint.")
- right_zone.addWidget(window._card("FINE-TUNE COMPATIBILITY", _single_widget_layout(window.fine_tune_preview)), 0)
-
- window.fine_tune_log = QTextEdit()
- window.fine_tune_log.setReadOnly(True)
- window.fine_tune_log.document().setMaximumBlockCount(1500)
- window.fine_tune_log.setMinimumHeight(320)
- window.fine_tune_log.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
- right_zone.addWidget(window._card("FINE-TUNE TELEMETRY", _single_widget_layout(window.fine_tune_log)), 1)
-
- window.fine_tune_button = QPushButton("Start Fine-Tune")
- window.fine_tune_button.setMaximumWidth(220)
- window.fine_tune_button.clicked.connect(window.start_fine_tuning)
- window._tip(window.fine_tune_button, "Start fine-tuning from the selected compatible base checkpoint.")
- window.stop_fine_tune_button = QPushButton("Stop")
- window.stop_fine_tune_button.setEnabled(False)
- window.stop_fine_tune_button.setMaximumWidth(120)
- window.stop_fine_tune_button.clicked.connect(window.stop_active_task)
- action_row = QHBoxLayout()
- action_row.addWidget(window.fine_tune_button)
- action_row.addWidget(window.stop_fine_tune_button)
- action_row.addStretch(1)
- layout.addLayout(action_row)
-
- window.fine_tune_progress = window._thin_progress()
- outer.addWidget(window.fine_tune_progress)
- QTimer.singleShot(0, window._update_training_mode_controls)
- return page
-
-
-def _single_widget_layout(widget: QWidget) -> QVBoxLayout:
- """Wrap one widget in a vertical layout.
-
- Args:
- widget: Widget to place in a layout.
-
- Returns:
- Layout containing the widget.
- """
-
- layout = QVBoxLayout()
- layout.addWidget(widget, 1)
- return layout
\ No newline at end of file
diff --git a/llm_trainer/ui/tabs/job_manager_tab.py b/llm_trainer/ui/tabs/job_manager_tab.py
deleted file mode 100644
index 6c6d204..0000000
--- a/llm_trainer/ui/tabs/job_manager_tab.py
+++ /dev/null
@@ -1,291 +0,0 @@
-from __future__ import annotations
-
-from pathlib import Path
-
-from PySide6.QtCore import Qt
-from PySide6.QtWidgets import (
- QCheckBox,
- QComboBox,
- QFormLayout,
- QHBoxLayout,
- QHeaderView,
- QLabel,
- QLineEdit,
- QPushButton,
- QScrollArea,
- QTableWidget,
- QTableWidgetItem,
- QTextEdit,
- QVBoxLayout,
- QWidget,
-)
-
-
-def build_job_manager_tab(window) -> QWidget:
- """Build the distributed job manager page.
-
- Args:
- window: Main window instance.
-
- Returns:
- Job manager page widget.
- """
-
- page = window._panel()
- page_layout = QVBoxLayout(page)
- page_layout.setContentsMargins(18, 18, 18, 10)
- page_layout.setSpacing(8)
-
- title_row = QHBoxLayout()
- title_row.addWidget(window._page_title("Job Manager"))
- title_row.addStretch(1)
- window.job_refresh_button = QPushButton("Refresh")
- window.job_refresh_button.clicked.connect(window.refresh_job_manager_tab)
- window._tip(window.job_refresh_button, "Reload worker and job status from the coordinator state store.")
- window.job_stale_button = QPushButton("Mark Stale Offline")
- window.job_stale_button.clicked.connect(window.mark_stale_workers_offline)
- window._tip(window.job_stale_button, "Mark remote workers offline when they have not sent a heartbeat recently.")
- window.job_pause_button = QPushButton("Pause All")
- window.job_pause_button.clicked.connect(window.pause_all_managed_jobs)
- window._tip(window.job_pause_button, "Ask remote workers to pause active jobs and hold queued jobs.")
- window.job_resume_button = QPushButton("Resume All")
- window.job_resume_button.clicked.connect(window.resume_all_managed_jobs)
- window._tip(window.job_resume_button, "Return paused jobs to the queue.")
- window.job_stop_button = QPushButton("Stop All Jobs")
- window.job_stop_button.clicked.connect(window.stop_all_managed_jobs)
- window._tip(window.job_stop_button, "Request cooperative stop for every queued or active managed job.")
- for button in (
- window.job_refresh_button,
- window.job_stale_button,
- window.job_pause_button,
- window.job_resume_button,
- window.job_stop_button,
- ):
- button.setMaximumWidth(170)
- title_row.addWidget(button)
- page_layout.addLayout(title_row)
-
- summary = QHBoxLayout()
- window.job_worker_count_label = QLabel("Workers: -")
- window.job_worker_count_label.setObjectName("Metric")
- window.job_active_count_label = QLabel("Active jobs: -")
- window.job_active_count_label.setObjectName("Metric")
- window.job_queue_count_label = QLabel("Queued jobs: -")
- window.job_queue_count_label.setObjectName("Metric")
- window.job_db_label = QLabel("State DB: -")
- window.job_db_label.setObjectName("Metric")
- for label in (
- window.job_worker_count_label,
- window.job_active_count_label,
- window.job_queue_count_label,
- window.job_db_label,
- ):
- label.setTextInteractionFlags(Qt.TextSelectableByMouse)
- summary.addWidget(label)
- summary.addStretch(1)
- page_layout.addLayout(summary)
-
- scroll = QScrollArea()
- scroll.setObjectName("PageScroll")
- scroll.setWidgetResizable(True)
- scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
- content = QWidget()
- root = QVBoxLayout(content)
- root.setContentsMargins(0, 0, 0, 0)
- root.setSpacing(10)
- scroll.setWidget(content)
- page_layout.addWidget(scroll, 1)
-
- connection_form = QFormLayout()
- window._configure_form(connection_form)
- window.coordinator_host = QLineEdit("0.0.0.0")
- window._tip(window.coordinator_host, "Network address used by the coordinator API. Use 0.0.0.0 to accept workers from other machines.")
- window.coordinator_port = window._spin(1, 65535, 8765)
- window._tip(window.coordinator_port, "Coordinator API port. Remote workers connect to this port.")
- window.coordinator_artifact_root = QLineEdit(str(Path.home() / ".drunkenbot_ide" / "artifacts"))
- window._tip(window.coordinator_artifact_root, "Folder where job input bundles and uploaded worker result bundles are stored.")
- window.coordinator_public_url = QLineEdit("http://127.0.0.1:8765")
- window._tip(window.coordinator_public_url, "URL workers use to reach this coordinator. Use this machine's LAN/IP address for remote machines.")
- connection_form.addRow("Host", window.coordinator_host)
- connection_form.addRow("Port", window.coordinator_port)
- connection_form.addRow("Artifact root", window._path_row(window.coordinator_artifact_root, directory=True))
- connection_form.addRow("Worker URL", window.coordinator_public_url)
-
- connection_buttons = QHBoxLayout()
- window.coordinator_start_button = QPushButton("Start Coordinator")
- window.coordinator_start_button.clicked.connect(window.start_coordinator_server)
- window._tip(window.coordinator_start_button, "Start the HTTP coordinator so remote workers can register, claim jobs, download inputs, and upload outputs.")
- window.coordinator_stop_button = QPushButton("Stop Coordinator")
- window.coordinator_stop_button.setEnabled(False)
- window.coordinator_stop_button.clicked.connect(window.stop_coordinator_server)
- window._tip(window.coordinator_stop_button, "Stop the coordinator API. Running remote workers will lose connection until it starts again.")
- window.publish_remote_job_button = QPushButton("Publish Remote Job")
- window.publish_remote_job_button.clicked.connect(window.publish_remote_training_job)
- window._tip(window.publish_remote_job_button, "Bundle the current dataset/checkpoints and queue the current training settings for remote workers.")
- connection_buttons.addWidget(window.coordinator_start_button)
- connection_buttons.addWidget(window.coordinator_stop_button)
- connection_buttons.addWidget(window.publish_remote_job_button)
- connection_buttons.addStretch(1)
- connection_form.addRow("", connection_buttons)
- window.coordinator_status_label = QLabel("Coordinator: stopped")
- window.coordinator_status_label.setObjectName("Metric")
- connection_form.addRow("Status", window.coordinator_status_label)
- coordinator_card = window._card("COORDINATOR API / ARTIFACT SYNC", connection_form)
- coordinator_card.setMinimumHeight(230)
- root.addWidget(coordinator_card, 0)
-
- runpod_form = QFormLayout()
- window._configure_form(runpod_form)
- window.runpod_api_key = QLineEdit()
- window.runpod_api_key.setEchoMode(QLineEdit.Password)
- window._tip(window.runpod_api_key, "RunPod API key. Stored in runpod_config.json inside the project folder.")
- window.runpod_gpu_type = QComboBox()
- window.runpod_gpu_type.setEditable(True)
- window.runpod_gpu_type.addItems([
- "NVIDIA GeForce RTX 4090",
- "NVIDIA RTX A5000",
- "NVIDIA A40",
- "NVIDIA L40S",
- "NVIDIA A100 80GB PCIe",
- "NVIDIA H100 80GB HBM3",
- ])
- window._tip(window.runpod_gpu_type, "Preferred RunPod GPU type. Availability depends on RunPod capacity.")
- window.runpod_cloud_type = QComboBox()
- window.runpod_cloud_type.addItems(["COMMUNITY", "SECURE"])
- window._tip(window.runpod_cloud_type, "Community is usually cheaper. Secure is more controlled and often more predictable.")
- window.runpod_image = QLineEdit("runpod/pytorch:2.4.0-py3.11-cuda12.4.1-devel-ubuntu22.04")
- window._tip(window.runpod_image, "Docker image used for the cloud worker. The default PyTorch image installs this app's worker bundle at startup.")
- window.runpod_container_disk = window._spin(20, 500, 80)
- window._tip(window.runpod_container_disk, "Temporary container disk in GB. Larger values help with dependency install and cache.")
- window.runpod_volume_gb = window._spin(20, 1000, 40)
- window._tip(window.runpod_volume_gb, "Persistent Pod volume in GB mounted at /workspace.")
- window.runpod_min_ram = window._spin(8, 512, 16)
- window._tip(window.runpod_min_ram, "Minimum system RAM per GPU in GB.")
- window.runpod_min_vcpu = window._spin(2, 128, 4)
- window._tip(window.runpod_min_vcpu, "Minimum virtual CPUs per GPU.")
- window.runpod_spot = QCheckBox("Use interruptible cheaper pods")
- window.runpod_spot.setChecked(True)
- window._tip(window.runpod_spot, "Interruptible pods cost less but may be reclaimed. Checkpoints make recovery easier.")
- window.runpod_auto_terminate = QCheckBox("Exit worker after one job")
- window.runpod_auto_terminate.setChecked(True)
- window._tip(window.runpod_auto_terminate, "Worker claims one job and exits. Verify the Pod has stopped in RunPod when the job completes.")
- window.runpod_save_button = QPushButton("Save RunPod Settings")
- window.runpod_save_button.clicked.connect(window.save_runpod_settings)
- window._tip(window.runpod_save_button, "Save RunPod settings to runpod_config.json.")
- window.runpod_launch_button = QPushButton("Launch RunPod Worker")
- window.runpod_launch_button.clicked.connect(window.launch_runpod_worker_for_current_training)
- window._tip(window.runpod_launch_button, "Publish the current training job and create a RunPod cloud worker to claim it.")
- runpod_buttons = QHBoxLayout()
- runpod_buttons.addWidget(window.runpod_save_button)
- runpod_buttons.addWidget(window.runpod_launch_button)
- runpod_buttons.addStretch(1)
- runpod_form.addRow("API key", window.runpod_api_key)
- runpod_form.addRow("GPU", window.runpod_gpu_type)
- runpod_form.addRow("Cloud", window.runpod_cloud_type)
- runpod_form.addRow("Image", window.runpod_image)
- runpod_form.addRow("Disk / volume", _inline_widgets(window.runpod_container_disk, window.runpod_volume_gb))
- runpod_form.addRow("RAM / vCPU", _inline_widgets(window.runpod_min_ram, window.runpod_min_vcpu))
- runpod_form.addRow("", window.runpod_spot)
- runpod_form.addRow("", window.runpod_auto_terminate)
- runpod_form.addRow("", runpod_buttons)
- window.runpod_status_label = QLabel("RunPod: not configured")
- window.runpod_status_label.setObjectName("Metric")
- window.runpod_status_label.setWordWrap(True)
- runpod_form.addRow("Status", window.runpod_status_label)
- runpod_card = window._card("RUNPOD CLOUD GPU", runpod_form)
- runpod_card.setMinimumHeight(330)
- root.addWidget(runpod_card, 0)
-
- window.job_worker_table = _table(
- ["Worker", "Status", "Backend", "Device", "Last Seen", "Active Job", "CPU/RAM/GPU", "Labels"]
- )
- root.addWidget(window._card("WORKERS / CONNECTIONS", _table_layout(window.job_worker_table)), 0)
-
- window.job_table = _table(
- ["Job", "Stage", "Status", "Worker", "Backend", "Epoch", "Step", "Batch", "Layers", "Loss", "Speed", "Updated"]
- )
- root.addWidget(window._card("JOBS / TRAINING ASSIGNMENTS", _table_layout(window.job_table)), 0)
-
- window.job_manager_log = QTextEdit()
- window.job_manager_log.setReadOnly(True)
- window.job_manager_log.setMinimumHeight(120)
- root.addWidget(window._card("COORDINATOR TELEMETRY", _table_layout(window.job_manager_log)), 0)
-
- window.job_manager_progress = window._thin_progress()
- page_layout.addWidget(window.job_manager_progress)
- window.load_runpod_settings()
- window.refresh_job_manager_tab()
- return page
-
-
-def _table(headers: list[str]) -> QTableWidget:
- """Create a table widget.
-
- Args:
- headers: Column labels.
-
- Returns:
- Table widget.
- """
-
- table = QTableWidget(0, len(headers))
- table.setHorizontalHeaderLabels(headers)
- table.setAlternatingRowColors(True)
- table.setSelectionBehavior(QTableWidget.SelectRows)
- table.setEditTriggers(QTableWidget.NoEditTriggers)
- table.verticalHeader().setVisible(False)
- table.horizontalHeader().setStretchLastSection(True)
- table.horizontalHeader().setSectionResizeMode(QHeaderView.ResizeToContents)
- table.setMinimumHeight(170)
- return table
-
-
-def _table_layout(widget: QWidget) -> QVBoxLayout:
- """Wrap a widget in a layout.
-
- Args:
- widget: Widget to wrap.
-
- Returns:
- Layout containing the widget.
- """
-
- layout = QVBoxLayout()
- layout.addWidget(widget)
- return layout
-
-
-def _inline_widgets(*widgets: QWidget) -> QWidget:
- """Create a compact inline row.
-
- Args:
- widgets: Widgets to place side by side.
-
- Returns:
- Container widget.
- """
-
- holder = QWidget()
- row = QHBoxLayout(holder)
- row.setContentsMargins(0, 0, 0, 0)
- row.setSpacing(8)
- for widget in widgets:
- row.addWidget(widget)
- row.addStretch(1)
- return holder
-
-
-def set_table_rows(table: QTableWidget, rows: list[list[str]]) -> None:
- """Replace all table rows.
-
- Args:
- table: Table widget.
- rows: Row values.
- """
-
- table.setRowCount(len(rows))
- for row_index, row in enumerate(rows):
- for column_index, value in enumerate(row):
- item = QTableWidgetItem(value)
- item.setToolTip(value)
- table.setItem(row_index, column_index, item)
diff --git a/llm_trainer/ui/tabs/live_tab.py b/llm_trainer/ui/tabs/live_tab.py
deleted file mode 100644
index ce2abb2..0000000
--- a/llm_trainer/ui/tabs/live_tab.py
+++ /dev/null
@@ -1,229 +0,0 @@
-from __future__ import annotations
-
-from PySide6.QtCore import Qt
-from PySide6.QtWidgets import (
- QGridLayout,
- QHBoxLayout,
- QLabel,
- QPushButton,
- QScrollArea,
- QSlider,
- QVBoxLayout,
- QWidget,
-)
-
-from llm_trainer.ui.charts import LossChartWidget
-from llm_trainer.ui.live_widgets import (
- LiveDistributionWidget,
- LiveGradientFlowWidget,
- LiveHeatmapWidget,
- LiveHistogramWidget,
- ModelFlowWidget,
-)
-
-
-def build_live_training_tab(window) -> QWidget:
- """Build the live training tracker page.
-
- Returns:
- Live training tracker page widget.
- """
-
- page = window._panel()
- outer = QVBoxLayout(page)
- outer.setContentsMargins(0, 0, 0, 0)
- outer.setSpacing(0)
- scroll = QScrollArea()
- scroll.setObjectName("PageScroll")
- scroll.setWidgetResizable(True)
- scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
- content = QWidget()
- content.setObjectName("Panel")
- layout = QVBoxLayout(content)
- layout.setContentsMargins(18, 18, 18, 10)
- layout.setSpacing(10)
- scroll.setWidget(content)
- outer.addWidget(scroll, 1)
-
- header = QHBoxLayout()
- title = window._page_title("Model Training Live")
- live_badge = QLabel("[*] LIVE")
- live_badge.setObjectName("Metric")
- window.live_epoch_metric = window._metric_chip("Epoch: -", "Current epoch and total epochs.")
- window.live_step_metric = window._metric_chip("Step: -", "Current optimizer step and total planned steps.")
- window.live_tokens_metric = window._metric_chip("Tokens/sec: -", "Current token throughput.")
- window.live_loss_metric = window._metric_chip("Loss: -", "Latest training loss.")
- window.live_lr_metric = window._metric_chip("LR: -", "Current learning rate.")
- window.live_data_metric = window._metric_chip("Data: -", "Estimated percentage of planned training steps completed.")
- for chip in (
- window.live_epoch_metric,
- window.live_step_metric,
- window.live_tokens_metric,
- window.live_loss_metric,
- window.live_lr_metric,
- window.live_data_metric,
- ):
- chip.setMaximumWidth(180)
- header.addWidget(title)
- header.addWidget(live_badge)
- header.addSpacing(10)
- header.addWidget(window.live_epoch_metric)
- header.addWidget(window.live_step_metric)
- header.addWidget(window.live_tokens_metric)
- header.addWidget(window.live_loss_metric)
- header.addWidget(window.live_lr_metric)
- header.addWidget(window.live_data_metric)
- header.addStretch(1)
- layout.addLayout(header)
-
- body = QGridLayout()
- body.setSpacing(10)
- layout.addLayout(body, 1)
-
- left_column = QVBoxLayout()
- left_column.setSpacing(10)
-
- status_layout = QVBoxLayout()
- window.live_model_status = QLabel("Model: Transformer decoder")
- window.live_layer_status = QLabel("Layers: -")
- window.live_head_status = QLabel("Heads: -")
- window.live_hidden_status = QLabel("Hidden size: -")
- window.live_batch_status = QLabel("Batch size: -")
- window.live_context_status = QLabel("Context: -")
- for label in (
- window.live_model_status,
- window.live_layer_status,
- window.live_head_status,
- window.live_hidden_status,
- window.live_batch_status,
- window.live_context_status,
- ):
- status_layout.addWidget(label)
- left_column.addWidget(window._card("TRAINING STATUS", status_layout), 0)
-
- hardware_layout = QVBoxLayout()
- hardware_layout.setSpacing(7)
- window.live_device_status = QLabel("Device: -")
- window.live_cpu_label = QLabel("CPU: -")
- window.live_cpu_bar = window._hardware_meter("CPU")
- window.live_gpu_label = QLabel("GPU memory: -")
- window.live_gpu_bar = window._hardware_meter("GPU memory")
- window.live_vram_label = QLabel("VRAM reserved: -")
- window.live_vram_bar = window._hardware_meter("VRAM reserved")
- window.live_ram_label = QLabel("System RAM: -")
- window.live_ram_bar = window._hardware_meter("System RAM")
- window.live_worker_status = QLabel("CPU workers: -")
- window.hardware_meter_labels[id(window.live_cpu_bar)] = window.live_cpu_label
- window.hardware_meter_labels[id(window.live_gpu_bar)] = window.live_gpu_label
- window.hardware_meter_labels[id(window.live_vram_bar)] = window.live_vram_label
- window.hardware_meter_labels[id(window.live_ram_bar)] = window.live_ram_label
- hardware_layout.addWidget(window.live_device_status)
- for label, meter in (
- (window.live_cpu_label, window.live_cpu_bar),
- (window.live_gpu_label, window.live_gpu_bar),
- (window.live_vram_label, window.live_vram_bar),
- (window.live_ram_label, window.live_ram_bar),
- ):
- hardware_layout.addWidget(label)
- hardware_layout.addWidget(meter)
- hardware_layout.addWidget(window.live_worker_status)
- hardware_layout.addStretch(1)
- left_column.addWidget(window._card("HARDWARE", hardware_layout), 1)
-
- center_column = QVBoxLayout()
- center_column.setSpacing(10)
- flow_column = QVBoxLayout()
- flow_column.setContentsMargins(0, 0, 0, 0)
- flow_column.setSpacing(4)
- window.live_flow = ModelFlowWidget()
- window._tip(window.live_flow, "Visual summary of forward and backward flow through the configured transformer layers.")
- window.live_sample_text = QLabel("Training text: -")
- window.live_sample_text.setObjectName("TrainingSampleLine")
- window.live_sample_text.setWordWrap(False)
- window.live_sample_text.setTextInteractionFlags(Qt.TextSelectableByMouse)
- window._tip(window.live_sample_text, "A compact preview of the current token window being used for training.")
- flow_column.addWidget(window.live_flow, 1)
- flow_column.addWidget(window.live_sample_text, 0)
-
- window.loss_chart = LossChartWidget("Train loss", "Validation loss", "Loss chart will appear during training", "Loss curve", "Cross entropy loss")
- window._tip(window.loss_chart, "Live training and validation loss. Falling values usually mean the model is learning.")
- window.optimization_chart = LossChartWidget("Learning rate", "Gradient norm", "Learning rate and gradient norm will appear during training", "Learning rate / gradient", "Value")
- window._tip(window.optimization_chart, "Learning rate and gradient norm. Watch for unstable spikes or gradients collapsing toward zero.")
- window.stability_chart = LossChartWidget("Weight norm", "Update ratio", "Weight norm and update ratio will appear during training", "Parameter stability", "Value")
- window._tip(window.stability_chart, "Weight norm and parameter update ratio. Large update ratios can destabilize training; tiny ratios can stall learning.")
- window.throughput_chart = LossChartWidget("Tokens/sec", "Samples/sec", "Throughput will appear during training", "Throughput", "Rate")
- window._tip(window.throughput_chart, "Training speed measured as tokens/sec and samples/sec.")
- window.memory_chart = LossChartWidget("VRAM allocated", "VRAM reserved", "VRAM usage will appear during CUDA training", "GPU memory", "GB")
- window._tip(window.memory_chart, "CUDA memory usage in GB. Helps diagnose memory bottlenecks.")
- charts_grid = QGridLayout()
- charts_grid.setHorizontalSpacing(8)
- charts_grid.setVerticalSpacing(8)
- charts_grid.addWidget(window.loss_chart, 0, 0)
- charts_grid.addWidget(window.optimization_chart, 0, 1)
- charts_grid.addWidget(window.stability_chart, 1, 0)
- charts_grid.addWidget(window.throughput_chart, 1, 1)
- charts_grid.addWidget(window.memory_chart, 2, 0, 1, 2)
- charts_grid.setColumnStretch(0, 1)
- charts_grid.setColumnStretch(1, 1)
- center_column.addWidget(window._card("TRAINING GRAPHS", charts_grid), 3)
-
- timeline_layout = QHBoxLayout()
- timeline_layout.setSpacing(8)
- window.live_time_slider = QSlider(Qt.Horizontal)
- window.live_time_slider.setRange(0, 0)
- window.live_time_slider.setValue(0)
- window.live_time_slider.sliderPressed.connect(window._begin_live_scrub)
- window.live_time_slider.sliderReleased.connect(window._end_live_scrub)
- window.live_time_slider.valueChanged.connect(window._scrub_live_timeline)
- window._tip(window.live_time_slider, "Drag to replay training graphs from recorded SQLite telemetry.")
- window.live_timeline_label = QLabel("Timeline: live")
- window.live_timeline_label.setObjectName("Metric")
- live_button = QPushButton("Live")
- live_button.setMaximumWidth(80)
- live_button.clicked.connect(window._jump_live_timeline_to_latest)
- window._tip(live_button, "Return the tracker to the latest live metrics.")
- timeline_layout.addWidget(QLabel("Time"))
- timeline_layout.addWidget(window.live_time_slider, 1)
- timeline_layout.addWidget(window.live_timeline_label)
- timeline_layout.addWidget(live_button)
- center_column.addWidget(window._card("TIMELINE", timeline_layout), 0)
-
- right_column = QVBoxLayout()
- right_column.setSpacing(10)
- window.live_prediction_chart = LiveDistributionWidget()
- window.live_attention_chart = LiveHeatmapWidget()
- window.live_activation_chart = LiveHistogramWidget()
- window.live_gradient_chart = LiveGradientFlowWidget()
- right_column.addWidget(window._card("PREDICTION DISTRIBUTION", single_widget_layout(window.live_prediction_chart)), 1)
- right_column.addWidget(window._card("ATTENTION", single_widget_layout(window.live_attention_chart)), 1)
- right_column.addWidget(window._card("ACTIVATION", single_widget_layout(window.live_activation_chart)), 1)
- right_column.addWidget(window._card("GRADIENT FLOW", single_widget_layout(window.live_gradient_chart)), 1)
-
- body.addLayout(flow_column, 0, 0, 1, 2)
- body.addLayout(left_column, 1, 0)
- body.addLayout(center_column, 1, 1)
- body.addLayout(right_column, 0, 2, 2, 1)
- body.setColumnStretch(0, 1)
- body.setColumnStretch(1, 4)
- body.setColumnStretch(2, 1)
- body.setRowStretch(0, 2)
- body.setRowStretch(1, 3)
-
- window.live_progress = window._thin_progress()
- outer.addWidget(window.live_progress)
- return page
-
-def single_widget_layout(widget: QWidget) -> QVBoxLayout:
- """Wrap a single widget in a vertical layout.
-
- Args:
- widget: Widget to place in the layout.
-
- Returns:
- Layout containing the widget.
- """
-
- layout = QVBoxLayout()
- layout.setContentsMargins(0, 0, 0, 0)
- layout.addWidget(widget, 1)
- return layout
diff --git a/llm_trainer/ui/tabs/training_tab.py b/llm_trainer/ui/tabs/training_tab.py
deleted file mode 100644
index 8d69bb9..0000000
--- a/llm_trainer/ui/tabs/training_tab.py
+++ /dev/null
@@ -1,428 +0,0 @@
-from __future__ import annotations
-
-from pathlib import Path
-
-from PySide6.QtCore import Qt, QTimer
-from PySide6.QtWidgets import (
- QCheckBox,
- QComboBox,
- QFormLayout,
- QGridLayout,
- QHBoxLayout,
- QLabel,
- QLineEdit,
- QPushButton,
- QScrollArea,
- QSizePolicy,
- QTextEdit,
- QVBoxLayout,
- QWidget,
-)
-
-
-def build_training_tab(window) -> QWidget:
- """Build the training configuration page.
-
- Returns:
- Training page widget.
- """
-
- page = window._panel()
- outer = QVBoxLayout(page)
- outer.setContentsMargins(0, 0, 0, 0)
- outer.setSpacing(0)
- scroll = QScrollArea()
- scroll.setObjectName("PageScroll")
- scroll.setWidgetResizable(True)
- scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
- content = QWidget()
- content.setObjectName("Panel")
- layout = QVBoxLayout(content)
- layout.setContentsMargins(18, 18, 18, 10)
- layout.setSpacing(10)
- scroll.setWidget(content)
- outer.addWidget(scroll, 1)
- layout.addWidget(window._page_title("Neural Forge"))
- training_body = QHBoxLayout()
- training_body.setSpacing(12)
- left_zone = QVBoxLayout()
- left_zone.setSpacing(10)
- right_zone = QVBoxLayout()
- right_zone.setSpacing(10)
- training_body.addLayout(left_zone, 2)
- training_body.addLayout(right_zone, 1)
- layout.addLayout(training_body, 1)
-
- left = QFormLayout()
- window._configure_form(left)
- window.train_data_dir = QLineEdit(str(Path.cwd() / "runs" / "dataset"))
- window._tip(window.train_data_dir, "Prepared dataset folder containing tokenizer.json and train/validation token files.")
- window.model_dir = QLineEdit(str(Path.cwd() / "runs" / "model"))
- window._tip(window.model_dir, "Folder where checkpoints, final model, tokenizer copy, and training summary are saved.")
- window.preset = QComboBox()
- window.preset.addItems(["Tiny", "Small", "Custom"])
- window.preset.setMaximumWidth(260)
- window._tip(window.preset, "Architecture preset. Tiny is faster; Small has more capacity but needs more memory and training data.")
- window.preset.currentTextChanged.connect(window._apply_preset)
- window.architecture_style = QComboBox()
- window.architecture_style.addItems(["Classic GPT", "Llama-like"])
- window.architecture_style.setMaximumWidth(260)
- window._tip(
- window.architecture_style,
- "Classic uses learned positions, LayerNorm, and GELU. Llama-like uses RoPE, RMSNorm, and SwiGLU.",
- )
- window.rope_theta = window._double_spin(1.0, 10_000_000.0, 10000.0, 1000.0, 1)
- window._tip(
- window.rope_theta,
- "RoPE frequency base. Only used when Block style is Llama-like. Higher values are commonly used when "
- "targeting longer context lengths; 10000 is the standard default.",
- )
- window.use_bias = QCheckBox("Use bias terms")
- window.use_bias.setChecked(True)
- window._tip(
- window.use_bias,
- "Whether linear and normalization layers include bias terms. Some modern architectures (e.g. PaLM, "
- "LLaMA) disable most biases to save a small amount of memory and compute with little quality impact. "
- "Leave checked unless you are deliberately experimenting with this.",
- )
- window.n_embd = window._spin(32, 4096, 128)
- window._tip(window.n_embd, "Embedding size, also called n_embd. Larger values increase model capacity and memory usage.")
- window.architecture_style.currentTextChanged.connect(
- lambda text: window.rope_theta.setEnabled(text == "Llama-like")
- )
- window.rope_theta.setEnabled(window.architecture_style.currentText() == "Llama-like")
- window.n_head = window._spin(1, 64, 4)
- window._tip(window.n_head, "Attention head count. More heads can model varied relationships, but n_embd must divide evenly by n_head.")
- window.attention_type = QComboBox()
- window.attention_type.addItems(["Multi-head", "Grouped-query", "Multi-query"])
- window.attention_type.setMaximumWidth(260)
- window._tip(
- window.attention_type,
- "Attention layout. Grouped-query and multi-query share key/value heads to reduce memory and speed up generation.",
- )
- window.kv_head_count = window._spin(1, 64, 2)
- window._tip(window.kv_head_count, "Key/value heads for grouped-query attention. Must divide n_head. Ignored by multi-head and multi-query.")
- window.attention_backend = QComboBox()
- window.attention_backend.addItems(["SDPA / Flash when available", "Manual"])
- window.attention_backend.setMaximumWidth(260)
- window._tip(
- window.attention_backend,
- "Attention kernel. SDPA lets PyTorch use Flash Attention on supported GPUs and falls back safely otherwise.",
- )
- window.attention_window = window._spin(0, 4096, 0)
- window._tip(window.attention_window, "Sliding attention window. 0 uses full context; higher values restrict attention to recent tokens.")
- window.n_layer = window._spin(1, 64, 4)
- window._tip(window.n_layer, "Transformer layer count. More layers improve capacity and reasoning patterns but slow training.")
- window.train_context_length = window._spin(16, 4096, 128)
- window._tip(window.train_context_length, "Training context length in tokens. Must fit your GPU/CPU memory.")
- window.dropout = window._double_spin(0.0, 0.9, 0.1, 0.01, 3)
- window._tip(window.dropout, "Dropout regularization. Higher values reduce overfitting but can slow learning.")
- left.addRow("Dataset", window._path_row(window.train_data_dir, directory=True))
- left.addRow("Model", window._path_row(window.model_dir, directory=True))
- left.addRow("Preset", window.preset)
- left.addRow("Block style", window.architecture_style)
- left.addRow("RoPE theta", window.rope_theta)
- left.addRow("", window.use_bias)
- left.addRow("n_embd", window.n_embd)
- left.addRow("n_head", window.n_head)
- left.addRow("Attention", window.attention_type)
- left.addRow("KV heads", window.kv_head_count)
- left.addRow("Backend", window.attention_backend)
- left.addRow("Window", window.attention_window)
- left.addRow("n_layer", window.n_layer)
- left.addRow("Context length", window.train_context_length)
- left.addRow("Dropout", window.dropout)
-
- right = QFormLayout()
- window._configure_form(right)
- window.epochs = window._spin(1, 10000, 5)
- window._tip(window.epochs, "Number of full passes over the training tokens. More epochs can improve learning or overfit small data.")
- window.batch_size = window._spin(1, 512, 16)
- window._tip(window.batch_size, "Sequences processed per step. Larger batches are smoother but require more memory.")
- window.learning_rate = window._double_spin(0.000001, 1.0, 0.0003, 0.0001, 6)
- window._tip(window.learning_rate, "Optimizer step size. Too high can destabilize training; too low trains slowly.")
- window.weight_decay = window._double_spin(0.0, 1.0, 0.1, 0.01, 4)
- window._tip(window.weight_decay, "Weight decay regularization. Helps control overfitting by discouraging large weights.")
- window.training_profile = QComboBox()
- window.training_profile.addItems(["Stable LLM", "Low-memory", "Code fine-tune", "Experimental Lion"])
- window.training_profile.setMaximumWidth(260)
- window._tip(window.training_profile, "Applies a practical optimizer, scheduler, precision, and regularization profile.")
- window.apply_training_profile_button = QPushButton("Apply Profile")
- window.apply_training_profile_button.setMaximumWidth(160)
- window.apply_training_profile_button.clicked.connect(window.apply_training_profile)
- window._tip(window.apply_training_profile_button, "Apply the selected training profile to the controls below.")
- window.optimizer_name = QComboBox()
- window.optimizer_name.addItems(["AdamW", "Adam", "Lion", "Adafactor"])
- window.optimizer_name.setMaximumWidth(260)
- window._tip(
- window.optimizer_name,
- "Optimizer algorithm. AdamW is the safest default; Lion can be efficient; Adafactor can reduce optimizer memory when supported.",
- )
- window.scheduler_name = QComboBox()
- window.scheduler_name.addItems(["Warmup linear", "Cosine decay", "Polynomial decay", "One-cycle", "Constant"])
- window.scheduler_name.setMaximumWidth(260)
- window._tip(
- window.scheduler_name,
- "Learning-rate schedule. Cosine and one-cycle are common for stable LLM training; constant is mostly for experiments.",
- )
- window.min_lr_ratio = window._double_spin(0.0, 1.0, 0.1, 0.01, 3)
- window._tip(window.min_lr_ratio, "Lowest learning-rate multiplier after decay. 0.1 means decay down to 10% of the base LR.")
- window.polynomial_power = window._double_spin(0.1, 10.0, 1.0, 0.1, 2)
- window._tip(window.polynomial_power, "Shape of polynomial decay. Higher values decay the learning rate more aggressively near the end.")
- window.gradient_accumulation = window._spin(1, 256, 1)
- window._tip(window.gradient_accumulation, "Accumulate gradients across batches before updating. Simulates larger batches with less memory.")
- window.warmup_steps = window._spin(0, 1_000_000, 100)
- window._tip(window.warmup_steps, "Steps used to ramp up learning rate. Warmup helps avoid unstable early training.")
- window.sample_stride = window._spin(1, 4096, 128)
- window._tip(window.sample_stride, "Token stride between sliding windows when preparing training samples. Larger stride reduces overlapping windows and lowers sample count.",)
- window.eval_interval = window._spin(0, 1_000_000, 100)
- window._tip(window.eval_interval, "Training steps between validation checks. Set 0 to skip interval validation.")
- window.max_eval_batches = window._spin(0, 1_000_000, 50)
- window._tip(window.max_eval_batches, "Maximum validation batches per validation check. Set 0 to evaluate the full validation split.")
- window.save_interval = window._spin(1, 1_000_000, 500)
- window._tip(window.save_interval, "Training steps between checkpoints. Lower values improve crash recovery but use more disk.")
- window.data_loader_workers = window._spin(0, 64, 0)
- window._tip(
- window.data_loader_workers,
- "CPU worker processes used to prepare training batches while the model trains. This is the safe GPU+CPU hybrid mode.",
- )
- window.max_grad_norm = window._double_spin(0.1, 100.0, 1.0, 0.1, 3)
- window._tip(window.max_grad_norm, "Gradient clipping limit. Helps prevent exploding gradients during training.")
- window.activation_checkpointing = QCheckBox("Activation checkpointing")
- window._tip(
- window.activation_checkpointing,
- "Recompute transformer activations during backpropagation to lower VRAM use. Training becomes slower, but this is useful when memory is the constraint.",
- )
- window.seed = window._spin(1, 2_147_483_647, 1337)
- window._tip(window.seed, "Random seed for reproducible initialization and sampling order.")
- window.device = QComboBox()
- window.device.setMaximumWidth(260)
- window._tip(window.device, "Hardware target. CUDA uses NVIDIA GPU when available; CPU is slower but broadly compatible.")
- window.device_info = QLabel()
- window.device_info.setObjectName("Metric")
- window.device_info.setWordWrap(True)
- window.device_info.setMaximumWidth(260)
- window._configure_device_options()
- window.use_amp = QCheckBox("Mixed precision")
- window.use_amp.setChecked(window.use_amp_default)
- window._tip(window.use_amp, "Use mixed precision on CUDA. Usually faster and lighter on GPU memory.")
- window.precision = QComboBox()
- window.precision.addItems(["FP16", "BF16", "FP32"])
- window.precision.setMaximumWidth(260)
- window._tip(
- window.precision,
- "Numeric precision. FP16 is fast on many NVIDIA GPUs; BF16 is more stable on supported GPUs; FP32 is safest but uses more memory.",
- )
- window.resume_training = QCheckBox("Resume latest")
- window.resume_training.setChecked(True)
- window._tip(window.resume_training, "Continue from the latest checkpoint if training was interrupted.")
- window.resume_safety = QCheckBox("Safe resume")
- window.resume_safety.setChecked(True)
- window._tip(
- window.resume_safety,
- "Before resuming, verify that the dataset tokenizer and model architecture match the checkpoint.",
- )
- window.early_stopping = QCheckBox("Early stopping")
- window.early_stopping.setChecked(True)
- window._tip(
- window.early_stopping,
- "Automatically stop training when validation loss stops improving. Uncheck to train for all remaining epochs.",
- )
- window.early_stopping_patience = window._spin(1, 100, 3)
- window._tip(
- window.early_stopping_patience,
- "Consecutive validation checks without improvement before early stopping triggers. Lower values stop "
- "sooner (saves compute, risks stopping on noisy validation loss); higher values are more tolerant of "
- "temporary plateaus. Has no effect if early stopping is unchecked.",
- )
- window.early_stopping.toggled.connect(window.early_stopping_patience.setEnabled)
- window.resume_checkpoint = QLineEdit()
- window._tip(window.resume_checkpoint, "Optional specific checkpoint file to resume from instead of the latest checkpoint.")
- window.resume_check_button = QPushButton("Check Resume")
- window.resume_check_button.setMaximumWidth(180)
- window.resume_check_button.clicked.connect(window.preview_resume_compatibility)
- window._tip(window.resume_check_button, "Inspect checkpoint compatibility before starting training.")
- right.addRow("Epochs", window.epochs)
- right.addRow("Batch", window.batch_size)
- right.addRow("Profile", window.training_profile)
- right.addRow("", window.apply_training_profile_button)
- right.addRow("LR", window.learning_rate)
- right.addRow("Decay", window.weight_decay)
- right.addRow("Optimizer", window.optimizer_name)
- right.addRow("Schedule", window.scheduler_name)
- right.addRow("Min LR", window.min_lr_ratio)
- right.addRow("Poly power", window.polynomial_power)
- right.addRow("Grad accum", window.gradient_accumulation)
- right.addRow("Stride samples", window.sample_stride)
- right.addRow("Warmup", window.warmup_steps)
- right.addRow("Eval every", window.eval_interval)
- right.addRow("Eval batches", window.max_eval_batches)
- right.addRow("Save every", window.save_interval)
- right.addRow("CPU workers", window.data_loader_workers)
- right.addRow("Max grad", window.max_grad_norm)
- right.addRow("VRAM saver", window.activation_checkpointing)
- right.addRow("Seed", window.seed)
- runtime = QFormLayout()
- window._configure_form(runtime)
- window.training_launch_target = QComboBox()
- window.training_launch_target.addItems(["Local machine", "Remote workers", "RunPod cloud"])
- window.training_launch_target.setMaximumWidth(260)
- window._tip(
- window.training_launch_target,
- "Local runs training on this computer. Remote publishes a job for workers. RunPod creates a cloud GPU worker automatically.",
- )
- runtime.addRow("Launch", window.training_launch_target)
- runtime.addRow("Device", window.device)
- runtime.addRow("Hardware", window.device_info)
- runtime.addRow("", window.use_amp)
- runtime.addRow("Precision", window.precision)
- runtime.addRow("", window.resume_training)
- runtime.addRow("", window.resume_safety)
- runtime.addRow("", window.early_stopping)
- runtime.addRow("Patience", window.early_stopping_patience)
- runtime.addRow("Checkpoint", window._path_row(window.resume_checkpoint, directory=False))
- runtime.addRow("", window.resume_check_button)
-
- window.train_button = QPushButton("Start Training")
- window._tip(window.train_button, "Start or resume training using the selected model and optimizer settings.")
- window.train_button.clicked.connect(window.start_training)
- window.train_button.setMaximumWidth(320)
- window.stop_training_button = QPushButton("Stop")
- window.stop_training_button.setEnabled(False)
- window.stop_training_button.setMaximumWidth(120)
- window.stop_training_button.clicked.connect(window.stop_active_task)
- window._tip(window.stop_training_button, "Request a graceful stop and save a resumable checkpoint.")
-
- action_row = QHBoxLayout()
- action_row.addWidget(window.train_button)
- action_row.addWidget(window.stop_training_button)
- action_row.addStretch(1)
-
- architecture_stack = QVBoxLayout()
- architecture_stack.setSpacing(10)
- architecture_card = window._card("MODEL ARCHITECTURE", left)
- architecture_card.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Maximum)
- architecture_stack.addWidget(architecture_card, 0)
- architecture_stack.addLayout(action_row)
- window.training_status_stack = QVBoxLayout()
- window.training_status_stack.setSpacing(10)
- architecture_stack.addLayout(window.training_status_stack)
- architecture_stack.addStretch(1)
-
- architecture_column = QWidget()
- architecture_column.setLayout(architecture_stack)
- architecture_column.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
- optimization_card = window._card("OPTIMIZATION ENGINE", right)
- optimization_card.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Maximum)
- controls_row = QHBoxLayout()
- controls_row.setSpacing(12)
- controls_row.addWidget(architecture_column, 1)
- controls_row.addWidget(optimization_card, 1)
- left_zone.addLayout(controls_row, 1)
- window.training_cards = []
- window.training_controls_grid = None
- window.training_controls_columns = 0
- right_zone.addWidget(window._card("RUNTIME CONTROL", runtime), 0)
-
- window.resume_training_preview = QTextEdit()
- window.resume_training_preview.setReadOnly(True)
- window.resume_training_preview.setMinimumHeight(110)
- window.resume_training_preview.setMaximumHeight(180)
- window.resume_training_preview.setText("No compatibility check has been run.")
- window.resume_preview = window.resume_training_preview
- window._tip(window.resume_training_preview, "Compatibility report for the selected or latest checkpoint.")
- resume_preview_layout = QVBoxLayout()
- resume_preview_layout.addWidget(window.resume_training_preview)
- right_zone.addWidget(window._card("RESUME COMPATIBILITY", resume_preview_layout), 0)
-
- metrics_grid = QGridLayout()
- metrics_grid.setHorizontalSpacing(8)
- metrics_grid.setVerticalSpacing(8)
- window.training_epoch_metric = window._metric_chip("Epoch: -", "Current epoch and total epochs.")
- window.training_step_metric = window._metric_chip("Step: -", "Current optimizer step and total planned steps.")
- window.training_loss_metric = window._metric_chip("Train loss: -", "Latest training loss. Lower is usually better.")
- window.training_val_metric = window._metric_chip("Val loss: -", "Latest validation loss when validation is enabled.")
- window.training_health_metric = window._metric_chip("Health: -", "Training diagnostic based on train and validation loss.")
- window.training_lr_metric = window._metric_chip("LR: -", "Current learning rate from the scheduler.")
- window.training_speed_metric = window._metric_chip("Speed: -", "Current training throughput.")
- window.training_grad_metric = window._metric_chip("Grad: -", "Current gradient norm.")
- window.training_vram_metric = window._metric_chip("VRAM: -", "Current CUDA memory usage when training on GPU.")
- window.training_eta_metric = window._metric_chip("ETA: -", "Estimated time remaining based on recent optimizer steps.")
- window.model_size_metric = window._metric_chip("Model: -", "Estimated model parameter count and checkpoint size.")
- window.vram_estimate_metric = window._metric_chip("VRAM est: -", "Rough training VRAM estimate for selected architecture and batch.")
- window.parameter_breakdown_metric = window._metric_chip(
- "Params: -",
- "Estimated parameters by embedding, attention, MLP, and normalization components.",
- )
- window.memory_breakdown_metric = window._metric_chip(
- "Memory: -",
- "Estimated training memory by weights, optimizer state, activations, and KV cache.",
- )
- window.architecture_advisor_metric = window._metric_chip(
- "Advisor: -",
- "Architecture advisor based on model size, dataset tokens, context length, and memory estimate.",
- )
- window.history_metric = window._metric_chip("Runs: -", "Training run history count in the current model folder.")
- for index, metric in enumerate((
- window.training_eta_metric,
- window.training_epoch_metric,
- window.training_step_metric,
- window.training_loss_metric,
- window.training_val_metric,
- window.training_health_metric,
- window.training_lr_metric,
- window.training_speed_metric,
- window.training_grad_metric,
- window.training_vram_metric,
- )):
- metrics_grid.addWidget(metric, index // 2, index % 2)
- metrics_grid.setColumnStretch(0, 1)
- metrics_grid.setColumnStretch(1, 1)
- metrics_layout = QVBoxLayout()
- metrics_layout.setSpacing(8)
- metrics_layout.addLayout(metrics_grid)
- estimate_grid = QGridLayout()
- estimate_grid.setHorizontalSpacing(8)
- estimate_grid.setVerticalSpacing(8)
- estimate_grid.addWidget(window.model_size_metric, 0, 0)
- estimate_grid.addWidget(window.vram_estimate_metric, 0, 1)
- estimate_grid.addWidget(window.parameter_breakdown_metric, 1, 0)
- estimate_grid.addWidget(window.memory_breakdown_metric, 1, 1)
- estimate_grid.addWidget(window.architecture_advisor_metric, 2, 0, 1, 2)
- estimate_grid.addWidget(window.history_metric, 3, 0)
- window.refresh_estimate_button = QPushButton("Refresh Estimate")
- window.refresh_estimate_button.clicked.connect(window.refresh_model_estimate)
- window.refresh_estimate_button.setMaximumWidth(180)
- window._tip(window.refresh_estimate_button, "Refresh model size, rough VRAM, and training history estimates without starting training.")
- estimate_grid.addWidget(window.refresh_estimate_button, 3, 1)
- estimate_grid.setColumnStretch(0, 1)
- estimate_grid.setColumnStretch(1, 1)
- estimate_layout = QVBoxLayout()
- estimate_layout.setSpacing(8)
- estimate_layout.addLayout(estimate_grid)
- estimate_layout.addStretch(1)
- estimate_card = window._card("MODEL ESTIMATE", estimate_layout)
- metrics_card = window._card("TRAINING METRICS", metrics_layout)
- estimate_card.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Maximum)
- metrics_card.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Maximum)
- status_row = QHBoxLayout()
- status_row.setSpacing(10)
- status_row.addWidget(estimate_card, 1)
- status_row.addWidget(metrics_card, 2)
- window.training_status_stack.addLayout(status_row)
-
- window.training_log = QTextEdit()
- window.training_log.setReadOnly(True)
- window.training_log.document().setMaximumBlockCount(1500)
- window.training_log.setMinimumHeight(306)
- window.training_log.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
- telemetry_layout = QVBoxLayout()
- telemetry_layout.addWidget(window.training_log, 1)
- telemetry_card = window._card("TRAINING TELEMETRY", telemetry_layout)
- telemetry_card.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
- right_zone.addWidget(telemetry_card, 1)
-
- window.training_progress = window._thin_progress()
- outer.addWidget(window.training_progress)
- QTimer.singleShot(0, window._refresh_training_layout)
- return page
\ No newline at end of file
diff --git a/llm_trainer/ui/workers.py b/llm_trainer/ui/workers.py
deleted file mode 100644
index d5aceae..0000000
--- a/llm_trainer/ui/workers.py
+++ /dev/null
@@ -1,216 +0,0 @@
-from __future__ import annotations
-
-import logging
-import multiprocessing as mp
-import time
-from queue import Queue
-from threading import Event
-from typing import Any, Optional
-
-from PySide6.QtCore import QObject, Signal
-
-
-LOGGER = logging.getLogger(__name__)
-
-
-def _process_worker_entry(
- fn: Any,
- args: tuple[Any, ...],
- result_queue: mp.Queue,
- progress_queue: mp.Queue,
- stop_event: mp.Event,
- with_progress: bool,
-) -> None:
- """Run a worker function inside a child process.
-
- Args:
- fn: Callable to run.
- args: Positional arguments.
- result_queue: Queue receiving the final result or failure.
- progress_queue: Queue receiving progress events.
- stop_event: Cross-process cooperative cancellation event.
- with_progress: Whether to pass progress and stop callbacks to ``fn``.
- """
-
- try:
- if with_progress:
- result = fn(
- *args,
- progress=lambda event: progress_queue.put(event),
- should_stop=stop_event.is_set,
- )
- else:
- result = fn(*args)
- result_queue.put(("finished", result))
- except Exception as exc:
- LOGGER.exception("Process worker failed while running %s", getattr(fn, "__name__", fn))
- result_queue.put(("failed", str(exc)))
-
-
-class TaskWorker(QObject):
- """Background worker used for long-running UI tasks."""
-
- finished = Signal(object)
- failed = Signal(str)
-
- def __init__(
- self,
- fn: Any,
- *args: Any,
- progress_queue: Optional[Queue] = None,
- with_progress: bool = False,
- stop_event: Optional[Event] = None,
- ) -> None:
- """Create a worker.
-
- Args:
- fn: Callable to execute in the worker thread.
- *args: Positional arguments passed to ``fn``.
- progress_queue: Optional queue for progress events.
- with_progress: Whether to pass a progress callback to ``fn``.
- stop_event: Optional event used for cooperative cancellation.
- """
-
- super().__init__()
- self.fn = fn
- self.args = args
- self.progress_queue = progress_queue
- self.with_progress = with_progress
- self.stop_event = stop_event
-
- def run(self) -> None:
- """Execute the worker function and emit completion or failure."""
-
- try:
- if self.with_progress:
- self.finished.emit(self.fn(*self.args, progress=self._queue_progress, should_stop=self._should_stop))
- else:
- self.finished.emit(self.fn(*self.args))
- except Exception as exc:
- LOGGER.exception("Background worker failed while running %s", getattr(self.fn, "__name__", self.fn))
- self.failed.emit(str(exc))
-
- def _queue_progress(self, event: Any) -> None:
- if self.progress_queue is not None:
- self.progress_queue.put(event)
-
- def _should_stop(self) -> bool:
- """Return whether the active task has been asked to stop.
-
- Returns:
- True when the cooperative stop event is set.
- """
-
- return bool(self.stop_event and self.stop_event.is_set())
-
-
-class WorkerSignalBridge(QObject):
- """Relay worker results through the GUI thread."""
-
- finished = Signal(object)
- failed = Signal(str)
-
-
-class ProcessTaskWorker(QObject):
- """Background worker that isolates heavy tasks in a child process."""
-
- finished = Signal(object)
- failed = Signal(str)
-
- def __init__(
- self,
- fn: Any,
- *args: Any,
- progress_queue: Optional[Queue] = None,
- with_progress: bool = False,
- stop_event: Optional[Event] = None,
- ) -> None:
- """Create a process-backed worker.
-
- Args:
- fn: Callable to execute in the child process.
- *args: Positional arguments passed to ``fn``.
- progress_queue: UI-thread progress queue.
- with_progress: Whether to pass progress callbacks to ``fn``.
- stop_event: Thread event used to request cancellation.
- """
-
- super().__init__()
- self.fn = fn
- self.args = args
- self.progress_queue = progress_queue
- self.with_progress = with_progress
- self.stop_event = stop_event
-
- def run(self) -> None:
- """Execute the worker function in a separate process."""
-
- context = mp.get_context("spawn")
- child_progress_queue: mp.Queue = context.Queue()
- result_queue: mp.Queue = context.Queue()
- child_stop_event: mp.Event = context.Event()
- process = context.Process(
- target=_process_worker_entry,
- args=(self.fn, self.args, result_queue, child_progress_queue, child_stop_event, self.with_progress),
- # Not daemonic: Python forbids daemonic processes from spawning
- # their own children, but some tasks run here (e.g.
- # build_dataset) spawn their own worker-process pool internally
- # for parallel CPU-bound work. Cleanup does not depend on the
- # daemon flag -- this method explicitly terminates and joins
- # the process below (on stop, on timeout, and in `finally`), so
- # dropping daemon=True does not leave anything unmanaged in the
- # normal shutdown paths.
- daemon=False,
- )
- process.start()
- stop_requested_at: Optional[float] = None
- try:
- while process.is_alive():
- if self.stop_event is not None and self.stop_event.is_set():
- child_stop_event.set()
- if stop_requested_at is None:
- stop_requested_at = time.monotonic()
- elif time.monotonic() - stop_requested_at > 5.0:
- process.terminate()
- break
- self._drain_child_progress(child_progress_queue)
- process.join(0.05)
- self._drain_child_progress(child_progress_queue)
- process.join()
- if stop_requested_at is not None and process.exitcode not in {0, None} and result_queue.empty():
- self.failed.emit("Dataset preparation stopped by user.")
- return
- if result_queue.empty():
- if process.exitcode == 0:
- self.failed.emit("Process finished without returning a result.")
- else:
- self.failed.emit(f"Process exited unexpectedly with code {process.exitcode}.")
- return
- status, payload = result_queue.get()
- if status == "finished":
- self.finished.emit(payload)
- else:
- self.failed.emit(str(payload))
- finally:
- if process.is_alive():
- child_stop_event.set()
- process.terminate()
- process.join(2)
- child_progress_queue.close()
- result_queue.close()
-
- def _drain_child_progress(self, child_progress_queue: mp.Queue) -> None:
- """Move child-process progress events into the UI progress queue.
-
- Args:
- child_progress_queue: Queue owned by the child process.
- """
-
- if self.progress_queue is None:
- return
- while True:
- try:
- event = child_progress_queue.get_nowait()
- except Exception:
- break
- self.progress_queue.put(event)
\ No newline at end of file
diff --git a/llm_trainer/wiki_download.py b/llm_trainer/wiki_download.py
deleted file mode 100644
index 13b6de2..0000000
--- a/llm_trainer/wiki_download.py
+++ /dev/null
@@ -1,19 +0,0 @@
-"""Compatibility facade for Wikipedia download APIs."""
-from PySide6.QtWidgets import QApplication, QMainWindow
-from .wiki_download_backend import WikipediaDownloaderBackend
-from .wiki_download_worker import DownloadWorker
-from .wiki_download_layout import _GuiLayout
-from .wiki_download_style import _GuiStyle
-from .wiki_download_actions import _GuiActions
-from .wiki_download_processing import *
-
-class WikipediaDownloaderGUI(_GuiLayout, _GuiStyle, _GuiActions, QMainWindow):
- """Qt GUI compatibility wrapper for the Wikipedia downloader."""
- pass
-
-def main() -> None:
- """Launch the Wikipedia downloader GUI."""
- app = QApplication.instance() or QApplication([])
- window = WikipediaDownloaderGUI()
- window.show()
- app.exec()
diff --git a/llm_trainer/wiki_download_actions.py b/llm_trainer/wiki_download_actions.py
deleted file mode 100644
index 2174d72..0000000
--- a/llm_trainer/wiki_download_actions.py
+++ /dev/null
@@ -1,312 +0,0 @@
-from __future__ import annotations
-from typing import List
-from PySide6.QtWidgets import *
-from PySide6.QtCore import *
-from PySide6.QtGui import *
-from .wiki_download_backend import WikipediaDownloaderBackend
-from .wiki_download_worker import DownloadWorker
-
-class _GuiActions:
- def setup_connections(self):
- """Setup signal/slot connections"""
- # These are set up in the UI initialization
-
- # ========================================================================
- # Search Methods
- # ========================================================================
-
- def search_pages(self):
- """Search for Wikipedia pages with size/wordcount filtering"""
- query = self.search_input.text().strip()
- if not query:
- QMessageBox.warning(self, "â ī¸ Warning",
- "Please enter a search query")
- return
-
- # Get filter thresholds from UI
- min_size_kb = self.min_size_spin.value() * 1024 # Convert KB to bytes
- min_wordcount = self.min_words_spin.value()
-
- self.search_button.setEnabled(False)
- self.results_text.clear()
- self.search_results_list.clear()
- self.status_bar.showMessage(f"đ Searching for '{query}'...")
-
- try:
- limit = self.limit_spin.value()
- pages = self.downloader.search_pages(query, limit)
-
- # Filter pages by size and wordcount
- filtered_pages = []
- for page in pages:
- size_bytes = page.get('size', 0)
- wordcount = page.get('wordcount', 0)
-
- # Apply filters
- if size_bytes >= min_size_kb and wordcount >= min_wordcount:
- filtered_pages.append(page)
-
- self.current_pages = filtered_pages
-
- if filtered_pages:
- self.results_text.setHtml(f"""
- â
Found {len(filtered_pages)} pages
- (Filtered from {len(pages)} total, min size: {min_size_kb / 1024:.0f}KB, min words: {min_wordcount})
- """)
-
- for page in filtered_pages:
- size_kb = page.get('size', 0) / 1024
- item = QListWidgetItem(
- f"đ {page['title']} | Size: {size_kb:.1f} KB | Words: {page.get('wordcount', 0)}"
- )
- item.setData(Qt.UserRole, page['title'])
- item.setCheckState(Qt.Unchecked)
- self.search_results_list.addItem(item)
-
- self.status_bar.showMessage(
- f"â
Found {len(filtered_pages)} pages meeting criteria")
- else:
- self.results_text.setHtml(f"""
- â No pages meeting criteria
- Try lowering the minimum size or word count thresholds.
- """)
- self.status_bar.showMessage("â No pages meeting criteria")
-
- except Exception as e:
- error_msg = f"Error searching: {str(e)}"
- self.results_text.setHtml(
- f"â {error_msg}")
- self.status_bar.showMessage(f"â {error_msg}")
- QMessageBox.critical(self, "â Error", error_msg)
-
- self.search_button.setEnabled(True)
-
- def select_all_pages(self):
- """Select all pages in search results"""
- for i in range(self.search_results_list.count()):
- item = self.search_results_list.item(i)
- item.setCheckState(Qt.Checked)
- self.status_bar.showMessage("â
All pages selected")
-
- def select_none_pages(self):
- """Deselect all pages in search results"""
- for i in range(self.search_results_list.count()):
- item = self.search_results_list.item(i)
- item.setCheckState(Qt.Unchecked)
- self.status_bar.showMessage("â All pages deselected")
-
- def add_selected_pages(self):
- """Add selected pages to download list"""
- added_count = 0
- existing_titles = set()
-
- # Get existing titles in download list
- for i in range(self.selected_pages_list.count()):
- item = self.selected_pages_list.item(i)
- existing_titles.add(item.text())
-
- for i in range(self.search_results_list.count()):
- item = self.search_results_list.item(i)
- if item.checkState() == Qt.Checked:
- title = item.data(Qt.UserRole)
- if title not in existing_titles:
- self.selected_pages_list.addItem(title)
- existing_titles.add(title)
- added_count += 1
-
- if added_count > 0:
- self.status_bar.showMessage(
- f"â
Added {added_count} pages to download list")
- self.update_download_button_state()
- else:
- QMessageBox.information(self, "âšī¸ Info",
- "No new pages added (may already be in list)")
-
- def clear_page_list(self):
- """Clear the download list"""
- if self.selected_pages_list.count() > 0:
- reply = QMessageBox.question(
- self, "â ī¸ Confirm Clear",
- "Are you sure you want to clear all pages from the download list?",
- QMessageBox.Yes | QMessageBox.No
- )
- if reply == QMessageBox.Yes:
- self.selected_pages_list.clear()
- self.status_bar.showMessage("đī¸ Download list cleared")
- self.update_download_button_state()
-
- # ========================================================================
- # Settings Methods
- # ========================================================================
-
- def update_output_dir(self, text: str):
- """Update output directory"""
- self.output_dir = text
-
- def browse_output_dir(self):
- """Browse for output directory"""
- dir_path = QFileDialog.getExistingDirectory(
- self,
- "đ Select Output Directory",
- self.output_dir,
- QFileDialog.ShowDirsOnly
- )
- if dir_path:
- self.output_dir_edit.setText(dir_path)
- self.output_dir = dir_path
-
- # ========================================================================
- # Download Methods
- # ========================================================================
-
- def get_pages_to_download(self) -> List[str]:
- """Get list of pages to download"""
- pages = []
- for i in range(self.selected_pages_list.count()):
- pages.append(self.selected_pages_list.item(i).text())
- return pages
-
- def update_download_button_state(self):
- """Update download button state based on list content"""
- count = self.selected_pages_list.count()
- has_pages = count > 0
- self.download_button.setEnabled(has_pages and not self.worker)
- self.page_count_label.setText(f"đ Pages in queue: {count}")
-
- def start_download(self):
- """Start the download process"""
- pages = self.get_pages_to_download()
- if not pages:
- QMessageBox.warning(self, "â ī¸ Warning", "No pages to download")
- return
-
- # Check output directory
- output_dir = self.output_dir_edit.text()
- if not output_dir:
- QMessageBox.warning(self, "â ī¸ Warning",
- "Please specify an output directory")
- return
-
- # Confirm
- reply = QMessageBox.question(
- self,
- "đ Confirm Download",
- f"Download {len(pages)} pages to:\n{output_dir}\n\nContinue?",
- QMessageBox.Yes | QMessageBox.No
- )
- if reply != QMessageBox.Yes:
- return
-
- # Disable UI
- self.download_button.setEnabled(False)
- self.cancel_button.setEnabled(True)
- self.search_button.setEnabled(False)
- self.progress_bar.setValue(0)
-
- # Create and start worker
- self.worker = DownloadWorker(
- pages,
- output_dir,
- self.save_metadata_check.isChecked()
- )
-
- # Connect signals
- self.worker.progress_updated.connect(self.update_progress)
- self.worker.page_downloaded.connect(self.on_page_downloaded)
- self.worker.status_updated.connect(self.update_status)
- self.worker.download_complete.connect(self.on_download_complete)
- self.worker.error_occurred.connect(self.on_error)
-
- self.worker.start()
- self.status_bar.showMessage("âŗ Downloading...")
-
- def cancel_download(self):
- """Cancel the download"""
- if self.worker and self.worker.isRunning():
- reply = QMessageBox.question(
- self,
- "âšī¸ Cancel Download",
- "Are you sure you want to cancel the download?",
- QMessageBox.Yes | QMessageBox.No
- )
- if reply == QMessageBox.Yes:
- self.worker.stop()
- self.status_bar.showMessage("âšī¸ Cancelling download...")
-
- def update_progress(self, current: int, total: int):
- """Update progress bar"""
- progress = int((current / total) * 100)
- self.progress_bar.setValue(progress)
- self.progress_label.setText(f"đ {current}/{total}")
-
- def on_page_downloaded(self, title: str, success: bool):
- """Handle page download status"""
- status = "â
" if success else "â"
- if success:
- self.status_bar.showMessage(f"{status} Downloaded: {title}")
- else:
- self.status_bar.showMessage(f"{status} Failed: {title}")
-
- def update_status(self, message: str):
- """Update status message"""
- self.status_bar.showMessage(message)
-
- def on_error(self, error_message: str):
- """Handle error"""
- self.status_bar.showMessage(f"â Error: {error_message}")
- # Log error but continue
- print(f"Error: {error_message}")
-
- def on_download_complete(self, summary: dict):
- """Handle download completion"""
-
- # Wait for worker to completely terminate
- if self.worker is not None:
- self.worker.wait()
- self.worker.deleteLater()
-
- # Enable UI
- self.download_button.setEnabled(True)
- self.cancel_button.setEnabled(False)
- self.search_button.setEnabled(True)
-
- self.progress_bar.setValue(100)
-
- msg = (
- f"đ Download Complete!\n\n"
- f"đ Total pages: {summary['total']}\n"
- f"â
Downloaded: {summary['downloaded']}\n"
- f"â Failed: {summary['failed']}\n"
- f"âī¸ Skipped: {summary['skipped']}\n\n"
- f"đ Output directory:\n{summary['output_dir']}"
- )
-
- QMessageBox.information(
- self,
- "Download Complete",
- msg
- )
-
- self.status_bar.showMessage("Download complete")
- self.progress_label.setText("Done")
-
- self.update_download_button_state()
-
- try:
- cleanup(
- INPUT_DIR=self.output_dir_edit.text(),
- OUTPUT_DIR=os.path.join(
- self.output_dir_edit.text(),
- "cleaned_files"
- )
- )
- except Exception:
- import traceback
- traceback.print_exc()
-
- self.worker = None
-
-# ============================================================================
-# Main Entry Point
-# ============================================================================
-
diff --git a/llm_trainer/wiki_download_backend.py b/llm_trainer/wiki_download_backend.py
deleted file mode 100644
index a27d458..0000000
--- a/llm_trainer/wiki_download_backend.py
+++ /dev/null
@@ -1,117 +0,0 @@
-from __future__ import annotations
-import re
-import time
-from typing import Dict, List, Optional
-import requests
-
-
-class WikipediaDownloaderBackend:
- """Backend class for downloading Wikipedia pages"""
-
- def __init__(self):
- self.api_url = "https://en.wikipedia.org/w/api.php"
- self.session = requests.Session()
- self.min_request_interval = 2.0
- self.last_request_time = 0
- self.is_running = False
-
- def _rate_limit(self):
- """Rate limiting for Wikipedia API"""
- current_time = time.time()
- time_since_last = current_time - self.last_request_time
- if time_since_last < self.min_request_interval:
- time.sleep(self.min_request_interval - time_since_last)
- self.last_request_time = time.time()
-
- def _make_request(self, params: Dict) -> Dict:
- """Make API request with rate limiting"""
- self._rate_limit()
- print(params)
- try:
- response = self.session.get(
- self.api_url,
- params=params,
- headers={'User-Agent': 'DrunkenBot-Wikipedia-GUI/1.0'}
- )
- response.raise_for_status()
- return response.json()
- except Exception as e:
- return {'error': str(e)}
-
- def search_pages(self, query: str, limit: int = 50) -> List[Dict]:
- """Search for Wikipedia pages"""
- params = {
- 'action': 'query',
- 'list': 'search',
- 'srsearch': query,
- 'format': 'json',
- 'srlimit': limit
- }
-
- data = self._make_request(params)
- if 'error' in data:
- return []
-
- results = data.get('query', {}).get('search', [])
- pages = []
- for result in results:
- pages.append({
- 'title': result['title'],
- 'pageid': result['pageid'],
- 'snippet': result.get('snippet', ''),
- 'size': result.get('size', 0),
- 'wordcount': result.get('wordcount', 0)
- })
- return pages
-
- def get_page_content(self, title: str) -> Optional[Dict]:
- """Get full page content"""
- params = {
- 'action': 'parse',
- 'page': title,
- 'format': 'json',
- 'prop': 'text|revid|categories|links',
- 'formatversion': 2
- }
-
- data = self._make_request(params)
- if 'error' in data:
- return None
-
- parse_data = data.get('parse', {})
- if not parse_data:
- return None
-
- html_content = parse_data.get('text', '')
- plain_text = self._clean_html(html_content)
-
- return {
- 'title': title,
- 'text': plain_text,
- 'revid': parse_data.get('revid', 0),
- 'categories': parse_data.get('categories', []),
- 'timestamp': datetime.utcnow().isoformat()
- }
-
- def _clean_html(self, html_content: str) -> str:
- """Extract plain text from HTML"""
- import html
- text = re.sub(r'<[^>]+>', ' ', html_content)
- text = html.unescape(text)
- text = re.sub(r'\s+', ' ', text)
- text = text.strip()
- text = re.sub(r'\[\d+\]', '', text)
- return text
-
- def sanitize_filename(self, title: str) -> str:
- """Create safe filename"""
- safe = re.sub(r'[<>:"/\\|?*]', '_', title)
- if len(safe) > 200:
- safe = safe[:200]
- return safe
-
-
-# ============================================================================
-# Worker Thread for Downloading
-# ============================================================================
-
diff --git a/llm_trainer/wiki_download_layout.py b/llm_trainer/wiki_download_layout.py
deleted file mode 100644
index f1def46..0000000
--- a/llm_trainer/wiki_download_layout.py
+++ /dev/null
@@ -1,355 +0,0 @@
-from __future__ import annotations
-from typing import List
-from PySide6.QtWidgets import *
-from PySide6.QtCore import *
-from PySide6.QtGui import *
-from .wiki_download_backend import WikipediaDownloaderBackend
-from .wiki_download_worker import DownloadWorker
-
-class _GuiLayout:
- def init_ui(self):
- """Initialize the user interface"""
- self.setWindowTitle("Wikipedia Dataset Downloader - DrunkenBot")
- self.setGeometry(100, 100, 1100, 800)
-
- # Apply modern color scheme
- self.apply_styles()
-
- # Central widget and main layout
- central_widget = QWidget()
- self.setCentralWidget(central_widget)
- main_layout = QVBoxLayout(central_widget)
- main_layout.setSpacing(15)
- main_layout.setContentsMargins(15, 15, 15, 15)
-
- # ====================================================================
- # Search Section
- # ====================================================================
- search_group = QGroupBox("đ Search Wikipedia")
- search_layout = QVBoxLayout()
-
- # Search input row
- input_layout = QHBoxLayout()
- self.search_input = QLineEdit()
- self.search_input.setPlaceholderText(
- "Enter topic to search (e.g., Artificial Intelligence)")
- self.search_input.returnPressed.connect(self.search_pages)
- self.search_input.setMinimumHeight(35)
-
- self.search_button = QPushButton("đ Search")
- self.search_button.clicked.connect(self.search_pages)
- self.search_button.setMinimumHeight(35)
-
- self.limit_spin = QSpinBox()
- self.limit_spin.setRange(5, 1000)
- self.limit_spin.setValue(20)
- self.limit_spin.setPrefix("Max results: ")
- self.limit_spin.setMinimumHeight(35)
-
- input_layout.addWidget(self.search_input, 3)
- input_layout.addWidget(self.limit_spin, 1)
- input_layout.addWidget(self.search_button, 1)
-
- search_layout.addLayout(input_layout)
-
- # Results display
- self.results_text = QTextEdit()
- self.results_text.setReadOnly(True)
- self.results_text.setMaximumHeight(80)
- self.results_text.setPlaceholderText(
- "Search results will appear here...")
- self.results_text.setStyleSheet("""
- QTextEdit {
- background-color: #f8f9fa;
- color: #212529;
- border: 1px solid #dee2e6;
- border-radius: 5px;
- padding: 8px;
- font-size: 12px;
- }
- """)
-
- search_layout.addWidget(self.results_text)
- search_group.setLayout(search_layout)
- main_layout.addWidget(search_group)
-
- # ====================================================================
- # Page Selection Section
- # ====================================================================
- selection_group = QGroupBox("đ Pages to Download")
- selection_layout = QVBoxLayout()
-
- # Control buttons for selection
- selection_controls = QHBoxLayout()
- self.select_all_button = QPushButton("â
Select All")
- self.select_all_button.clicked.connect(self.select_all_pages)
- self.select_none_button = QPushButton("â Select None")
- self.select_none_button.clicked.connect(self.select_none_pages)
- self.add_selected_button = QPushButton("â Add Selected to Download")
- self.add_selected_button.clicked.connect(self.add_selected_pages)
- self.clear_list_button = QPushButton("đī¸ Clear List")
- self.clear_list_button.clicked.connect(self.clear_page_list)
- self.clear_list_button.setObjectName("danger")
-
- for btn in [self.select_all_button, self.select_none_button,
- self.add_selected_button, self.clear_list_button]:
- btn.setMinimumHeight(30)
-
- selection_controls.addWidget(self.select_all_button)
- selection_controls.addWidget(self.select_none_button)
- selection_controls.addWidget(self.add_selected_button)
- selection_controls.addWidget(self.clear_list_button)
- selection_controls.addStretch()
-
- selection_layout.addLayout(selection_controls)
-
- # Split view for search results and selected pages
- splitter = QSplitter(Qt.Horizontal)
-
- # Search results list with checkboxes
- self.search_results_list = QListWidget()
- self.search_results_list.setSelectionMode(
- QListWidget.ExtendedSelection)
- self.search_results_list.setMinimumHeight(200)
- self.search_results_list.setStyleSheet("""
- QListWidget {
- background-color: white;
- color: #212529;
- border: 1px solid #dee2e6;
- border-radius: 5px;
- padding: 5px;
- }
- QListWidget::item {
- padding: 5px;
- border-bottom: 1px solid #f0f0f0;
- color: #212529;
- }
- QListWidget::item:selected {
- background-color: #e3f2fd;
- color: #212529;
- }
- QListWidget::item:hover {
- background-color: #f8f9fa;
- }
- """)
-
- # Selected pages list
- self.selected_pages_list = QListWidget()
- self.selected_pages_list.setMinimumHeight(200)
- self.selected_pages_list.setStyleSheet("""
- QListWidget {
- background-color: #f8f9fa;
- color: #212529;
- border: 2px solid #4CAF50;
- border-radius: 5px;
- padding: 5px;
- }
- QListWidget::item {
- padding: 5px;
- border-bottom: 1px solid #e0e0e0;
- color: #212529;
- }
- QListWidget::item:selected {
- background-color: #c8e6c9;
- color: #212529;
- }
- QListWidget::item:hover {
- background-color: #e8f5e9;
- }
- """)
-
- # Labels for lists
- left_widget = QWidget()
- left_layout = QVBoxLayout(left_widget)
- left_layout.setContentsMargins(0, 0, 0, 0)
- left_label = QLabel("đ Search Results")
- left_label.setStyleSheet(
- "font-weight: bold; color: #212529; padding: 5px;")
- left_layout.addWidget(left_label)
- left_layout.addWidget(self.search_results_list)
-
- right_widget = QWidget()
- right_layout = QVBoxLayout(right_widget)
- right_layout.setContentsMargins(0, 0, 0, 0)
- right_label = QLabel("đĨ Download Queue")
- right_label.setStyleSheet(
- "font-weight: bold; color: #212529; padding: 5px;")
- right_layout.addWidget(right_label)
- right_layout.addWidget(self.selected_pages_list)
-
- splitter.addWidget(left_widget)
- splitter.addWidget(right_widget)
- splitter.setSizes([500, 500])
-
- selection_layout.addWidget(splitter)
- selection_group.setLayout(selection_layout)
- main_layout.addWidget(selection_group)
-
- # ====================================================================
- # Settings Section
- # ====================================================================
- settings_group = QGroupBox("âī¸ Download Settings")
- settings_layout = QGridLayout()
- settings_layout.setSpacing(10)
-
- # Output directory
- settings_layout.addWidget(QLabel("đ Output Directory:"), 0, 0)
- self.output_dir_edit = QLineEdit(self.output_dir)
- self.output_dir_edit.textChanged.connect(self.update_output_dir)
- self.output_dir_edit.setStyleSheet("""
- QLineEdit {
- padding: 8px;
- border: 1px solid #dee2e6;
- border-radius: 4px;
- background-color: white;
- color: #212529;
- }
- """)
- settings_layout.addWidget(self.output_dir_edit, 0, 1)
-
- self.browse_button = QPushButton("đ Browse...")
- self.browse_button.clicked.connect(self.browse_output_dir)
- self.browse_button.setMinimumHeight(30)
- settings_layout.addWidget(self.browse_button, 0, 2)
-
- # Options
- self.save_metadata_check = QCheckBox("đž Save metadata (JSON)")
- self.save_metadata_check.setChecked(True)
- self.save_metadata_check.setStyleSheet("color: #212529;")
- settings_layout.addWidget(self.save_metadata_check, 1, 0, 1, 2)
-
- self.overwrite_check = QCheckBox("đ Overwrite existing files")
- self.overwrite_check.setChecked(False)
- self.overwrite_check.setStyleSheet("color: #212529;")
- settings_layout.addWidget(self.overwrite_check, 1, 2)
-
- settings_group.setLayout(settings_layout)
- main_layout.addWidget(settings_group)
-
- # Filter controls
- filters_layout = QHBoxLayout()
- filters_layout.addWidget(QLabel("Min Size (KB):"))
- self.min_size_spin = QDoubleSpinBox()
- self.min_size_spin.setRange(0, 10000)
- self.min_size_spin.setValue(100)
- self.min_size_spin.setSuffix(" KB")
- filters_layout.addWidget(self.min_size_spin)
-
- filters_layout.addWidget(QLabel("Min Words:"))
- self.min_words_spin = QSpinBox()
- self.min_words_spin.setRange(0, 100000)
- self.min_words_spin.setValue(15000)
- filters_layout.addWidget(self.min_words_spin)
-
- # Add to your settings layout
- settings_layout.addLayout(filters_layout, 2, 0, 1, 3)
-
- # ====================================================================
- # Download Controls
- # ====================================================================
- download_group = QGroupBox("âŦī¸ Download")
- download_layout = QVBoxLayout()
-
- # Progress bar
- self.progress_bar = QProgressBar()
- self.progress_bar.setMinimumHeight(25)
- self.progress_bar.setStyleSheet("""
- QProgressBar {
- border: 1px solid #dee2e6;
- border-radius: 5px;
- text-align: center;
- background-color: white;
- color: #212529;
- }
- QProgressBar::chunk {
- background-color: #4CAF50;
- border-radius: 5px;
- }
- """)
- download_layout.addWidget(self.progress_bar)
-
- # Control buttons
- control_layout = QHBoxLayout()
- self.download_button = QPushButton("đ Start Download")
- self.download_button.clicked.connect(self.start_download)
- self.download_button.setMinimumHeight(40)
- self.download_button.setStyleSheet("""
- QPushButton {
- background-color: #2196F3;
- color: white;
- font-size: 14px;
- font-weight: bold;
- padding: 10px 20px;
- border: none;
- border-radius: 5px;
- }
- QPushButton:hover {
- background-color: #1976D2;
- }
- QPushButton:disabled {
- background-color: #b0bec5;
- color: #ffffff;
- }
- """)
-
- self.cancel_button = QPushButton("âšī¸ Cancel")
- self.cancel_button.clicked.connect(self.cancel_download)
- self.cancel_button.setObjectName("danger")
- self.cancel_button.setMinimumHeight(40)
- self.cancel_button.setStyleSheet("""
- QPushButton {
- background-color: #f44336;
- color: white;
- font-size: 14px;
- font-weight: bold;
- padding: 10px 20px;
- border: none;
- border-radius: 5px;
- }
- QPushButton:hover {
- background-color: #d32f2f;
- }
- QPushButton:disabled {
- background-color: #ef9a9a;
- color: #ffffff;
- }
- """)
- self.cancel_button.setEnabled(False)
-
- control_layout.addWidget(self.download_button)
- control_layout.addWidget(self.cancel_button)
- control_layout.addStretch()
-
- # Page count label
- self.page_count_label = QLabel("Pages in queue: 0")
- self.page_count_label.setStyleSheet(
- "color: #212529; font-weight: bold;")
- control_layout.addWidget(self.page_count_label)
-
- download_layout.addLayout(control_layout)
- download_group.setLayout(download_layout)
- main_layout.addWidget(download_group)
-
- # ====================================================================
- # Status Bar
- # ====================================================================
- self.status_bar = QStatusBar()
- self.status_bar.setStyleSheet("""
- QStatusBar {
- background-color: #f8f9fa;
- color: #212529;
- border-top: 1px solid #dee2e6;
- padding: 5px;
- }
- """)
- self.setStatusBar(self.status_bar)
- self.status_bar.showMessage("â
Ready")
-
- # Add progress label to status bar
- self.progress_label = QLabel("")
- self.progress_label.setStyleSheet("color: #212529; font-weight: bold;")
- self.status_bar.addPermanentWidget(self.progress_label)
-
- # Update initial state
- self.update_download_button_state()
-
diff --git a/llm_trainer/wiki_download_processing.py b/llm_trainer/wiki_download_processing.py
deleted file mode 100644
index 44b5092..0000000
--- a/llm_trainer/wiki_download_processing.py
+++ /dev/null
@@ -1,243 +0,0 @@
-from __future__ import annotations
-import json
-import re
-from pathlib import Path
-from typing import List
-
-
-def remove_sections(text):
-
- for section in REMOVE_SECTIONS:
-
- pattern = (
- rf"\n{section}\n.*"
- )
-
- text = re.sub(
- pattern,
- "",
- text,
- flags=re.IGNORECASE | re.DOTALL,
- )
-
- return text
-
-
-def clean_text(text):
- import re
-
- # ---------------------------------------------------------
- # Remove CSS
- # ---------------------------------------------------------
- text = re.sub(
- r"\.mw-parser-output.*?(?=The |\# |\n[A-Z])",
- "",
- text,
- flags=re.DOTALL,
- )
-
- text = re.sub(
- r"@media.*?(?=The |\# |\n[A-Z])",
- "",
- text,
- flags=re.DOTALL,
- )
-
- # ---------------------------------------------------------
- # Remove references like [1], [23], [a]
- # ---------------------------------------------------------
- text = re.sub(r"\[[^\]]+\]", "", text)
-
- # ---------------------------------------------------------
- # Remove edit markers
- # ---------------------------------------------------------
- text = text.replace("[edit]", "")
-
- # ---------------------------------------------------------
- # Collapse whitespace first
- # ---------------------------------------------------------
- text = re.sub(r"\s+", " ", text).strip()
-
- # ---------------------------------------------------------
- # Remove everything before the first real paragraph.
- # Most Wikipedia pages begin with
- #
- # "The ..."
- # "A ..."
- # "An ..."
- #
- # This removes infoboxes/navigation.
- # ---------------------------------------------------------
- m = re.search(r"\b(The|A|An)\b.+", text)
-
- if m:
- text = text[m.start():]
-
- # ---------------------------------------------------------
- # Sentence splitting
- # ---------------------------------------------------------
- text = re.sub(
- r"([.!?])\s+",
- r"\1\n",
- text
- )
-
- # ---------------------------------------------------------
- # Rebuild paragraphs
- # ---------------------------------------------------------
- paragraph_starters = (
- "The ",
- "In ",
- "On ",
- "At ",
- "After ",
- "Before ",
- "During ",
- "By ",
- "Following ",
- "Meanwhile ",
- "However ",
- "Although ",
- "Later ",
- "Since ",
- "From ",
- "As ",
- "When ",
- "While ",
- )
-
- paragraphs = []
- current = ""
-
- for line in text.splitlines():
-
- line = line.strip()
-
- if not line:
- continue
-
- if current == "":
- current = line
- continue
-
- if line.startswith(paragraph_starters):
- paragraphs.append(current.strip())
- current = line
- else:
- current += " " + line
-
- if current:
- paragraphs.append(current.strip())
-
- # ---------------------------------------------------------
- # Remove obvious junk paragraphs
- # ---------------------------------------------------------
- cleaned = []
-
- junk_words = (
- "Belligerents",
- "Campaign",
- "Atlantic Theater",
- "West Indies",
- "Result",
- "Date",
- "Location",
- "Combatants",
- "Casualties",
- "Commander",
- "References",
- "External links",
- "Bibliography",
- "Further reading",
- "See also",
- )
-
- for p in paragraphs:
-
- if len(p) < 40:
- continue
-
- if any(word in p for word in junk_words):
- continue
-
- cleaned.append(p)
-
- return "\n\n".join(cleaned)
-
-
-def chunk_text(text, words_per_chunk):
-
- words = text.split()
-
- chunks = []
-
- for i in range(0, len(words), words_per_chunk):
-
- chunks.append(
- " ".join(words[i:i + words_per_chunk])
- )
-
- return chunks
-
-
-def process_file(file_path, output_dir):
-
- out = Path(output_dir) / file_path.name
-
- # Skip if already cleaned
- if out.exists():
- print(f"Skipping (already cleaned): {file_path.name}")
- return
-
- text = file_path.read_text(
- encoding="utf8",
- errors="ignore",
- )
-
- cleaned = clean_text(text)
-
- out.write_text(
- cleaned,
- encoding="utf8",
- )
-
- print(f"Cleaned: {file_path.name}")
-
-
-def cleanup(INPUT_DIR, OUTPUT_DIR):
-
- input_dir = Path(INPUT_DIR)
- output_dir = Path(OUTPUT_DIR)
-
- output_dir.mkdir(
- exist_ok=True,
- parents=True,
- )
-
- files = list(input_dir.glob("*.txt"))
-
- print(f"Found {len(files)} files")
-
- cleaned_count = 0
- skipped_count = 0
-
- for i, file in enumerate(files, 1):
-
- print(f"[{i}/{len(files)}] {file.name}")
-
- out = output_dir / file.name
-
- if out.exists():
- print(" -> Already cleaned, skipping.")
- skipped_count += 1
- continue
-
- process_file(file, output_dir)
- cleaned_count += 1
-
- print()
- print(f"Cleanup Done. Cleaned: {cleaned_count}, Skipped: {skipped_count}")
-
-
-if __name__ == "__main__":
- main()
diff --git a/llm_trainer/wiki_download_style.py b/llm_trainer/wiki_download_style.py
deleted file mode 100644
index ebc886e..0000000
--- a/llm_trainer/wiki_download_style.py
+++ /dev/null
@@ -1,206 +0,0 @@
-from __future__ import annotations
-from typing import List
-from PySide6.QtWidgets import *
-from PySide6.QtCore import *
-from PySide6.QtGui import *
-from .wiki_download_backend import WikipediaDownloaderBackend
-from .wiki_download_worker import DownloadWorker
-
-class _GuiStyle:
- def apply_styles(self):
- """Apply modern stylesheet to the application with proper colors"""
- self.setStyleSheet("""
- QMainWindow {
- background-color: #f0f2f5;
- }
- QGroupBox {
- font-weight: bold;
- border: 2px solid #d0d7de;
- border-radius: 8px;
- margin-top: 10px;
- padding-top: 15px;
- padding-bottom: 15px;
- background-color: #ffffff;
- color: #1a1a1a;
- }
- QGroupBox::title {
- subcontrol-origin: margin;
- left: 10px;
- padding: 0 10px 0 10px;
- color: #1a1a1a;
- background-color: #ffffff;
- }
- QLabel {
- color: #1a1a1a;
- }
- QCheckBox {
- color: #1a1a1a;
- background-color: transparent;
- }
- QSpinBox {
- padding: 5px;
- border: 1px solid #d0d7de;
- border-radius: 4px;
- background-color: #ffffff;
- color: #1a1a1a;
- }
- QSpinBox::up-button, QSpinBox::down-button {
- background-color: #f0f2f5;
- }
- QLineEdit {
- padding: 5px;
- border: 1px solid #d0d7de;
- border-radius: 4px;
- background-color: #ffffff;
- color: #1a1a1a;
- }
- QTextEdit {
- background-color: #f8f9fa;
- color: #1a1a1a;
- border: 1px solid #d0d7de;
- border-radius: 5px;
- }
- QPushButton {
- background-color: #2ea44f;
- color: #ffffff;
- border: none;
- padding: 8px 16px;
- border-radius: 4px;
- font-weight: bold;
- }
- QPushButton:hover {
- background-color: #22863a;
- }
- QPushButton:disabled {
- background-color: #d0d7de;
- color: #8b949e;
- }
- QPushButton#danger {
- background-color: #da3633;
- }
- QPushButton#danger:hover {
- background-color: #b62324;
- }
- QSplitter::handle {
- background-color: #d0d7de;
- width: 2px;
- }
- QSplitter::handle:hover {
- background-color: #2ea44f;
- }
- QListWidget {
- background-color: #ffffff;
- color: #1a1a1a;
- border: 1px solid #d0d7de;
- border-radius: 5px;
- padding: 5px;
- }
- QListWidget::item {
- color: #1a1a1a;
- padding: 8px;
- border-bottom: 1px solid #f0f2f5;
- }
- QListWidget::item:selected {
- background-color: #ddf4ff;
- color: #1a1a1a;
- border: none;
- }
- QListWidget::item:hover {
- background-color: #f6f8fa;
- }
- QProgressBar {
- border: 1px solid #d0d7de;
- border-radius: 5px;
- text-align: center;
- background-color: #ffffff;
- color: #1a1a1a;
- }
- QProgressBar::chunk {
- background-color: #2ea44f;
- border-radius: 5px;
- }
- QStatusBar {
- background-color: #f8f9fa;
- color: #1a1a1a;
- border-top: 1px solid #d0d7de;
- padding: 5px;
- }
- QScrollBar:vertical {
- background-color: #f6f8fa;
- width: 12px;
- border-radius: 6px;
- }
- QScrollBar::handle:vertical {
- background-color: #d0d7de;
- border-radius: 6px;
- min-height: 20px;
- }
- QScrollBar::handle:vertical:hover {
- background-color: #8b949e;
- }
- QScrollBar:horizontal {
- background-color: #f6f8fa;
- height: 12px;
- border-radius: 6px;
- }
- QScrollBar::handle:horizontal {
- background-color: #d0d7de;
- border-radius: 6px;
- min-width: 20px;
- }
- QScrollBar::handle:horizontal:hover {
- background-color: #8b949e;
- }
- QMenuBar {
- background-color: #ffffff;
- color: #1a1a1a;
- }
- QMenuBar::item:selected {
- background-color: #f0f2f5;
- }
- QMenu {
- background-color: #ffffff;
- color: #1a1a1a;
- }
- QMenu::item:selected {
- background-color: #f0f2f5;
- }
- QMessageBox {
- background-color: #ffffff;
- color: #1a1a1a;
- }
- QMessageBox QLabel {
- color: #1a1a1a;
- }
- QMessageBox QPushButton {
- background-color: #2ea44f;
- color: #ffffff;
- min-width: 80px;
- padding: 8px;
- }
- QMessageBox QPushButton:hover {
- background-color: #22863a;
- }
- QDialog {
- background-color: #ffffff;
- color: #1a1a1a;
- }
- QDialog QLabel {
- color: #1a1a1a;
- }
- QCheckBox::indicator {
- width: 18px;
- height: 18px;
- }
- QCheckBox::indicator:unchecked {
- background-color: #ffffff;
- border: 2px solid #d0d7de;
- border-radius: 4px;
- }
- QCheckBox::indicator:checked {
- background-color: #2ea44f;
- border: 2px solid #2ea44f;
- border-radius: 4px;
- }
- """)
-
diff --git a/llm_trainer/wiki_download_worker.py b/llm_trainer/wiki_download_worker.py
deleted file mode 100644
index 192ea5f..0000000
--- a/llm_trainer/wiki_download_worker.py
+++ /dev/null
@@ -1,136 +0,0 @@
-from __future__ import annotations
-import json
-import os
-import threading
-from pathlib import Path
-from typing import List
-from PySide6.QtCore import QThread, Signal
-from .wiki_download_backend import WikipediaDownloaderBackend
-
-class DownloadWorker(QThread):
- """Worker thread for downloading pages without blocking UI"""
-
- # Signals
- progress_updated = Signal(int, int) # current, total
- page_downloaded = Signal(str, bool) # title, success
- status_updated = Signal(str) # status message
- download_complete = Signal(dict) # summary stats
- error_occurred = Signal(str) # error message
-
- def __init__(self, pages: List[str], output_dir: str,
- save_metadata: bool = False):
- super().__init__()
- self.pages = pages
- self.output_dir = output_dir
- self.save_metadata = save_metadata
- self.is_running = True
- self.downloader = WikipediaDownloaderBackend()
-
- def run(self):
- """Main download process"""
- total_pages = len(self.pages)
- downloaded = 0
- failed = 0
- skipped = 0
- successful_titles = []
- failed_titles = []
-
- output_path = Path(self.output_dir)
- output_path.mkdir(parents=True, exist_ok=True)
-
- self.status_updated.emit(
- f"Starting download of {total_pages} pages...")
-
- for idx, title in enumerate(self.pages, 1):
- if not self.is_running:
- self.status_updated.emit("Download cancelled")
- break
-
- self.progress_updated.emit(idx, total_pages)
- self.status_updated.emit(
- f"Downloading: {title} ({idx}/{total_pages})")
-
- # Check if already exists
- safe_title = self.downloader.sanitize_filename(title)
- file_path = output_path / f"{safe_title}.txt"
-
- if file_path.exists():
- skipped += 1
- self.page_downloaded.emit(title, False)
- self.status_updated.emit(f"Skipped {title} (already exists)")
- continue
-
- # Download page
- content = self.downloader.get_page_content(title)
-
- if content and content.get('text'):
- try:
- # Save text
- with open(file_path, 'w', encoding='utf-8') as f:
- f.write(content['text'])
-
- # Save metadata if requested
- if self.save_metadata:
- meta_path = output_path / f"{safe_title}.meta.json"
- with open(meta_path, 'w', encoding='utf-8') as f:
- json.dump(content, f, indent=2)
-
- downloaded += 1
- successful_titles.append(title)
- self.page_downloaded.emit(title, True)
-
- except Exception as e:
- failed += 1
- failed_titles.append(title)
- self.error_occurred.emit(f"Error saving {title}: {str(e)}")
- else:
- failed += 1
- failed_titles.append(title)
- self.page_downloaded.emit(title, False)
-
- # Small delay between requests
- time.sleep(0.5)
-
- # Save index file
- self._save_index(successful_titles, failed_titles, output_path)
-
- # Emit completion signal
- summary = {
- 'total': total_pages,
- 'downloaded': downloaded,
- 'failed': failed,
- 'skipped': skipped,
- 'successful_titles': successful_titles,
- 'failed_titles': failed_titles,
- 'output_dir': str(output_path)
- }
-
- self.download_complete.emit(summary)
- self.status_updated.emit(
- f"Download complete! Downloaded: {downloaded}, Failed: {failed}, Skipped: {skipped}")
-
- def _save_index(self, successful_titles: List[str],
- failed_titles: List[str], output_path: Path):
- """Save index file"""
- index = {
- 'download_date': datetime.utcnow().isoformat(),
- 'total_pages': len(successful_titles) + len(failed_titles),
- 'successful': len(successful_titles),
- 'failed': len(failed_titles),
- 'successful_titles': successful_titles,
- 'failed_titles': failed_titles
- }
-
- index_path = output_path / 'download_index.json'
- with open(index_path, 'w', encoding='utf-8') as f:
- json.dump(index, f, indent=2)
-
- def stop(self):
- """Stop the download process"""
- self.is_running = False
-
-
-# ============================================================================
-# Main GUI Application
-# ============================================================================
-
diff --git a/llm_trainer/worker/__init__.py b/llm_trainer/worker/__init__.py
deleted file mode 100644
index 2179365..0000000
--- a/llm_trainer/worker/__init__.py
+++ /dev/null
@@ -1,5 +0,0 @@
-from __future__ import annotations
-
-from .client import RemoteWorkerClient, WorkerClientConfig, run_worker_client
-
-__all__ = ["RemoteWorkerClient", "WorkerClientConfig", "run_worker_client"]
diff --git a/llm_trainer/worker/client.py b/llm_trainer/worker/client.py
deleted file mode 100644
index e245886..0000000
--- a/llm_trainer/worker/client.py
+++ /dev/null
@@ -1,6 +0,0 @@
-"""Compatibility facade for remote worker clients."""
-
-from .client_core import CoordinatorHttpClient, WorkerClientConfig
-from .client_impl import RemoteWorkerClient, detect_worker_capabilities, run_worker_client
-
-__all__ = ["CoordinatorHttpClient", "WorkerClientConfig", "RemoteWorkerClient", "detect_worker_capabilities", "run_worker_client"]
diff --git a/llm_trainer/worker/client_core.py b/llm_trainer/worker/client_core.py
deleted file mode 100644
index f448ddc..0000000
--- a/llm_trainer/worker/client_core.py
+++ /dev/null
@@ -1,166 +0,0 @@
-from __future__ import annotations
-
-import json
-import os
-import platform
-import shutil
-import socket
-import time
-import zipfile
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Any, Optional
-from urllib.parse import urljoin
-from urllib.request import Request, urlopen
-
-import torch
-
-from llm_trainer.coordinator.artifacts import create_result_artifact_bundle
-from llm_trainer.contracts import (
- ArtifactSpec,
- BackendKind,
- ClaimJobRequest,
- ClaimJobResponse,
- CompleteJobRequest,
- DatasetSpec,
- FailJobRequest,
- HeartbeatRequest,
- ProgressReportRequest,
- ProtocolStatus,
- RegisterWorkerRequest,
- TrainingMetrics,
- TrainingResultSpec,
- WorkerAvailability,
- WorkerCapabilities,
-)
-from llm_trainer.contracts.jobs import JobStatus, TrainingJobSpec
-from llm_trainer.training_orchestrator import train_from_dataset
-
-try:
- import psutil
-except ImportError:
- psutil = None
-
-
-@dataclass
-class WorkerClientConfig:
- """Configuration for a remote worker client.
-
- Attributes:
- coordinator_url: Base URL for the coordinator API.
- worker_id: Stable worker identifier.
- device: Preferred training device.
- labels: Worker scheduling labels.
- heartbeat_interval_seconds: Seconds between heartbeats.
- execute_jobs: Whether to execute claimed jobs.
- claim_once: Whether to claim at most one job and exit.
- workspace_dir: Local folder used for downloaded jobs and outputs.
- """
-
- coordinator_url: str = "http://127.0.0.1:8765"
- worker_id: str = field(default_factory=lambda: f"{socket.gethostname()}-{os.getpid()}")
- device: str = "cuda" if torch.cuda.is_available() else "cpu"
- labels: list[str] = field(default_factory=list)
- heartbeat_interval_seconds: int = 10
- execute_jobs: bool = False
- claim_once: bool = False
- workspace_dir: Path = field(default_factory=lambda: Path.home() / ".drunkenbot_ide" / "worker_workspace")
-
-
-class CoordinatorHttpClient:
- """Small JSON HTTP client for the coordinator API."""
-
- def __init__(self, base_url: str) -> None:
- """Create an HTTP client.
-
- Args:
- base_url: Coordinator base URL.
- """
-
- self.base_url = base_url.rstrip("/")
-
- def get(self, path: str) -> dict[str, Any]:
- """Send a GET request.
-
- Args:
- path: API path.
-
- Returns:
- JSON response payload.
- """
-
- with urlopen(f"{self.base_url}{path}", timeout=10) as response:
- return json.loads(response.read().decode("utf-8"))
-
- def post(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
- """Send a JSON POST request.
-
- Args:
- path: API path.
- payload: Request payload.
-
- Returns:
- JSON response payload.
- """
-
- request = Request(
- f"{self.base_url}{path}",
- data=json.dumps(payload).encode("utf-8"),
- headers={"Content-Type": "application/json"},
- method="POST",
- )
- with urlopen(request, timeout=30) as response:
- return json.loads(response.read().decode("utf-8"))
-
- def download(self, path_or_url: str, output_path: Path) -> None:
- """Download a binary artifact.
-
- Args:
- path_or_url: Absolute URL or coordinator-relative path.
- output_path: Destination file path.
- """
-
- output_path.parent.mkdir(parents=True, exist_ok=True)
- with urlopen(self.absolute_url(path_or_url), timeout=300) as response, output_path.open("wb") as output:
- while chunk := response.read(1024 * 1024):
- output.write(chunk)
-
- def upload(self, path_or_url: str, input_path: Path) -> dict[str, Any]:
- """Upload a binary artifact.
-
- Args:
- path_or_url: Absolute URL or coordinator-relative path.
- input_path: Source file path.
-
- Returns:
- JSON response payload.
- """
-
- data = input_path.read_bytes()
- request = Request(
- self.absolute_url(path_or_url),
- data=data,
- headers={
- "Content-Type": "application/octet-stream",
- "Content-Length": str(len(data)),
- },
- method="PUT",
- )
- with urlopen(request, timeout=300) as response:
- return json.loads(response.read().decode("utf-8"))
-
- def absolute_url(self, path_or_url: str) -> str:
- """Build an absolute coordinator URL.
-
- Args:
- path_or_url: Absolute URL or coordinator-relative path.
-
- Returns:
- Absolute URL string.
- """
-
- if path_or_url.startswith(("http://", "https://")):
- return path_or_url
- return urljoin(f"{self.base_url}/", path_or_url.lstrip("/"))
-
-
diff --git a/llm_trainer/worker/client_impl.py b/llm_trainer/worker/client_impl.py
deleted file mode 100644
index 3016286..0000000
--- a/llm_trainer/worker/client_impl.py
+++ /dev/null
@@ -1,403 +0,0 @@
-from __future__ import annotations
-
-import json
-import os
-import platform
-import shutil
-import socket
-import time
-import zipfile
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Any, Optional
-from urllib.parse import urljoin
-from urllib.request import Request, urlopen
-
-import torch
-
-from llm_trainer.coordinator.artifacts import create_result_artifact_bundle
-from llm_trainer.contracts import (
- ArtifactSpec,
- BackendKind,
- ClaimJobRequest,
- ClaimJobResponse,
- CompleteJobRequest,
- DatasetSpec,
- FailJobRequest,
- HeartbeatRequest,
- ProgressReportRequest,
- ProtocolStatus,
- RegisterWorkerRequest,
- TrainingMetrics,
- TrainingResultSpec,
- WorkerAvailability,
- WorkerCapabilities,
-)
-from llm_trainer.contracts.jobs import JobStatus, TrainingJobSpec
-from llm_trainer.training_orchestrator import train_from_dataset
-
-from .client_core import CoordinatorHttpClient, WorkerClientConfig
-class RemoteWorkerClient:
- """Remote worker client that talks to the coordinator API."""
-
- def __init__(self, config: WorkerClientConfig) -> None:
- """Create a remote worker client.
-
- Args:
- config: Worker client configuration.
- """
-
- self.config = config
- self.http = CoordinatorHttpClient(config.coordinator_url)
- self.stop_requested = False
- self.pause_requested = False
- self.active_job_id: Optional[str] = None
-
- def register(self) -> dict[str, Any]:
- """Register this worker with the coordinator.
-
- Returns:
- Register response payload.
- """
-
- request = RegisterWorkerRequest(
- worker_id=self.config.worker_id,
- backend=BackendKind.REMOTE_CLIENT,
- device=self.config.device,
- capabilities=detect_worker_capabilities(),
- labels=self.config.labels,
- )
- response = self.http.post("/register", request.to_jsonable())
- if response.get("heartbeat_interval_seconds"):
- self.config.heartbeat_interval_seconds = int(response["heartbeat_interval_seconds"])
- return response
-
- def heartbeat(self, availability: WorkerAvailability, metrics: Optional[dict[str, Any]] = None) -> dict[str, Any]:
- """Send a heartbeat.
-
- Args:
- availability: Worker availability.
- metrics: Optional runtime metrics.
-
- Returns:
- Heartbeat response payload.
- """
-
- request = HeartbeatRequest(
- worker_id=self.config.worker_id,
- availability=availability,
- backend=BackendKind.REMOTE_CLIENT,
- active_job_id=self.active_job_id,
- device=self.config.device,
- metrics=metrics or {},
- )
- response = self.http.post("/heartbeat", request.to_jsonable())
- self.stop_requested = bool(response.get("should_stop_job"))
- self.pause_requested = bool(response.get("should_pause_job"))
- return response
-
- def claim_job(self) -> Optional[TrainingJobSpec]:
- """Ask the coordinator for a compatible job.
-
- Returns:
- Assigned job when available.
- """
-
- request = ClaimJobRequest(
- worker_id=self.config.worker_id,
- backend=BackendKind.REMOTE_CLIENT,
- capabilities=detect_worker_capabilities(),
- )
- response = ClaimJobResponse.from_jsonable(self.http.post("/claim-job", request.to_jsonable()))
- if response.status != ProtocolStatus.OK:
- raise RuntimeError(response.message)
- return response.job
-
- def run_forever(self) -> None:
- """Run the worker loop."""
-
- self.register()
- while True:
- self.heartbeat(WorkerAvailability.AVAILABLE)
- if not self.config.execute_jobs and not self.config.claim_once:
- time.sleep(self.config.heartbeat_interval_seconds)
- continue
- job = self.claim_job()
- if job is None:
- if self.config.claim_once:
- return
- time.sleep(self.config.heartbeat_interval_seconds)
- continue
- job = self.sync_job_artifacts(job)
- self.active_job_id = job.job_id
- try:
- if self.config.execute_jobs:
- self.execute_job(job)
- else:
- self.fail_job(job.job_id, "Worker execution disabled. Run with --execute to train jobs.", retryable=True)
- finally:
- self.active_job_id = None
- if self.config.claim_once:
- return
-
- def execute_job(self, job: TrainingJobSpec) -> None:
- """Execute a claimed training job.
-
- Args:
- job: Claimed training job.
- """
-
- self.stop_requested = False
- self.pause_requested = False
-
- def progress(event: Any) -> None:
- metrics = _event_to_metrics(event)
- response = self.http.post(
- "/progress",
- ProgressReportRequest(self.config.worker_id, job.job_id, metrics).to_jsonable(),
- )
- self.stop_requested = bool(response.get("should_stop_job"))
- self.pause_requested = bool(response.get("should_pause_job"))
- while self.pause_requested and not self.stop_requested:
- time.sleep(self.config.heartbeat_interval_seconds)
- heartbeat_response = self.heartbeat(WorkerAvailability.BUSY, {"paused": True})
- self.pause_requested = bool(heartbeat_response.get("should_pause_job"))
-
- try:
- result = train_from_dataset(
- job.dataset.dataset_dir,
- job.model.config,
- job.training,
- progress=progress,
- should_stop=lambda: self.stop_requested,
- )
- except Exception as exc:
- self.fail_job(job.job_id, str(exc), retryable=False)
- print(f"Job {job.job_id} failed on worker {self.config.worker_id}: {exc}")
- return
- status = JobStatus.CANCELLED if result.stopped else JobStatus.COMPLETED
- artifact_bundle_url = self.upload_result_artifacts(job)
- self.http.post(
- "/complete",
- CompleteJobRequest(
- self.config.worker_id,
- TrainingResultSpec(
- job_id=job.job_id,
- status=status,
- checkpoint_path=result.checkpoint_path,
- summary_path=result.summary_path,
- final_train_loss=result.final_train_loss,
- final_val_loss=result.final_val_loss,
- stopped=result.stopped,
- artifact_bundle_url=artifact_bundle_url,
- ),
- ).to_jsonable(),
- )
-
- def sync_job_artifacts(self, job: TrainingJobSpec) -> TrainingJobSpec:
- """Download and localize remote job artifacts.
-
- Args:
- job: Claimed job from the coordinator.
-
- Returns:
- Job rewritten to worker-local paths.
- """
-
- bundle_url = str(job.metadata.get("artifact_bundle_url") or "")
- if not bundle_url:
- return job
- workspace = self._job_workspace(job.job_id)
- bundle_path = workspace / "input_bundle.zip"
- extract_dir = workspace / "input"
- self.http.download(bundle_url, bundle_path)
- if extract_dir.exists():
- shutil.rmtree(extract_dir)
- extract_dir.mkdir(parents=True, exist_ok=True)
- _safe_extract_zip(bundle_path, extract_dir)
- dataset_dir = extract_dir / "dataset"
- if not dataset_dir.exists():
- raise FileNotFoundError(f"Downloaded job bundle does not contain dataset/: {bundle_url}")
- output_dir = workspace / "model"
- output_dir.mkdir(parents=True, exist_ok=True)
- job.dataset = DatasetSpec.from_dataset_dir(dataset_dir)
- job.training.output_dir = output_dir
- job.artifacts = ArtifactSpec.from_output_dir(output_dir)
- resume_artifact = job.metadata.get("resume_checkpoint_artifact")
- if resume_artifact:
- resume_path = extract_dir / str(resume_artifact)
- if resume_path.is_file():
- job.training.resume_from_checkpoint = resume_path
- base_artifact = job.metadata.get("base_checkpoint_artifact")
- if base_artifact:
- base_path = extract_dir / str(base_artifact)
- if base_path.is_file():
- job.training.fine_tune_from_checkpoint = base_path
- job.model.base_checkpoint = base_path
- return job
-
- def upload_result_artifacts(self, job: TrainingJobSpec) -> Optional[str]:
- """Upload worker output artifacts to the coordinator.
-
- Args:
- job: Completed job.
-
- Returns:
- Coordinator artifact URL when upload succeeds.
- """
-
- output_dir = Path(job.training.output_dir)
- if not output_dir.exists():
- return None
- bundle_path = self._job_workspace(job.job_id) / "result_bundle.zip"
- create_result_artifact_bundle(job.job_id, output_dir, bundle_path)
- remote_path = f"/artifacts/results/{job.job_id}/{bundle_path.name}"
- response = self.http.upload(remote_path, bundle_path)
- return str(response.get("artifact_url") or remote_path)
-
- def _job_workspace(self, job_id: str) -> Path:
- """Return the worker-local workspace for a job.
-
- Args:
- job_id: Training job identifier.
-
- Returns:
- Worker-local job workspace.
- """
-
- workspace = Path(self.config.workspace_dir) / job_id
- workspace.mkdir(parents=True, exist_ok=True)
- return workspace
-
- def fail_job(self, job_id: str, error: str, retryable: bool) -> None:
- """Report job failure to the coordinator.
-
- Args:
- job_id: Job identifier.
- error: Failure text.
- retryable: Whether the job may be retried.
- """
-
- self.http.post("/fail", FailJobRequest(self.config.worker_id, job_id, error, retryable).to_jsonable())
-
-
-def detect_worker_capabilities() -> WorkerCapabilities:
- """Detect local worker hardware capabilities.
-
- Returns:
- Worker capabilities.
- """
-
- gpu_names: list[str] = []
- total_vram_gb: Optional[float] = None
- if torch.cuda.is_available():
- total_vram_bytes = 0
- for index in range(torch.cuda.device_count()):
- properties = torch.cuda.get_device_properties(index)
- gpu_names.append(properties.name)
- total_vram_bytes += int(properties.total_memory)
- total_vram_gb = total_vram_bytes / (1024**3)
- system_ram_gb = None
- if psutil is not None:
- system_ram_gb = psutil.virtual_memory().total / (1024**3)
- return WorkerCapabilities(
- hostname=socket.gethostname(),
- platform=f"{platform.system()} {platform.release()}",
- cpu_count=os.cpu_count(),
- system_ram_gb=system_ram_gb,
- gpu_names=gpu_names,
- total_vram_gb=total_vram_gb,
- supports_cuda=torch.cuda.is_available(),
- supports_bf16=bool(torch.cuda.is_available() and torch.cuda.is_bf16_supported()),
- supports_fp16=torch.cuda.is_available(),
- )
-
-
-def run_worker_client(config: WorkerClientConfig) -> None:
- """Run a remote worker client.
-
- Args:
- config: Worker client configuration.
- """
-
- RemoteWorkerClient(config).run_forever()
-
-
-def _event_to_metrics(event: Any) -> TrainingMetrics:
- """Convert a training progress event into protocol metrics.
-
- Args:
- event: Progress event.
-
- Returns:
- Training metrics.
- """
-
- if not isinstance(event, dict):
- return TrainingMetrics(message=str(event))
- return TrainingMetrics(
- step=_int_or_none(event.get("step")),
- total_steps=_int_or_none(event.get("total_steps")),
- epoch=_int_or_none(event.get("epoch")),
- total_epochs=_int_or_none(event.get("epochs") or event.get("total_epochs")),
- train_loss=_float_or_none(event.get("loss") or event.get("train_loss")),
- val_loss=_float_or_none(event.get("val_loss")),
- learning_rate=_float_or_none(event.get("learning_rate") or event.get("lr")),
- tokens_per_second=_float_or_none(event.get("tokens_per_second") or event.get("tokens_per_sec")),
- samples_per_second=_float_or_none(event.get("samples_per_second") or event.get("samples_per_sec")),
- gpu_memory_percent=_float_or_none(event.get("gpu_memory_percent")),
- system_ram_percent=_float_or_none(event.get("system_ram_percent")),
- message=str(event.get("message")) if event.get("message") is not None else None,
- )
-
-
-def _int_or_none(value: Any) -> Optional[int]:
- """Convert a value to int when possible.
-
- Args:
- value: Input value.
-
- Returns:
- Integer or None.
- """
-
- try:
- return int(value) if value is not None else None
- except (TypeError, ValueError):
- return None
-
-
-def _float_or_none(value: Any) -> Optional[float]:
- """Convert a value to float when possible.
-
- Args:
- value: Input value.
-
- Returns:
- Float or None.
- """
-
- try:
- return float(value) if value is not None else None
- except (TypeError, ValueError):
- return None
-
-
-def _safe_extract_zip(zip_path: Path, target_dir: Path) -> None:
- """Extract a zip file without allowing path traversal.
-
- Args:
- zip_path: Zip file path.
- target_dir: Destination directory.
-
- Raises:
- ValueError: If a zip member would escape the target directory.
- """
-
- root = Path(target_dir).resolve()
- with zipfile.ZipFile(zip_path) as archive:
- for member in archive.infolist():
- member_path = (root / member.filename).resolve()
- if root not in member_path.parents and member_path != root:
- raise ValueError(f"Unsafe artifact member path: {member.filename}")
- archive.extractall(root)
diff --git a/pyproject.toml b/pyproject.toml
index fb8e6f6..dc423d4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,6 +30,7 @@ ignore = [
[tool.ruff.lint.per-file-ignores]
"tests/**/*.py" = ["D", "ANN"]
"tools/**/*.py" = ["D", "ANN"]
+"interface/**/*.py" = ["E402", "E501", "F401", "F403", "F405", "F821", "F541", "F841"]
[tool.ruff.lint.pydocstyle]
convention = "google"
diff --git a/tools/check_dependency_boundaries.py b/tools/check_dependency_boundaries.py
index cf30eae..944fb66 100644
--- a/tools/check_dependency_boundaries.py
+++ b/tools/check_dependency_boundaries.py
@@ -44,10 +44,7 @@ def _imports(path: Path) -> list[tuple[int, str]]:
def find_violations(root: Path) -> list[BoundaryViolation]:
"""Find imports that violate the engine/interface dependency boundary."""
- rules = (
- ("engine", "interface"),
- ("interface", "llm_trainer"),
- )
+ rules = (("engine", "interface"),)
violations: list[BoundaryViolation] = []
for package_name, forbidden_root in rules:
package_root = root / package_name
@@ -86,8 +83,7 @@ def main() -> int:
print(f" {violation.path}:{violation.line}: {violation.imported}")
return 1
print(
- "Dependency boundaries are clean: engine -> interface and "
- "interface -> llm_trainer are forbidden."
+ "Dependency boundaries are clean: engine cannot import interface."
)
return 0
diff --git a/tools/curriculum_finetune.py b/tools/curriculum_finetune.py
index 0ccc835..26cf23d 100644
--- a/tools/curriculum_finetune.py
+++ b/tools/curriculum_finetune.py
@@ -1,6 +1,5 @@
from __future__ import annotations
-from pathlib import Path
try:
from .curriculum_shared import *
except ImportError:
@@ -461,4 +460,3 @@ def instruction_fine_tune_blocks(count: int, topic: str) -> list[str]:
return blocks
-
diff --git a/tools/curriculum_shared.py b/tools/curriculum_shared.py
index a147a79..845bf4b 100644
--- a/tools/curriculum_shared.py
+++ b/tools/curriculum_shared.py
@@ -1,5 +1,7 @@
from __future__ import annotations
+from pathlib import Path
+
NAMES = [
"Mina", "Ravi", "Lena", "Omar", "Sara", "Tara", "Jin", "Ada",
"Leo", "Nia", "Sam", "Priya",
@@ -196,4 +198,3 @@ def write_blocks(path: Path, blocks: list[str], min_unique_ratio: float = 0.9) -
temp_path = path.with_suffix(path.suffix + ".tmp")
temp_path.write_text("\n\n".join(blocks).strip() + "\n", encoding="utf-8")
temp_path.replace(path)
-
diff --git a/tools/curriculum_subjects.py b/tools/curriculum_subjects.py
index 0be8a2f..1b5305a 100644
--- a/tools/curriculum_subjects.py
+++ b/tools/curriculum_subjects.py
@@ -1,6 +1,5 @@
from __future__ import annotations
-from pathlib import Path
try:
from .curriculum_shared import *
except ImportError:
@@ -395,4 +394,3 @@ def code_blocks(count: int) -> list[str]:
return blocks
-