diff --git a/.github/workflows/quality.yml b/.github/workflows/quality.yml new file mode 100644 index 0000000..5a8375d --- /dev/null +++ b/.github/workflows/quality.yml @@ -0,0 +1,24 @@ +name: Code quality + +on: + push: + pull_request: + +jobs: + quality: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Install quality tools + run: python -m pip install -r requirements-dev.txt + - name: Check formatting and lint + run: ruff check interface run_app.py tools tests --select E,F --ignore E501,F403,F405 + - name: Compile interface modules + run: python -m compileall -q interface run_app.py + - name: Check Python file sizes + run: python tools/check_code_standards.py + - name: Check package dependency boundaries + run: python tools/check_dependency_boundaries.py diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..3eb249d --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "engine"] + path = engine + url = https://github.com/drunkenbot-ai/engine.git diff --git a/README.md b/README.md index d6d5731..c9117f9 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,10 @@ creating tokenizers, training small GPT-style language models, benchmarking checkpoints, exporting model artifacts, and testing local GGUF models in a streamed Markdown chat interface. +The source is split into the non-Qt `engine/` submodule and the Qt desktop +`interface/` package. The dependency direction is one-way: `engine/` never +imports `interface/`. + For development, use Python 3.12 or newer. Distribution builds include a private Python runtime and do not require Python on the user's machine. @@ -74,22 +78,28 @@ license data, logs, and machine identifier under the user's ## First backend commands +The non-Qt command-line engine is available directly: + +```powershell +python -m engine.cli --help +``` + Prepare text/PDF/JSONL files: ```powershell -python -m llm_trainer.cli prepare --input_dir .\examples\tiny_corpus --output_dir .\runs\tiny_data --context_length 16 +python -m engine.cli prepare --input_dir .\examples\tiny_corpus --output_dir .\runs\tiny_data --context_length 16 ``` Prepare programming PDFs plus source files in code-aware mode: ```powershell -python -m llm_trainer.cli prepare --input_dir .\examples\tiny_corpus --output_dir .\runs\code_data --context_length 128 --code_training_mode +python -m engine.cli prepare --input_dir .\examples\tiny_corpus --output_dir .\runs\code_data --context_length 128 --code_training_mode ``` For large corpora, enable faster preview/build scanning: ```powershell -python -m llm_trainer.cli prepare --input_dir .\my_big_data --output_dir .\runs\big_data --fast_scan_mode +python -m engine.cli prepare --input_dir .\my_big_data --output_dir .\runs\big_data --fast_scan_mode ``` `--fast_scan_mode` uses cheaper file fingerprints and cached preview statistics, @@ -98,7 +108,7 @@ which is significantly faster on very large datasets. To reduce false duplicate matches in fast mode, add strict verification: ```powershell -python -m llm_trainer.cli prepare --input_dir .\my_big_data --output_dir .\runs\big_data --fast_scan_mode --strict_duplicate_verification --fast_scan_sample_bytes 65536 +python -m engine.cli prepare --input_dir .\my_big_data --output_dir .\runs\big_data --fast_scan_mode --strict_duplicate_verification --fast_scan_sample_bytes 65536 ``` `--strict_duplicate_verification` only runs full SHA-256 hashing on suspected @@ -111,7 +121,7 @@ tries to extract code-like blocks from PDFs/text. Train a very small smoke-test model: ```powershell -python -m llm_trainer.cli train --data_dir .\runs\tiny_data --output_dir .\runs\tiny_model --epochs 1 --batch_size 2 --context_length 16 --embedding_size 32 --head_count 4 --layer_count 2 --device cpu --no_resume +python -m engine.cli train --data_dir .\runs\tiny_data --output_dir .\runs\tiny_model --epochs 1 --batch_size 2 --context_length 16 --embedding_size 32 --head_count 4 --layer_count 2 --device cpu --no_resume ``` Training saves checkpoints in the model folder and can resume from the latest @@ -163,13 +173,19 @@ contains a real Hugging Face-compatible `hf_model` directory. For MicroGPT checkpoints, use the HF-style package export first: ```bash -python -m llm_trainer.cli export-hf --model_dir runs/model +python -m engine.cli export-hf --model_dir runs/model ``` That creates `runs/model/hf_model` with config, weights, tokenizer metadata, lineage, and a README. It is portable MicroGPT packaging, not a claim that the checkpoint is already a llama.cpp-supported Llama/Mistral/Gemma model. +Before submitting changes, run the package-boundary check: + +```powershell +python tools/check_dependency_boundaries.py +``` + ## Current IDE Features - Dataset Blueprint with dynamic bundled corpus discovery. diff --git a/UserGuide.md b/UserGuide.md index 648830b..a5425c3 100644 --- a/UserGuide.md +++ b/UserGuide.md @@ -37,9 +37,14 @@ their own private Python runtime and do not require Python to be installed system-wide. For installer builds, see [build.md](build.md). The installer excludes the -bundled `llm_trainer/default_data` training corpus; projects receive only the +bundled `engine/default_data` training corpus; projects receive only the training data that you explicitly copy or add. +The application code is organized into `engine/` (non-Qt data, training, and +worker services) and `interface/` (the Qt desktop UI). Existing +`llm_trainer.*` imports and `python -m llm_trainer.cli` remain compatibility +aliases, but new code should import from the canonical packages. + The app has five main work areas: - `IN`: prepare datasets. diff --git a/build.md b/build.md index d7e5294..3ae5625 100644 --- a/build.md +++ b/build.md @@ -39,8 +39,10 @@ python packager.py --gpu ``` The packager creates a private runtime, installs the pinned dependencies, -builds the launcher bundle, copies application assets and fonts, runs Inno -Setup, and writes the installer to `packaging/artifacts/`. +builds the launcher bundle, copies `engine/` and `interface/` directly along +with application assets and fonts, runs Inno Setup, and writes the installer to +`packaging/artifacts/`. The historical `llm_trainer/` compatibility package is +kept in the source tree but is not the primary packaged application path. The CUDA build detects the NVIDIA driver using `nvidia-smi`. It selects the supported PyTorch wheel index and falls back to CPU when no compatible driver @@ -62,6 +64,10 @@ The installer includes the private Python runtime, application code, third party packages, fonts, logo images, and other application assets. Training corpus data is not bundled; users download or select it through the Dataset Sources page after installation. +Run `python tools/check_dependency_boundaries.py` before packaging to verify +that the non-Qt engine does not import the desktop interface and that new +interface code does not depend on the legacy package. + Build intermediates and installers are written under `packaging/` and are ignored by Git. diff --git a/engine b/engine new file mode 160000 index 0000000..8d55bc8 --- /dev/null +++ b/engine @@ -0,0 +1 @@ +Subproject commit 8d55bc88dac2ca13c64fa8796f4975fc345c2036 diff --git a/llm_trainer/ui/__init__.py b/interface/__init__.py similarity index 98% rename from llm_trainer/ui/__init__.py rename to interface/__init__.py index 2d4cdc6..58819b3 100644 --- a/llm_trainer/ui/__init__.py +++ b/interface/__init__.py @@ -1 +1,2 @@ """PySide6 desktop interface for Micro Trainer.""" + diff --git a/interface/app.py b/interface/app.py new file mode 100644 index 0000000..356bde5 --- /dev/null +++ b/interface/app.py @@ -0,0 +1,338 @@ +from __future__ import annotations + +import ctypes +from datetime import datetime +import html +import importlib +import json +from functools import partial +import logging +import math +import os +from queue import Empty, Queue +import re +import shutil +import signal +import sqlite3 +import subprocess +import sys +from pathlib import Path +from threading import Event, Thread +from typing import Any, Optional, Union + +import torch +from PySide6.QtCore import QObject, QEvent, QPoint, Qt, QThread, QTimer, Slot, qInstallMessageHandler +from PySide6.QtGui import QBrush, QColor, QFont, QIcon, QPainter, QPen, QPixmap, QPolygon +from PySide6.QtGui import QFontDatabase +from PySide6.QtWidgets import ( + QApplication, + QAbstractButton, + QComboBox, + QDoubleSpinBox, + QDialog, + QFileDialog, + QFormLayout, + QGridLayout, + QHBoxLayout, + QInputDialog, + QLabel, + QLineEdit, + QListWidget, + QListWidgetItem, + QTreeWidget, + QTreeWidgetItem, + QMainWindow, + QMessageBox, + QProgressBar, + QPushButton, + QSizePolicy, + QStackedWidget, + QSpinBox, + QTextBrowser, + QTextEdit, + QVBoxLayout, + QWidget, +) + +from engine.app_logging import qt_message_handler, setup_logging +from engine.app_logging import DEFAULT_LOG_DIR +from engine.config import DatasetConfig, ModelConfig, TrainingConfig +from engine.conversation_datasets import CONVERSATION_DATASET_PRESETS, dataset_ids_for_stage, dataset_stage_label +from engine.contracts import BackendKind +from engine.contracts.jobs import RuntimeSpec, TrainingJobSpec +from engine.coordinator import CoordinatorApiServer, JobManager, create_job_artifact_bundle +from engine.evaluation import DEFAULT_BENCHMARK_PROMPTS, evaluate_checkpoint, normalize_prompts +from engine.export import export_gguf_with_llama_cpp, export_hf_microgpt_package, export_llama_adapter_package, export_project_bundle, quantize_checkpoint +from engine.fine_tuning_service import run_fine_tuning_job +from engine.llama_chat import LlamaChatSession, load_llama_chat_session, stream_chat_reply +from engine.lineage import read_json +from engine.microgpt_chat import load_microgpt_chat_session, stream_microgpt_chat_reply +from engine.notifier import NotificationManager, default_notifier_config_path, ensure_notifier_config +from engine.runpod_cloud import ( + RunPodClient, + RunPodConfig, + create_runpod_worker_bundle, + default_runpod_config_path, + ensure_runpod_config, + load_runpod_config, + public_url_is_cloud_reachable, + save_runpod_config, +) +from engine.dataset_build import build_dataset +from engine.dataset_preview import check_project_health, scan_dataset_preview +from engine.telemetry_store import initialize_store, insert_metric, latest_run, rows_until, telemetry_db_path +from engine.training import check_resume_compatibility, latest_checkpoint +from engine.training_planning import estimate_training_resources, format_bytes +from engine.training_service import run_training_job +from engine.external_dataset import ( + DEFAULT_MANIFEST_URL, + download_latest_dataset, + is_newer_version, + load_manifest, +) +from interface.chat_widgets import ChatMessageWidget +from interface.markdown_renderer import markdown_to_html +from interface.workers import ProcessTaskWorker, TaskWorker, WorkerSignalBridge +from interface.startup_splash import StartupSplash +from interface.tabs.benchmark_tab import build_benchmark_tab +from interface.tabs.chat_tab import build_chat_tab +from interface.tabs.dataset_tab import build_dataset_tab +from interface.tabs.dataset_plan_tab import ( + build_dataset_plan_tab, + default_data_root, + default_data_stage, + dataset_plan_defaults, + iter_default_data_files, + populate_default_data_tree, +) +from interface.tabs.live_tab import build_live_training_tab +from interface.tabs.training_tab import build_training_tab +from interface.tabs.export_tab import build_export_tab +from interface.tabs.fine_tuning_tab import build_fine_tuning_tab +from interface.tabs.job_manager_tab import build_job_manager_tab, set_table_rows +from engine.license_client import load_stored_license_key +from interface.license_activation_dialog import LicenseActivationDialog, run_license_check_responsively + +try: + import psutil +except ImportError: + psutil = None + + +APP_NAME = "DrunkenBot LLM-IDE" +# Bump on every release that should require a version-ceiling check against +# licenses -- this is what license_client.check_license_at_launch compares +# against a license's version_ceiling/grace_period_until. +APP_VERSION = "1.0.0" +# TODO: point at the real deployed cloud-service URL once it has one. +# Overridable via env var so ops can point a build at a different +# deployment (dev/staging/prod) without a code change or rebuild. +# LICENSE_SERVER_URL = os.environ.get("DRUNKENBOT_LICENSE_SERVER_URL", "https://license.drunkenbot.ai") +# LICENSE_SERVER_URL = "http://127.0.0.1:8000/" +LICENSE_SERVER_URL = "https://drunkenbot.store" +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 + + +from interface.startup import ( + ProjectChoiceDialog, StartupValidationSplash, _apply_windows_taskbar_icon, + _load_recent_projects, _register_recent_project, _run_startup_tests, + _run_startup_validations, _validate_writable_directory, +) +from interface.main_window_part1 import MainWindowPart1 +from interface.main_window_part2 import MainWindowPart2 +from interface.main_window_part3 import MainWindowPart3 +from interface.main_window_part4 import MainWindowPart4 +from interface.main_window_part5 import MainWindowPart5 +from interface.main_window_part6 import MainWindowPart6 +from interface.main_window_part7 import MainWindowPart7 +from interface.main_window_part8 import MainWindowPart8 +from interface.main_window_part9 import MainWindowPart9 +from interface.main_window_part10 import MainWindowPart10 +from interface.main_window_part11 import MainWindowPart11 +from interface.main_window_part12 import MainWindowPart12 +from interface.main_window_part13 import MainWindowPart13 +from interface.main_window_part14 import MainWindowPart14 +from interface.main_window_part15 import MainWindowPart15 +from interface.main_window_part16 import MainWindowPart16 +from interface.main_window_part17 import MainWindowPart17 +from interface.main_window_part18 import MainWindowPart18 + +class MainWindow(MainWindowPart1, MainWindowPart2, MainWindowPart3, MainWindowPart4, MainWindowPart5, MainWindowPart6, MainWindowPart7, MainWindowPart8, MainWindowPart9, MainWindowPart10, MainWindowPart11, MainWindowPart12, MainWindowPart13, MainWindowPart14, MainWindowPart15, MainWindowPart16, MainWindowPart17, MainWindowPart18, QMainWindow): + """Main application window composed from focused UI mixins.""" + +def _ensure_valid_license(splash: "StartupValidationSplash") -> bool: + """Block app launch until a valid license is confirmed. + + Checks the currently stored license key (if any). On failure, shows + :class:`LicenseActivationDialog` in a loop -- unlike the general + startup-validation flow elsewhere in ``main()``, there is deliberately + no "continue anyway" option here: an unlicensed launch is not a + degraded-but-usable state, it's the one thing this app must not do. + + Args: + splash: Startup splash screen, used to show progress. + + Returns: + True if the app is licensed to proceed, False if the user cancelled + activation and the app should exit. + """ + + splash.append_log("Checking license...") + QApplication.processEvents() + + stored_key = load_stored_license_key() + if stored_key: + result = run_license_check_responsively(APP_VERSION, LICENSE_SERVER_URL) + if result.valid: + splash.append_log( + "[OK] License valid" + + (" (offline grace period)" if result.used_offline_grace else "") + ) + return True + initial_message = result.reason + else: + initial_message = "No license activated on this machine yet." + + # A QSplashScreen-style window is designed to stay on top of other + # windows during startup -- which means it can end up covering a newly + # created dialog instead of the other way around. Hide it while the + # dialog is up rather than fight window-stacking order; it isn't doing + # anything useful to look at during activation anyway. + splash.hide() + try: + while True: + dialog = LicenseActivationDialog(APP_VERSION, LICENSE_SERVER_URL, initial_message) + dialog.setWindowIcon(MainWindow._static_app_icon()) + dialog.show() + dialog.raise_() + dialog.activateWindow() + if dialog.exec() != QDialog.Accepted: + LOGGER.info("License activation cancelled by user; exiting.") + QApplication.instance().setProperty("startup_aborted", True) + return False + splash.append_log("[OK] License activated") + return True + finally: + splash.show() + splash.raise_() + + +def main(app: Optional[QApplication] = None, splash: Optional[StartupSplash] = None) -> None: + """Launch the PySide6 desktop application.""" + + owns_app = app is None + log_file = setup_logging() + qInstallMessageHandler(qt_message_handler) + LOGGER.info("Starting %s. Log file: %s", APP_NAME, log_file) + if sys.platform == "win32": + try: + ctypes.windll.shell32.SetCurrentProcessExplicitAppUserModelID(WINDOWS_APP_ID) + except Exception: + LOGGER.exception("Could not set Windows app user model ID") + app = app or QApplication(sys.argv) + app.setFont(QFont("Arial", 10)) + app.setWindowIcon(MainWindow._static_app_icon()) + splash = splash or StartupValidationSplash() + splash.setWindowIcon(MainWindow._static_app_icon()) + splash.show() + QTimer.singleShot(0, lambda: _apply_windows_taskbar_icon(splash)) + QApplication.processEvents() + if not _ensure_valid_license(splash): + splash.close() + if not owns_app: + app.quit() + app.setProperty("startup_aborted", True) + return + try: + _run_startup_validations(splash) + except Exception as exc: + LOGGER.exception("Startup validation failed") + splash.append_log(f"[FAIL] Startup blocked: {exc}") + splash.close() + proceed = QMessageBox.question( + None, + "Startup validation failed", + "One or more startup checks failed.\n\n" + f"{exc}\n\n" + "Do you want to continue anyway?", + QMessageBox.Yes | QMessageBox.No, + QMessageBox.No, + ) + if proceed != QMessageBox.Yes: + if not owns_app: + app.quit() + app.setProperty("startup_aborted", True) + return + LOGGER.warning("User chose to continue after failed startup validation.") + splash.close() + while True: + chooser = ProjectChoiceDialog() + chooser.setWindowIcon(MainWindow._static_app_icon()) + QTimer.singleShot(0, lambda dialog=chooser: _apply_windows_taskbar_icon(dialog)) + if chooser.exec() != QDialog.Accepted: + LOGGER.info("Startup closed at project selection screen") + if not owns_app: + app.quit() + app.setProperty("startup_aborted", True) + return + window = MainWindow() + try: + if chooser.choice == "new": + base_dir = QFileDialog.getExistingDirectory( + None, + "Choose folder where the new project will be created", + str(DEFAULT_PROJECTS_DIR), + ) + if not base_dir: + window.deleteLater() + continue + project_name, ok = QInputDialog.getText(None, "Project name", "Enter project name:", text="MicroLLMProject") + if not ok: + window.deleteLater() + continue + project_name = project_name.strip() or "MicroLLMProject" + window._create_project_at(project_name, Path(base_dir)) + elif chooser.choice == "open": + project_file, _ = QFileDialog.getOpenFileName( + None, + "Open Micro LLM project", + str(DEFAULT_PROJECTS_DIR), + "Micro LLM project (project.json *.json);;All files (*)", + ) + if not project_file: + window.deleteLater() + continue + window._open_project_file(Path(project_file)) + elif chooser.choice == "recent": + if chooser.selected_project_file is None: + window.deleteLater() + continue + window._open_project_file(chooser.selected_project_file) + elif chooser.choice == "test_local_llm": + window.show_chat_only_mode() + except Exception as exc: + LOGGER.exception("Project setup failed during startup") + QMessageBox.critical(None, "Project setup failed", f"Could not complete project setup.\n\n{exc}") + window.deleteLater() + continue + break + window.show() + QTimer.singleShot(0, window.apply_windows_taskbar_icon) + interrupt_timer = QTimer() + interrupt_timer.timeout.connect(lambda: None) + interrupt_timer.start(200) + window.interrupt_timer = interrupt_timer + signal.signal(signal.SIGINT, lambda *_: QTimer.singleShot(0, window.request_shutdown_from_signal)) + if owns_app: + sys.exit(app.exec()) + + +if __name__ == "__main__": + main() diff --git a/llm_trainer/ui/charts.py b/interface/charts.py similarity index 99% rename from llm_trainer/ui/charts.py rename to interface/charts.py index b16fe67..eb8c9f1 100644 --- a/llm_trainer/ui/charts.py +++ b/interface/charts.py @@ -397,3 +397,4 @@ def _nearest_point(self, x_value: float, y_value: float) -> Optional[tuple[str, ) distance = ((nearest[1] - x_value) / x_span) ** 2 + ((nearest[2] - y_value) / y_span) ** 2 return nearest if distance < 0.01 else None + diff --git a/llm_trainer/ui/chat_widgets.py b/interface/chat_widgets.py similarity index 98% rename from llm_trainer/ui/chat_widgets.py rename to interface/chat_widgets.py index 317248e..c38c732 100644 --- a/llm_trainer/ui/chat_widgets.py +++ b/interface/chat_widgets.py @@ -97,12 +97,12 @@ def __init__( footer = QHBoxLayout() footer.setContentsMargins(6, 0, 6, 0) - self.copy_button = QPushButton("⧉") + self.copy_button = QPushButton("Copy") self.copy_button.setObjectName("MessageAction") self.copy_button.setFixedWidth(28) self.copy_button.setToolTip("Copy this message.") self.copy_button.clicked.connect(lambda: QApplication.clipboard().setText(self.browser.toPlainText())) - self.resend_button = QPushButton("↻") + self.resend_button = QPushButton("Resend") self.resend_button.setObjectName("MessageAction") self.resend_button.setFixedWidth(28) self.resend_button.setToolTip("Send this message again.") diff --git a/llm_trainer/ui/drunkenbot_llm_ide.ico b/interface/drunkenbot_llm_ide.ico similarity index 100% rename from llm_trainer/ui/drunkenbot_llm_ide.ico rename to interface/drunkenbot_llm_ide.ico diff --git a/llm_trainer/ui/license_activation_dialog.py b/interface/license_activation_dialog.py similarity index 98% rename from llm_trainer/ui/license_activation_dialog.py rename to interface/license_activation_dialog.py index 6a17962..0956945 100644 --- a/llm_trainer/ui/license_activation_dialog.py +++ b/interface/license_activation_dialog.py @@ -13,7 +13,7 @@ QVBoxLayout, ) -from llm_trainer.license_client import ( +from engine.license_client import ( LicenseCheckResult, check_license_at_launch, store_license_key, @@ -170,4 +170,4 @@ def _show_status(self, message: str) -> None: """ self._status.setVisible(True) - self._status.setPlainText(message) \ No newline at end of file + self._status.setPlainText(message) diff --git a/llm_trainer/ui/live_widgets.py b/interface/live_widgets.py similarity index 99% rename from llm_trainer/ui/live_widgets.py rename to interface/live_widgets.py index 56330b6..82f567c 100644 --- a/llm_trainer/ui/live_widgets.py +++ b/interface/live_widgets.py @@ -374,3 +374,4 @@ def update_flow(self, layer_count: int, grad_norm: Optional[float], step: int) - self.bars.setOpts(x0=[0] * count, x1=values, y=self.layers, height=0.55, brush="#ff4ca8") self.plot.setYRange(0, count + 1) self.plot.setXRange(0, max_value * 1.15) + diff --git a/interface/main_window_part1.py b/interface/main_window_part1.py new file mode 100644 index 0000000..c149431 --- /dev/null +++ b/interface/main_window_part1.py @@ -0,0 +1,437 @@ +from __future__ import annotations + +# MainWindow implementation mixin. Runtime names are provided by app.py. +from typing import Any +from . import app as _app + +globals().update({name: value for name, value in vars(_app).items() if not name.startswith("__")}) + +class MainWindowPart1: + def __init__(self) -> None: + """Create the main application window.""" + + super().__init__() + self.log_file_path = setup_logging() + LOGGER.info("Creating %s main window", APP_NAME) + if QApplication.instance(): + QApplication.instance().setFont(QFont("Arial", 10)) + self.setWindowTitle(APP_NAME) + self.setWindowIcon(self._app_icon()) + self._windows_icon_handles: list[int] = [] + self.resize(1240, 820) + self.thread: Optional[QThread] = None + self.worker: Optional[TaskWorker] = None + self.result_bridge: Optional[WorkerSignalBridge] = None + self.stop_event: Optional[Event] = None + self.progress_queue: Optional[Queue] = None + self.active_log: Optional[QTextEdit] = None + self.active_progress_bar: Optional[QProgressBar] = None + self.active_button: Optional[QPushButton] = None + self.active_stop_button: Optional[QPushButton] = None + self.active_button_text = "" + self.active_button_restore_text = "" + self.active_task_kind = "" + self.notification_manager: Optional[NotificationManager] = None + self.current_project_file: Optional[Path] = None + self.telemetry_db_path: Optional[Path] = None + self.telemetry_run_id = "" + self.telemetry_latest_id = 0 + self.telemetry_latest_index = 0 + self.live_scrub_active = False + self.hardware_meter_labels: dict[int, QLabel] = {} + self.training_cards: list[QWidget] = [] + self.training_controls_grid: Optional[QGridLayout] = None + self.training_controls_columns = 3 + self.training_health_points: list[tuple[int, Optional[float], Optional[float]]] = [] + self.active_training_log: Optional[QTextEdit] = None + self.active_training_progress: Optional[QProgressBar] = None + self.active_training_final_button_text = "Start Training" + self.active_training_output_dir: Optional[Path] = None + self.interrupt_count = 0 + self.chat_session: Optional[LlamaChatSession] = None + self.chat_markdown = "" + self.chat_stream_prefix = "" + self.chat_stream_reply = "" + self.current_assistant_browser: Optional[QTextBrowser] = None + self.current_assistant_meta: Optional[QLabel] = None + self.current_assistant_message: Optional[ChatMessageWidget] = None + self.pending_user_message = "" + self.spinner_index = 0 + self.spinner_timer = QTimer(self) + self.spinner_timer.timeout.connect(self._tick_spinner) + self.progress_timer = QTimer(self) + self.progress_timer.timeout.connect(self._drain_progress_queue) + self.job_manager = JobManager() + self.coordinator_server: Optional[CoordinatorApiServer] = None + self.coordinator_thread: Optional[Thread] = None + self.job_manager_timer = QTimer(self) + self.job_manager_timer.setInterval(2500) + self.job_manager_timer.timeout.connect(self.refresh_job_manager_tab) + + self._apply_style() + + shell = self._build_shell() + self.setCentralWidget(shell) + self._install_ui_event_logging(shell) + self._install_wheel_guard(shell) + self._refresh_notification_manager() + self.job_manager_timer.start() + + def eventFilter(self, watched: QObject, event: QEvent) -> bool: + """Prevent accidental wheel changes on compact option widgets. + + Args: + watched: Widget receiving the event. + event: Qt event. + + Returns: + True when the event is handled by the filter. + """ + + guarded_types = (QSpinBox, QDoubleSpinBox, QComboBox) + if isinstance(watched, guarded_types): + if event.type() == QEvent.Type.MouseButtonPress: + watched.setProperty("_wheel_enabled_after_click", True) + elif event.type() == QEvent.Type.FocusOut: + watched.setProperty("_wheel_enabled_after_click", False) + elif event.type() == QEvent.Type.Wheel and not watched.property("_wheel_enabled_after_click"): + return True + return super().eventFilter(watched, event) + + def _install_wheel_guard(self, root: QWidget) -> None: + """Require a click before spin boxes and combos react to mouse wheel. + + Args: + root: Root widget to scan for child controls. + """ + + for widget in root.findChildren(QWidget): + if not isinstance(widget, (QSpinBox, QDoubleSpinBox, QComboBox)): + continue + widget.setFocusPolicy(Qt.FocusPolicy.ClickFocus) + widget.setProperty("_wheel_enabled_after_click", False) + widget.installEventFilter(self) + + def _install_ui_event_logging(self, root: QWidget) -> None: + """Log user-facing widget actions and parameter changes. + + Args: + root: Root widget to scan for child controls. + """ + + for widget in root.findChildren(QWidget): + if isinstance(widget, QAbstractButton): + if widget.isCheckable(): + widget.toggled.connect( + lambda checked, item=widget: self._log_ui_event("toggled", item, checked) + ) + else: + widget.clicked.connect( + lambda checked=False, item=widget: self._log_ui_event("clicked", item, checked) + ) + elif isinstance(widget, QComboBox): + widget.currentTextChanged.connect( + lambda value, item=widget: self._log_ui_event("changed", item, value) + ) + elif isinstance(widget, QSpinBox): + widget.valueChanged.connect( + lambda value, item=widget: self._log_ui_event("changed", item, value) + ) + elif isinstance(widget, QDoubleSpinBox): + widget.valueChanged.connect( + lambda value, item=widget: self._log_ui_event("changed", item, value) + ) + elif isinstance(widget, QLineEdit): + widget.editingFinished.connect( + lambda item=widget: self._log_ui_event("edited", item, item.text()) + ) + + def _log_ui_event(self, action: str, widget: QWidget, value: Any) -> None: + """Log a UI action or parameter value. + + Args: + action: Event label. + widget: Widget that emitted the event. + value: Current value. + """ + + if action == "clicked" and isinstance(widget, QAbstractButton) and not widget.isCheckable(): + LOGGER.info("UI clicked: %s", self._widget_log_name(widget)) + return + LOGGER.info("UI %s: %s = %s", action, self._widget_log_name(widget), value) + + @staticmethod + def _widget_log_name(widget: QWidget) -> str: + """Return a useful log label for a widget. + + Args: + widget: Widget to describe. + + Returns: + Human-readable widget label. + """ + + if isinstance(widget, QAbstractButton) and widget.text(): + return widget.text().replace("\n", " ") + if isinstance(widget, QLineEdit) and widget.placeholderText(): + return widget.placeholderText() + if widget.objectName(): + return widget.objectName() + return widget.__class__.__name__ + + def _apply_style(self) -> None: + """Load the application stylesheet from the QSS module file.""" + + qss_path = Path(__file__).with_name("styles.qss") + self.setStyleSheet(qss_path.read_text(encoding="utf-8")) + + def _build_shell(self) -> QWidget: + """Build the top-level dashboard shell. + + Returns: + Root shell widget. + """ + + shell = QWidget() + shell.setObjectName("AppShell") + root = QVBoxLayout(shell) + root.setContentsMargins(8, 8, 8, 8) + root.setSpacing(0) + + top = QWidget() + top.setObjectName("TopBar") + self.top_bar = top + top_layout = QHBoxLayout(top) + top_layout.setContentsMargins(16, 8, 16, 8) + top_layout.setSpacing(8) + logo = QLabel() + logo.setObjectName("Logo") + logo_pixmap = self._app_logo_pixmap(36) + if logo_pixmap.isNull(): + logo.setText("DB") + else: + logo.setPixmap(logo_pixmap) + logo.setFixedSize(42, 42) + logo.setScaledContents(False) + self.search_box = QLineEdit() + self.search_box.setPlaceholderText("Project name...") + self.search_box.setMaximumWidth(260) + self._tip(self.search_box, f"Project name used when saving or reopening a {APP_NAME} project.") + self.new_project_button = QPushButton("New Project") + self.new_project_button.setMaximumWidth(130) + self.new_project_button.clicked.connect(self.new_project) + self._tip(self.new_project_button, f"Start a fresh {APP_NAME} project with default paths and settings.") + self.save_project_button = QPushButton("Save Project") + self.save_project_button.setMaximumWidth(130) + self.save_project_button.clicked.connect(self.save_project) + self._tip(self.save_project_button, "Save all current paths and settings into a project.json file.") + self.open_project_button = QPushButton("Open Project") + self.open_project_button.setMaximumWidth(130) + self.open_project_button.clicked.connect(self.open_project) + self._tip(self.open_project_button, "Open a saved project.json file and restore the UI settings.") + self.dataset_status = QLabel("Dataset: not prepared") + self.train_status = QLabel("Training: idle") + self.export_status = QLabel("Export: waiting") + self.chat_status = QLabel("Chat: no model loaded") + for label in (self.dataset_status, self.train_status, self.export_status, self.chat_status): + label.setObjectName("TopStatus") + label.setMinimumWidth(0) + label.setMaximumWidth(180) + label.setSizePolicy(QSizePolicy.Ignored, QSizePolicy.Fixed) + label.setWordWrap(False) + self.project_state = QLabel("Ready") + self.project_state.setObjectName("Metric") + self.project_state.setMinimumWidth(0) + self.project_state.setSizePolicy(QSizePolicy.Ignored, QSizePolicy.Fixed) + top_layout.addWidget(logo) + top_layout.addSpacing(12) + top_layout.addWidget(self.search_box) + top_layout.addWidget(self.new_project_button) + top_layout.addWidget(self.save_project_button) + top_layout.addWidget(self.open_project_button) + top_layout.addSpacing(10) + top_layout.addWidget(self.dataset_status) + top_layout.addWidget(self.train_status) + top_layout.addWidget(self.export_status) + top_layout.addWidget(self.chat_status) + top_layout.addStretch(1) + top_layout.addWidget(self.project_state) + root.addWidget(top) + + body = QHBoxLayout() + body.setContentsMargins(0, 0, 0, 0) + body.setSpacing(0) + rail = QWidget() + rail.setObjectName("SideRail") + self.side_rail = rail + rail.setFixedWidth(82) + rail_layout = QVBoxLayout(rail) + rail_layout.setContentsMargins(12, 18, 12, 18) + rail_layout.setSpacing(12) + self.dataset_plan_nav = self._nav_button("PLAN") + self.dataset_nav = self._nav_button("IN") + self.training_nav = self._nav_button("AI") + self.training_nav.setText("AI") + self.fine_tune_nav = self._nav_button("FT") + self.live_nav = self._nav_button("LIVE") + self.jobs_nav = self._nav_button("JOB") + self.benchmark_nav = self._nav_button("Bench") + self.export_nav = self._nav_button("X") + self.chat_nav = self._nav_button("Chat") + self._tip(self.dataset_plan_nav, "Open Dataset Blueprint: plan the target data mix before ingestion.") + self._tip(self.dataset_nav, "Open dataset preparation: load text/PDF files and build tokenizer data.") + self._tip(self.training_nav, "Open model training: configure architecture and optimization settings.") + self._tip(self.fine_tune_nav, "Open fine-tuning: adapt checkpoints with instruction, conversation, or LoRA settings.") + self._tip(self.live_nav, "Open the live training tracker with model flow, charts, metrics, and telemetry.") + self._tip(self.jobs_nav, "Open Job Manager: monitor workers, remote connections, assignments, and job controls.") + self._tip(self.benchmark_nav, "Open benchmark prompts: test checkpoint quality with repeatable prompts.") + self._tip(self.export_nav, "Open export tools: bundle or quantize the trained model artifacts.") + self._tip(self.chat_nav, "Open Chat: load a GGUF or native MicroGPT model once and send prompts.") + self.dataset_plan_nav.setChecked(True) + self.dataset_plan_nav.clicked.connect(lambda: self._switch_page(0)) + self.dataset_nav.clicked.connect(lambda: self._switch_page(1)) + self.training_nav.clicked.connect(lambda: self._switch_page(2)) + self.fine_tune_nav.clicked.connect(lambda: self._switch_page(3)) + self.live_nav.clicked.connect(lambda: self._switch_page(4)) + self.jobs_nav.clicked.connect(lambda: self._switch_page(5)) + self.benchmark_nav.clicked.connect(lambda: self._switch_page(6)) + self.export_nav.clicked.connect(lambda: self._switch_page(7)) + self.chat_nav.clicked.connect(lambda: self._switch_page(8)) + rail_layout.addWidget(self.dataset_plan_nav) + rail_layout.addWidget(self.dataset_nav) + rail_layout.addWidget(self.training_nav) + rail_layout.addWidget(self.fine_tune_nav) + rail_layout.addWidget(self.live_nav) + rail_layout.addWidget(self.jobs_nav) + rail_layout.addWidget(self.benchmark_nav) + rail_layout.addWidget(self.export_nav) + rail_layout.addWidget(self.chat_nav) + rail_layout.addStretch(1) + + self.pages = QStackedWidget() + self.pages.addWidget(self._build_dataset_plan_tab()) + self.pages.addWidget(self._build_dataset_tab()) + self.pages.addWidget(self._build_training_tab()) + self.pages.addWidget(self._build_fine_tuning_tab()) + self.pages.addWidget(self._build_live_training_tab()) + self.pages.addWidget(self._build_job_manager_tab()) + self.pages.addWidget(self._build_benchmark_tab()) + self.pages.addWidget(self._build_export_tab()) + self.pages.addWidget(self._build_chat_tab()) + + body.addWidget(rail) + body.addWidget(self.pages, 1) + root.addLayout(body, 1) + return shell + + def _nav_button(self, text: str) -> QPushButton: + """Create a left-rail navigation button. + + Args: + text: Button label. + + Returns: + Configured navigation button. + """ + + button = QPushButton(text) + button.setObjectName("NavButton") + button.setCheckable(True) + return button + + def _switch_page(self, index: int) -> None: + """Switch the visible page. + + Args: + index: Page index in the stacked widget. + """ + + self.pages.setCurrentIndex(index) + buttons = [ + self.dataset_plan_nav, + self.dataset_nav, + self.training_nav, + self.fine_tune_nav, + self.live_nav, + self.jobs_nav, + self.benchmark_nav, + self.export_nav, + self.chat_nav, + ] + for button_index, button in enumerate(buttons): + button.setChecked(button_index == index) + self._refresh_training_layout() + if index == 5: + QTimer.singleShot(20, self.refresh_job_manager_tab) + + def show_chat_only_mode(self) -> None: + """Collapse the UI to chat-only view for quick local LLM testing.""" + + if hasattr(self, "top_bar"): + self.top_bar.hide() + if hasattr(self, "side_rail"): + self.side_rail.hide() + self._switch_page(8) + self.setWindowTitle("DrunkenBot - Chat") + self.resize(980, 760) + + def resizeEvent(self, event: Any) -> None: + """Refresh responsive layouts when the main window changes size. + + Args: + event: Qt resize event. + """ + + super().resizeEvent(event) + self._refresh_training_layout() + + def _refresh_training_layout(self) -> None: + """Apply responsive card columns on the training page.""" + + if not self.training_cards or self.training_controls_grid is None: + return + width = self.pages.width() if hasattr(self, "pages") else self.width() + if width >= 900: + columns = 2 + else: + columns = 1 + if columns == self.training_controls_columns: + return + self._set_training_card_columns(columns) + + def _set_training_card_columns(self, columns: int) -> None: + """Reflow the training cards into the requested column count. + + Args: + columns: Number of columns to use. + """ + + if self.training_controls_grid is None: + return + while self.training_controls_grid.count(): + self.training_controls_grid.takeAt(0) + for index, card in enumerate(self.training_cards): + row = index // columns + column = index % columns + self.training_controls_grid.addWidget(card, row, column) + for column in range(2): + self.training_controls_grid.setColumnStretch(column, 1 if column < columns else 0) + self.training_controls_columns = columns + + def _build_dataset_plan_tab(self) -> QWidget: + """Build the dataset blueprint page. + + Returns: + Dataset blueprint page widget. + """ + + return build_dataset_plan_tab(self) + + def _build_dataset_tab(self) -> QWidget: + """Build the dataset preparation page. + + Returns: + Dataset page widget. + """ + + return build_dataset_tab(self) diff --git a/interface/main_window_part10.py b/interface/main_window_part10.py new file mode 100644 index 0000000..4063ac3 --- /dev/null +++ b/interface/main_window_part10.py @@ -0,0 +1,444 @@ +from __future__ import annotations + +# MainWindow implementation mixin. Runtime names are provided by app.py. +from typing import Any +from . import app as _app + +globals().update({name: value for name, value in vars(_app).items() if not name.startswith("__")}) + +class MainWindowPart10: + def _system_cpu_value(self) -> Optional[float]: + """Read system CPU utilization for live telemetry. + + Returns: + System CPU percentage, or None when unavailable. + """ + + if psutil is None: + return None + return float(psutil.cpu_percent(interval=None)) + + @staticmethod + def _format_duration(seconds: float) -> str: + """Format a duration for compact UI display. + + Args: + seconds: Duration in seconds. + + Returns: + Human-readable compact duration. + """ + + seconds = max(0, int(seconds)) + hours, remainder = divmod(seconds, 3600) + minutes, secs = divmod(remainder, 60) + if hours: + return f"{hours}h {minutes:02d}m" + if minutes: + return f"{minutes}m {secs:02d}s" + return f"{secs}s" + + def _apply_chat_delta(self, event: dict[str, Any]) -> None: + """Apply one streamed chat chunk to the rendered conversation. + + Args: + event: Chat stream progress event. + """ + + self.chat_stream_reply += str(event.get("content", "")) + should_follow = self._is_chat_near_bottom() + self._render_chat_markdown(self.chat_stream_reply) + if should_follow: + self.chat_scroll.verticalScrollBar().setValue(self.chat_scroll.verticalScrollBar().maximum()) + self._set_chat_stats( + float(event.get("elapsed_seconds", 0.0)), + int(event.get("token_count", 0)), + float(event.get("tokens_per_second", 0.0)), + ) + + def _drain_progress_queue(self) -> None: + """Drain queued worker progress events on the UI thread.""" + + if self.progress_queue is None or self.active_log is None or self.active_progress_bar is None: + return + drained = 0 + last_percent = None + while drained < 12: + try: + event = self.progress_queue.get_nowait() + except Empty: + break + notification_event = event + if isinstance(event, dict) and event.get("percent") is not None: + last_percent = event.get("percent") + event = {**event, "percent": None} + if ( + isinstance(notification_event, dict) + and self.active_task_kind == "dataset_download" + and notification_event.get("button_text") + and self.active_button is not None + ): + self.active_button_text = str(notification_event["button_text"]) + self.active_button.setText(self.active_button_text) + self._handle_progress(event, self.active_log, self.active_progress_bar) + if isinstance(notification_event, dict): + self._notify_progress(notification_event) + drained += 1 + if last_percent is not None: + self.active_progress_bar.setValue(max(0, min(100, int(last_percent)))) + + def _thread_finished(self) -> None: + """Clean up thread bookkeeping after a worker finishes.""" + + LOGGER.info("Background task thread finished") + self._drain_progress_queue() + if self.progress_timer.isActive(): + self.progress_timer.stop() + self.thread = None + self.worker = None + if self.result_bridge is not None: + self.result_bridge.deleteLater() + self.result_bridge = None + self.stop_event = None + self.progress_queue = None + self.active_log = None + self.active_progress_bar = None + if self.active_stop_button is not None: + self.active_stop_button.setEnabled(False) + self.active_stop_button = None + if self.active_button is not None: + self._clear_button_busy() + self.active_task_kind = "" + + def _task_failed(self, message: str, log: QTextEdit, progress_bar: QProgressBar) -> None: + """Handle background task failure. + + Args: + message: Error message. + log: Log widget to append to. + progress_bar: Progress bar to reset. + """ + + stopped_by_user = "stopped by user" in message.lower() + if stopped_by_user: + LOGGER.info("Background task stopped by user: %s", message) + else: + LOGGER.error("Background task error: %s", message) + log.append(f"Stopped: {message}" if stopped_by_user else f"Error: {message}") + self._notify_failure("Task stopped" if stopped_by_user else "Task failed", message) + progress_bar.setRange(0, 100) + progress_bar.setValue(0) + if stopped_by_user: + self.project_state.setText("Stopped") + self._clear_button_busy() + + def _set_button_busy(self, button: QPushButton, text: str) -> None: + """Disable a button and start its spinner text. + + Args: + button: Button to mark busy. + text: Busy label. + """ + + self.active_button = button + self.active_button_text = text + self.active_button_restore_text = button.text() + self.spinner_index = 0 + button.setEnabled(False) + button.setText(f"| {text}") + self.spinner_timer.start(150) + + def _clear_button_busy(self, final_text: Optional[str] = None) -> None: + """Restore the active busy button. + + Args: + final_text: Optional final button text. + """ + + if self.spinner_timer.isActive(): + self.spinner_timer.stop() + if self.active_button: + self.active_button.setEnabled(True) + self.active_button.setText(final_text or self.active_button_restore_text) + if self.active_stop_button: + self.active_stop_button.setEnabled(False) + self.active_button = None + self.active_button_text = "" + self.active_button_restore_text = "" + + def _tick_spinner(self) -> None: + """Advance the active button spinner frame.""" + + if not self.active_button: + return + frames = "|/-\\" + self.spinner_index = (self.spinner_index + 1) % len(frames) + self.active_button.setText(f"{frames[self.spinner_index]} {self.active_button_text}") + + def _dataset_config_from_ui(self) -> DatasetConfig: + """Collect dataset options from the current UI controls. + + Returns: + Dataset preparation configuration. + """ + + conversation_paths: list[Path] = [] + instruction_paths: list[Path] = [] + dataset_stage = self._dataset_stage_value() + return DatasetConfig( + input_dir=Path(self.input_dir.text()), + output_dir=Path(self.dataset_dir.text()), + vocab_size=None if self.auto_vocab.isChecked() else self.manual_vocab_size.value(), + conversation_datasets=self._selected_conversation_datasets(), + conversation_sample_limit=self.conversation_sample_limit.value(), + conversation_dataset_path=conversation_paths[0] if conversation_paths else None, + instruction_dataset_path=instruction_paths[0] if instruction_paths else None, + conversation_dataset_paths=conversation_paths, + instruction_dataset_paths=instruction_paths, + default_data_paths=self._selected_default_data_paths_for_stage(dataset_stage), + mixture_weights=self._mixture_weights_from_ui(), + min_frequency=self.min_frequency.value(), + context_length=self.context_length.value(), + validation_split=self.validation_split.value(), + lowercase=False, + max_workers=self.max_workers.value(), + code_training_mode=self.code_training_mode.isChecked(), + include_prose=self.include_prose.isChecked(), + include_source_code=self.include_source_code.isChecked(), + extract_code_blocks=self.extract_code_blocks.isChecked(), + preserve_indentation=self.preserve_indentation.isChecked(), + generate_instruction_samples=self.instruction_samples.isChecked(), + reasoning_sample_mode=self._reasoning_sample_mode_value(), + prepare_mode=self._prepare_mode_value(), + tokenizer_strategy=self._tokenizer_strategy_value(), + tokenizer_path=Path(self.tokenizer_path.text()) if self.tokenizer_path.text().strip() else None, + dataset_stage=dataset_stage, + tokenizer_training_max_gb=self.tokenizer_training_max_gb.value(), + ) + + def _selected_default_data_paths_for_stage(self, stage: str) -> list[Path]: + """Return selected bundled files that match the dataset purpose. + + Args: + stage: Dataset preparation stage. + + Returns: + Selected paths suitable for the requested stage. + """ + + # Folder selection is the workflow configuration. Do not apply a + # second hardcoded stage filter here; the Dataset Sources tree already + # contains exactly the files selected by the user. + return self._selected_default_data_paths() + + @staticmethod + def _split_path_list(text: str) -> list[Path]: + """Split a semicolon-delimited path field. + + Args: + text: Raw path field text. + + Returns: + Parsed paths. + """ + + return [Path(item.strip().strip('"')) for item in text.split(";") if item.strip()] + + def check_project_health(self) -> None: + """Run a project health check in the background.""" + + self.dataset_log.clear() + self.dataset_progress.setValue(0) + self.dataset_log.append("Checking project health...") + self.project_state.setText("Checking health") + self._run_task( + check_project_health, + ( + Path(self.input_dir.text()), + Path(self.dataset_dir.text()), + Path(self.model_dir.text()), + Path(self.export_dir.text()), + Path(self.gguf_path.text()) if self.gguf_path.text().strip() else None, + Path(self.llama_cpp_dir.text()) if self.llama_cpp_dir.text().strip() else None, + self.device.currentText(), + ), + self._health_check_finished, + self.dataset_log, + self.dataset_progress, + with_progress=True, + button=self.health_check_button, + stop_button=self.stop_dataset_button, + busy_text="Checking Health", + isolate_process=True, + ) + + @Slot(object) + def _health_check_finished(self, result: Any) -> None: + """Display project health check results. + + Args: + result: Project health result. + """ + + self.dataset_progress.setValue(100) + self.dataset_log.append("") + self.dataset_log.append(f"Project health: {result.status.upper()} ({result.summary})") + for check in result.checks: + marker = {"ok": "OK", "warning": "WARN", "error": "ERROR"}.get(check.get("status"), "INFO") + self.dataset_log.append(f"[{marker}] {check.get('name')}: {check.get('detail')}") + self.project_state.setText("Health checked") + self._clear_button_busy("Check Health") + + def preview_dataset(self) -> None: + """Run a dataset preview and quality scan in the background.""" + + self.dataset_log.clear() + self.dataset_progress.setValue(0) + self.dataset_log.append("Previewing dataset...") + self.project_state.setText("Previewing dataset") + self._run_task( + scan_dataset_preview, + (self._dataset_config_from_ui(),), + self._dataset_preview_finished, + self.dataset_log, + self.dataset_progress, + with_progress=True, + button=self.preview_dataset_button, + stop_button=self.stop_dataset_button, + busy_text="Previewing Dataset", + isolate_process=True, + ) + + @Slot(object) + def _dataset_preview_finished(self, result: Any) -> None: + """Finish preview UI cleanup even when rendering a result fails.""" + try: + self._render_dataset_preview_result(result) + finally: + # Keep the action available after both successful and malformed + # worker results; otherwise the spinner can leave it disabled. + self._clear_button_busy("Preview Dataset") + + def _render_dataset_preview_result(self, result: Any) -> None: + """Display dataset preview and quality scan results. + + Args: + result: Dataset preview result. + """ + + self.dataset_progress.setValue(100) + suffix_text = ", ".join(f"{suffix}: {count}" for suffix, count in + result.suffix_counts.items()) or "none" + self.dataset_log.append("") + self.dataset_log.append( + f"Source files: {result.source_file_count:,}; size: {result.total_bytes / (1024 * 1024):.2f} MB") + self.dataset_log.append(f"File types: {suffix_text}") + self.dataset_log.append( + f"Prepared dataset artifacts: {'found' if result.prepared else 'not complete'}") + self.dataset_log.append( + f"Duplicate scan: {result.duplicate_count:,} file entries in {len(result.duplicate_groups):,} likely group(s).") + self.dataset_log.append( + f"Bad extraction scan: {result.bad_extraction_count:,} suspicious file(s).") + self.dataset_log.append( + f"Code/prose balance: {result.balance_label} ({result.code_preview_count:,}/{result.prose_preview_count:,}).") + self.dataset_log.append( + f"Training readiness: {result.readiness_label} ({result.readiness_score}/100).") + for reason in result.readiness_reasons[:8]: + self.dataset_log.append(f"- {reason}") + self.dataset_quality_duplicates.setText( + f"Duplicates: {result.duplicate_count:,}") + self.dataset_quality_extraction.setText( + f"Extraction: {result.bad_extraction_count:,} flagged") + self.dataset_quality_balance.setText( + f"Balance: {result.balance_label}") + self.dataset_quality_readiness.setText( + f"Readiness: {result.readiness_label} {result.readiness_score}/100") + if result.summary: + self._update_dataset_quality_report(result.summary) + # dataset_quality_duplicates is intentionally left alone here: + # _update_dataset_quality_report() just set it to the block-level + # duplication percentage from the prepared corpus (the more useful, + # actionable metric). Re-setting it to result.duplicate_count (a + # raw duplicate *file* count from the earlier preview scan) would + # silently discard that and always show the old metric instead. + self.dataset_quality_extraction.setText( + f"Extraction: {result.bad_extraction_count:,} flagged") + self.dataset_quality_balance.setText( + f"Balance: {result.balance_label}") + self.dataset_quality_readiness.setText( + f"Readiness: {result.readiness_label} {result.readiness_score}/100") + tokens = int(result.summary.get("token_count", 0) or 0) + vocab = int(result.summary.get("tokenizer_vocab_size", 0) or 0) + self.dataset_log.append( + f"Prepared summary: {tokens:,} tokens, vocab {vocab:,}.") + else: + self.dataset_quality_samples.setText( + f"Preview: {len(result.sample_previews):,} shown") + self.dataset_quality_tokens.setText("Tokens: not prepared") + self.dataset_quality_windows.setText("Windows: not prepared") + self.dataset_quality_vocab.setText("Vocab: not prepared") + self.dataset_quality_code.setText( + f"Code/prose: {result.code_preview_count:,}/{result.prose_preview_count:,}") + self.dataset_quality_cache.setText( + f"Files: {result.source_file_count:,} source") + if result.duplicate_groups: + self.dataset_log.append("") + self.dataset_log.append("Likely duplicates:") + for group in result.duplicate_groups[:8]: + self.dataset_log.append( + f"- {group.get('type')}: {group.get('count')} file(s)") + for path in group.get("files", [])[:4]: + self.dataset_log.append(f" {Path(path).name}") + if result.bad_extraction_files: + self.dataset_log.append("") + self.dataset_log.append("Suspicious extraction files:") + for item in result.bad_extraction_files[:12]: + self.dataset_log.append( + f"- {Path(item.get('path', '')).name}: {item.get('reasons')}") + suggestions: list[str] = [] + if result.duplicate_groups: + suggestions.append( + "Remove or move duplicate files before preparing the final dataset.") + if result.bad_extraction_files: + suggestions.append( + "Replace flagged PDFs with text/source versions, or remove files with bad extraction.") + if result.balance_label == "Prose heavy" and self.code_training_mode.isChecked(): + suggestions.append( + "Add real source-code folders or enable source-file inclusion for a stronger coding model.") + if result.balance_label == "Code heavy": + suggestions.append( + "Add README/tutorial/prose explanations if you want the model to explain code well.") + if result.readiness_label in {"Needs cleanup", "Not ready"}: + suggestions.append( + "Run Preview Dataset again after cleanup and only train once readiness improves.") + if hasattr(self, "dataset_advisor"): + if suggestions: + self.dataset_advisor.setPlainText( + "\n".join(f"- {suggestion}" for suggestion in suggestions)) + else: + self.dataset_advisor.setPlainText( + "No immediate cleanup suggestions. Dataset looks acceptable for the current preview.") + if suggestions: + self.dataset_log.append("") + self.dataset_log.append("Cleanup suggestions:") + for suggestion in suggestions: + self.dataset_log.append(f"- {suggestion}") + if result.issues: + self.dataset_quality_warning.setText( + f"Warnings: {len(result.issues)}") + self.dataset_log.append("") + self.dataset_log.append("Quality notes:") + for issue in result.issues[:12]: + self.dataset_log.append(f"- {issue}") + else: + self.dataset_quality_warning.setText("Warnings: none") + if result.sample_previews: + self.dataset_log.append("") + self.dataset_log.append("Preview samples:") + for index, sample in enumerate(result.sample_previews, start=1): + label = sample.get("language") or sample.get("kind") or "text" + self.dataset_log.append( + f"\n[{index}] {Path(sample.get('path', '')).name} ({label}, {sample.get('characters')} chars)") + self.dataset_log.append( + sample.get("preview", "").replace("\n", "\n ")[:1400]) + self.project_state.setText("Dataset previewed") diff --git a/interface/main_window_part11.py b/interface/main_window_part11.py new file mode 100644 index 0000000..5f09168 --- /dev/null +++ b/interface/main_window_part11.py @@ -0,0 +1,426 @@ +from __future__ import annotations + +# MainWindow implementation mixin. Runtime names are provided by app.py. +from typing import Any +from . import app as _app + +globals().update({name: value for name, value in vars(_app).items() if not name.startswith("__")}) + +class MainWindowPart11: + def prepare_dataset(self) -> None: + """Collect dataset options and start dataset preparation.""" + + config = self._dataset_config_from_ui() + self.dataset_log.clear() + self.dataset_progress.setValue(0) + self._reset_dataset_quality_report() + self.dataset_log.append("Preparing dataset...") + self.dataset_log.append(f"App log file: {self.log_file_path}") + self.dataset_log.append(f"Dataset purpose: {dataset_stage_label(config.dataset_stage)}") + if config.conversation_dataset_paths: + self.dataset_log.append(f"Local conversation JSON/JSONL: {len(config.conversation_dataset_paths)} path(s)") + LOGGER.info("Local conversation JSON/JSONL datasets: %s", "; ".join(str(path) for path in config.conversation_dataset_paths)) + if config.instruction_dataset_paths: + self.dataset_log.append(f"Local instruction JSON/JSONL: {len(config.instruction_dataset_paths)} path(s)") + LOGGER.info("Local instruction JSON/JSONL datasets: %s", "; ".join(str(path) for path in config.instruction_dataset_paths)) + if config.default_data_paths: + self.dataset_log.append(f"Bundled default data: {len(config.default_data_paths)} file(s)") + LOGGER.info("Bundled default data files: %s", "; ".join(str(path) for path in config.default_data_paths)) + if self.include_conversation_datasets.isChecked(): + selected_labels = [ + action.text() + for action in getattr(self, "conversation_dataset_actions", {}).values() + if action.isChecked() and action.isVisible() + ] + if selected_labels: + hf_cache = config.output_dir / "cache" / "huggingface" + self.dataset_log.append(f"Online training datasets: {', '.join(selected_labels)}") + self.dataset_log.append(f"Downloading/loading online data at: {hf_cache}") + LOGGER.info("Online training datasets: %s", ", ".join(selected_labels)) + LOGGER.info("Downloading/loading online data at: %s", hf_cache) + else: + self.dataset_log.append("Online training datasets are enabled, but no dataset is selected for this purpose.") + LOGGER.warning("Online training datasets enabled, but no dataset is selected") + else: + self.dataset_log.append("Online training datasets: off. Local source files only.") + LOGGER.info("Online training datasets: off. Local source files only.") + checked_count = sum( + 1 + for action in getattr(self, "conversation_dataset_actions", {}).values() + if action.isChecked() + ) + if checked_count: + self.dataset_log.append("Checked online dataset choices are ignored until the master checkbox is enabled.") + LOGGER.info("Checked online dataset choices are ignored until the master checkbox is enabled") + LOGGER.info( + "Preparing dataset: input=%s output=%s stage=%s online_datasets=%s conversation_json=%s instruction_json=%s", + config.input_dir, + config.output_dir, + config.dataset_stage, + ",".join(config.conversation_datasets) or "off", + ";".join(str(path) for path in config.conversation_dataset_paths) or "off", + ";".join(str(path) for path in config.instruction_dataset_paths) or "off", + ) + self.project_state.setText("Preparing dataset") + self.dataset_status.setText("Dataset: preparing") + self.auto_vocab_label.setText("Calculating...") + self._run_task( + build_dataset, + (config,), + self._dataset_finished, + self.dataset_log, + self.dataset_progress, + with_progress=True, + button=self.prepare_button, + stop_button=self.stop_dataset_button, + busy_text="Preparing Dataset", + task_kind="dataset", + isolate_process=True, + ) + + @Slot(object) + def _dataset_finished(self, result: Any) -> None: + """Update UI after dataset preparation finishes. + + Args: + result: Dataset build result. + """ + + self.dataset_progress.setValue(100) + self.auto_vocab_label.setText(f"{result.vocab_size:,}") + + LOGGER.info( + "Dataset prepared: documents=%s tokens=%s vocab=%s code=%s prose=%s conversation=%s output=%s", + result.document_count, + result.token_count, + result.vocab_size, + result.code_sample_count, + result.prose_sample_count, + getattr(result, "conversation_sample_count", 0), + result.output_dir, + ) + + self.dataset_log.append( + f"Prepared {result.document_count} documents, " + f"{result.character_count:,} characters, " + f"{result.token_count:,} tokens, " + f"vocab {result.vocab_size:,}." + ) + + if getattr(result, "train_window_count", 0) or getattr(result, + "val_window_count", + 0): + self.dataset_log.append( + f"Training windows: {result.train_window_count:,}; " + f"validation windows: {result.val_window_count:,}." + ) + + self.dataset_log.append( + f"Cache summary: reused {result.cached_file_count:,} file(s), " + f"processed {result.processed_file_count:,} file(s)." + ) + + if getattr(result, "dataset_version_id", ""): + self.dataset_log.append( + f"Dataset version: {result.dataset_version_id}" + ) + + if result.warning: + self.dataset_log.append(f"Recommendation: {result.warning}") + + self._update_dataset_quality_report( + { + "document_count": result.document_count, + "token_count": result.token_count, + "train_window_count": getattr(result, "train_window_count", 0), + "val_window_count": getattr(result, "val_window_count", 0), + "character_count": result.character_count, + "tokenizer_vocab_size": result.vocab_size, + "code_sample_count": result.code_sample_count, + "prose_sample_count": result.prose_sample_count, + "conversation_sample_count": getattr(result, + "conversation_sample_count", + 0), + "cached_file_count": result.cached_file_count, + "processed_file_count": result.processed_file_count, + "skipped_file_count": result.skipped_file_count, + "failed_file_count": result.failed_file_count, + "warning": result.warning, + "sequence_token_stats": getattr(result, "sequence_token_stats", + {}), + "duplicate_block_count": getattr(result, + "duplicate_block_count", 0), + "unique_block_count": getattr(result, "unique_block_count", 0), + "corpus_block_count": getattr(result, "corpus_block_count", 0), + "duplicate_block_ratio": getattr(result, + "duplicate_block_ratio", 0.0), + "unique_block_ratio": getattr(result, "unique_block_ratio", + 1.0), + } + ) + + self.train_data_dir.setText(str(result.output_dir)) + self.project_state.setText("Dataset ready") + + self.dataset_status.setText( + f"Dataset: {result.document_count} files, {result.token_count:,} tokens" + ) + + if result.code_sample_count: + self.dataset_status.setText( + f"Dataset: {result.code_sample_count:,} code, " + f"{result.prose_sample_count:,} prose, " + f"{result.token_count:,} tokens" + ) + + self.refresh_model_estimate() + self.refresh_fine_tune_workflow() + + self._notify_complete( + "dataset", + "Dataset preparation complete", + [ + f"Output: {result.output_dir}", + f"Documents: {result.document_count:,}", + f"Characters: {result.character_count:,}", + f"Tokens: {result.token_count:,}", + f"Vocabulary: {result.vocab_size:,}", + ( + "Windows: " + f"{getattr(result, 'train_window_count', 0):,} training, " + f"{getattr(result, 'val_window_count', 0):,} validation" + ), + ( + "Content mix: " + f"{result.code_sample_count:,} code, " + f"{result.prose_sample_count:,} prose, " + f"{getattr(result, 'conversation_sample_count', 0):,} conversation" + ), + ( + "Files: " + f"{result.processed_file_count:,} processed, " + f"{result.cached_file_count:,} cached, " + f"{result.skipped_file_count:,} skipped, " + f"{result.failed_file_count:,} failed" + ), + f"Dataset version: {getattr(result, 'dataset_version_id', '') or '-'}", + f"Health: {'warning - ' + result.warning if result.warning else 'ready'}", + ], + ) + + self._clear_button_busy("DataSet Prepared") + + def _prepare_mode_value(self) -> str: + """Return the selected dataset preparation mode. + + Returns: + Internal mode value. + """ + + label = self.prepare_mode.currentText() + if label == "Full rebuild": + return "full_rebuild" + if label == "Force reprocess": + return "force_reprocess" + return "incremental" + + def _tokenizer_strategy_value(self) -> str: + """Return the selected tokenizer strategy. + + Returns: + Internal tokenizer strategy value. + """ + + label = self.tokenizer_strategy.currentText() + if label == "Train new tokenizer": + return "train_new" + if label == "Reuse dataset tokenizer": + return "reuse_dataset" + if label == "Import tokenizer.json": + return "import_tokenizer" + return "auto" + + def _reasoning_sample_mode_value(self) -> str: + """Return the selected reasoning sample mode. + + Returns: + Internal reasoning sample mode. + """ + + label = self.reasoning_sample_mode.currentText() + if label == "Detailed code reasoning": + return "detailed" + if label == "No reasoning wrapper": + return "none" + return "scaffold" + + def _dataset_stage_value(self) -> str: + """Return the selected dataset preparation stage. + + Returns: + Dataset stage identifier. + """ + + return self.dataset_stage.currentText().strip().lower().replace(" ", "_") or "base" + + def _set_dataset_stage(self, stage: str) -> None: + """Set the dataset stage combo from an internal stage value. + + Args: + stage: Dataset stage identifier. + """ + + index = self.dataset_stage.findText(stage, Qt.MatchFixedString) + if index < 0: + self.dataset_stage.addItem(stage) + index = self.dataset_stage.count() - 1 + self.dataset_stage.setCurrentIndex(index) + self._update_online_dataset_stage_controls() + + def _update_online_dataset_stage_controls(self) -> None: + """Show and enable online datasets for the selected training stage.""" + + if not hasattr(self, "dataset_stage"): + return + stage = self._dataset_stage_value() + allowed = set(CONVERSATION_DATASET_PRESETS) + include_online = self.include_conversation_datasets.isChecked() + for dataset_id, action in getattr(self, "conversation_dataset_actions", {}).items(): + visible = dataset_id in allowed + action.setVisible(visible) + action.setEnabled(include_online and visible) + if not visible: + action.setChecked(False) + if hasattr(self, "conversation_dataset_button"): + self.conversation_dataset_button.setEnabled(include_online) + self.conversation_sample_limit.setEnabled(include_online) + self._update_conversation_dataset_button_text() + self.conversation_datasets_status.setText( + f"{self.dataset_stage.currentText()}: choose optional online datasets." + if include_online else "Choose optional online datasets, or use local folders only." + ) + + def _selected_conversation_datasets(self) -> list[str]: + """Return selected built-in conversation dataset IDs. + + Returns: + Selected dataset identifiers. + """ + + allowed = set(CONVERSATION_DATASET_PRESETS) + selected = [ + dataset_id + for dataset_id, action in getattr(self, "conversation_dataset_actions", {}).items() + if dataset_id in allowed and action.isChecked() and self.include_conversation_datasets.isChecked() + ] + custom = self.custom_huggingface_dataset.text().strip() + if custom and self.include_conversation_datasets.isChecked(): + selected.append(f"hf_custom:{custom}") + return selected + + def _download_custom_huggingface_dataset(self) -> None: + """Enable the entered Hugging Face dataset for the next preparation run.""" + value = self.custom_huggingface_dataset.text().strip() + if not value: + self.conversation_datasets_status.setText("Enter a Hugging Face dataset ID or URL first.") + return + self.include_conversation_datasets.setChecked(True) + self.conversation_datasets_status.setText( + f"Custom dataset queued: {value}. It will download during dataset preparation." + ) + self._update_conversation_dataset_button_text() + + def _set_selected_conversation_datasets(self, dataset_ids: list[str]) -> None: + """Restore selected conversation dataset actions. + + Args: + dataset_ids: Dataset IDs to select. + """ + + selected = set(dataset_ids) + allowed = set(dataset_ids_for_stage(self._dataset_stage_value())) + for dataset_id, action in getattr(self, "conversation_dataset_actions", {}).items(): + action.setChecked(dataset_id in selected and dataset_id in allowed) + action.setEnabled(self.include_conversation_datasets.isChecked() and dataset_id in allowed) + if hasattr(self, "custom_huggingface_dataset"): + self.custom_huggingface_dataset.setText( + next((value[10:] for value in dataset_ids if value.startswith("hf_custom:")), "") + ) + if hasattr(self, "conversation_sample_limit"): + self.conversation_sample_limit.setEnabled(self.include_conversation_datasets.isChecked()) + self._update_conversation_dataset_button_text() + if hasattr(self, "conversation_datasets_status"): + self._update_online_dataset_stage_controls() + + def _update_conversation_dataset_button_text(self) -> None: + """Refresh the compact online dataset selector label.""" + + if not hasattr(self, "conversation_dataset_button"): + return + allowed = set(dataset_ids_for_stage(self._dataset_stage_value())) if hasattr(self, "dataset_stage") else set() + selected_labels = [ + action.text() + for dataset_id, action in getattr(self, "conversation_dataset_actions", {}).items() + if dataset_id in allowed and action.isChecked() + ] + if not self.include_conversation_datasets.isChecked(): + self.conversation_dataset_button.setText("Online datasets off") + elif not selected_labels: + self.conversation_dataset_button.setText("Choose online datasets") + elif len(selected_labels) == 1: + self.conversation_dataset_button.setText(selected_labels[0]) + else: + self.conversation_dataset_button.setText(f"{len(selected_labels)} online datasets selected") + + def configure_fine_tune_dataset_builder(self) -> None: + """Configure the Ingest tab for the selected fine-tune dataset type.""" + + stage_label = self.fine_tune_dataset_builder_stage.currentText() + stage = { + "Instruction fine-tune": "instruction", + "Conversation fine-tune": "conversation", + "Code fine-tune": "code", + }.get(stage_label, "instruction") + starter_datasets = { + "instruction": ["alpaca_52k"], + "conversation": ["dailydialog"], + "code": ["codealpaca_20k"], + } + self._set_dataset_stage(stage) + self.include_conversation_datasets.setChecked(True) + self._set_selected_conversation_datasets(starter_datasets.get(stage, [])) + if stage == "code": + self.code_training_mode.setChecked(True) + self.include_source_code.setChecked(True) + self.extract_code_blocks.setChecked(True) + self.preserve_indentation.setChecked(True) + self._set_mixture_weights({}) + self._switch_page(0) + self.dataset_log.append(f"Configured Ingest for {dataset_stage_label(stage)}. Import the base tokenizer before preparing.") + self.project_state.setText(f"Configured {dataset_stage_label(stage)} data") + + def _dataset_plan_from_ui(self) -> dict[str, float]: + """Return dataset blueprint state. + + Returns: + Empty mapping because category percentages are disabled. + """ + + return {} + + def _selected_default_data_paths(self) -> list[Path]: + """Return bundled default data files selected in the Dataset Blueprint. + + Returns: + Selected bundled data paths. + """ + + if not hasattr(self, "default_data_actions"): + return [path for path, _category in iter_default_data_files()] + return [ + Path(path) + for path, item in self.default_data_actions.items() + if item.checkState(0) == Qt.Checked + ] + + diff --git a/interface/main_window_part12.py b/interface/main_window_part12.py new file mode 100644 index 0000000..b0ee5f4 --- /dev/null +++ b/interface/main_window_part12.py @@ -0,0 +1,438 @@ +from __future__ import annotations + +# MainWindow implementation mixin. Runtime names are provided by app.py. +from typing import Any +from . import app as _app + +globals().update({name: value for name, value in vars(_app).items() if not name.startswith("__")}) + +class MainWindowPart12: + def _set_selected_default_data_paths(self, paths: Optional[list[Any]]) -> None: + """Restore bundled default data checkbox selections. + + Args: + paths: Saved bundled data file paths. ``None`` means no + preference was ever saved (a brand-new project), and every + file is selected by default. An explicit empty list means + the user deliberately deselected everything, and that + choice is restored as-is rather than falling back to + "select everything" -- previously the two cases were + indistinguishable, so saving a project with nothing + selected silently reset to everything selected on reload. + """ + + if not hasattr(self, "default_data_actions"): + return + if paths is None: + selected = set(self.default_data_actions) + else: + selected = {str(Path(path)) for path in paths} + self.default_data_tree_updating = True + try: + for path, item in self.default_data_actions.items(): + item.setCheckState(0, Qt.Checked if path in selected else Qt.Unchecked) + self._refresh_default_data_category_states() + finally: + self.default_data_tree_updating = False + + def _set_dataset_blueprint_refresh_busy(self, busy: bool) -> None: + """Toggle refresh busy state indicators for the Dataset Sources page.""" + + if hasattr(self, "dataset_plan_refresh_button"): + self.dataset_plan_refresh_button.setEnabled(not busy) + self.dataset_plan_refresh_button.setText("Refreshing..." if busy else "Refresh") + if hasattr(self, "dataset_plan_progress"): + if busy: + self.dataset_plan_progress.setRange(0, 0) + self.dataset_plan_progress.setVisible(True) + else: + self.dataset_plan_progress.setRange(0, 100) + self.dataset_plan_progress.setValue(0) + self.dataset_plan_progress.setVisible(False) + + def refresh_dataset_blueprint_files(self) -> None: + """Reload the Dataset Blueprint file tree from disk.""" + + root = getattr(self, "blueprint_data_root", default_data_root()) + self._refresh_external_dataset_status() + selected_paths = [str(path) for path in self._selected_default_data_paths()] + self._set_dataset_blueprint_refresh_busy(True) + QApplication.processEvents() + try: + self._refresh_dataset_blueprint_source( + Path(root), + saved_paths=selected_paths, + saved_plan=self._dataset_plan_from_ui(), + preset="Custom", + ) + self.project_state.setText("Blueprint refreshed") + LOGGER.info("Dataset blueprint tree refreshed from %s", root) + finally: + self._set_dataset_blueprint_refresh_busy(False) + + def _handle_default_data_tree_changed(self, item: Any, column: int) -> None: + """Handle category and file toggles in the bundled data tree. + + Args: + item: Changed tree item. + column: Changed column index. + """ + + if column != 0 or getattr(self, "default_data_tree_updating", False): + return + data = item.data(0, Qt.UserRole) or {} + if data.get("kind") != "category": + self.default_data_tree_updating = True + try: + self._refresh_default_data_category_states() + finally: + self.default_data_tree_updating = False + if hasattr(self, "_mixture_weights_state"): + delattr(self, "_mixture_weights_state") + return + state = item.checkState(0) + if state == Qt.PartiallyChecked: + return + self.default_data_tree_updating = True + try: + for index in range(item.childCount()): + item.child(index).setCheckState(0, state) + finally: + self.default_data_tree_updating = False + if hasattr(self, "_mixture_weights_state"): + delattr(self, "_mixture_weights_state") + + def _refresh_default_data_category_states(self) -> None: + """Refresh category checkbox states from child file selections.""" + + if not hasattr(self, "default_data_category_items"): + return + for category_item in self.default_data_category_items.values(): + checked = 0 + partial = False + for index in range(category_item.childCount()): + state = category_item.child(index).checkState(0) + if state == Qt.Checked: + checked += 1 + elif state == Qt.PartiallyChecked: + partial = True + if partial or 0 < checked < category_item.childCount(): + category_item.setCheckState(0, Qt.PartiallyChecked) + elif checked == category_item.childCount() and category_item.childCount() > 0: + category_item.setCheckState(0, Qt.Checked) + else: + category_item.setCheckState(0, Qt.Unchecked) + + def _set_dataset_plan(self, plan: dict[str, Any], preset: str = "Custom") -> None: + """Restore high-level dataset blueprint controls. + + Args: + plan: Saved dataset domain percentages. + preset: Saved preset label. + """ + + if not hasattr(self, "dataset_plan_spins"): + return + self._restoring_dataset_plan = True + try: + values = {**dataset_plan_defaults(), **(plan or {})} + for key, widget in self.dataset_plan_spins.items(): + widget.blockSignals(True) + try: + widget.setValue(float(values.get(key, 0.0))) + except (TypeError, ValueError): + widget.setValue(0.0) + widget.blockSignals(False) + self.dataset_plan_preset.blockSignals(True) + if preset == "Custom": + self.dataset_plan_preset.setCurrentText(preset) + else: + self.dataset_plan_preset.setCurrentText("Custom") + self.dataset_plan_preset.blockSignals(False) + finally: + self._restoring_dataset_plan = False + self._update_dataset_plan_total() + + def _dataset_plan_mark_custom(self, *_args: Any) -> None: + """Mark the dataset blueprint as custom after manual edits.""" + + if getattr(self, "_restoring_dataset_plan", False): + return + if hasattr(self, "_mixture_weights_state"): + delattr(self, "_mixture_weights_state") + if hasattr(self, "dataset_plan_preset") and self.dataset_plan_preset.currentText() != "Custom": + self.dataset_plan_preset.blockSignals(True) + self.dataset_plan_preset.setCurrentText("Custom") + self.dataset_plan_preset.blockSignals(False) + + def _update_dataset_plan_total(self) -> None: + """No-op retained for compatibility after blueprint percentage removal.""" + + return + + def normalize_dataset_plan(self) -> None: + """No-op retained for compatibility after blueprint percentage removal.""" + + return + + def apply_dataset_plan_preset(self, preset: str) -> None: + """No-op retained for compatibility after blueprint percentage removal. + + Args: + preset: Preset label from the Dataset Blueprint combo box. + """ + + return + + def apply_dataset_plan_to_ingestion(self) -> None: + """Clear ingestion mixture overrides (category percentages are disabled).""" + + self._set_mixture_weights({}) + if hasattr(self, "dataset_log"): + self.dataset_log.append("Dataset blueprint applied: category percentages are disabled.") + self.project_state.setText("Blueprint applied") + LOGGER.info("Dataset blueprint applied with category percentages disabled") + + def _mixture_weights_from_ui(self) -> dict[str, float]: + """Return dataset mixture weights from the Ingest tab. + + Returns: + Empty mapping because category percentages are disabled. + """ + + if not hasattr(self, "_mixture_weights_state"): + self._mixture_weights_state = {} + return {} + + def _set_mixture_weights(self, weights: dict[str, Any]) -> None: + """Restore dataset mixture weights. + + Args: + weights: Saved mixture weights by source family. + """ + + self._mixture_weights_state = {} + + def _update_mixture_total(self) -> None: + """No-op retained for compatibility after mixture percentage removal.""" + + return + + def _normalize_mixture_weights(self) -> None: + """No-op retained for compatibility after mixture percentage removal.""" + + return + + def _training_launch_target_value(self) -> str: + """Return whether training should launch locally or remotely. + + Returns: + ``local`` or ``remote``. + """ + + if self.training_launch_target.currentText() == "RunPod cloud": + return "runpod" + return "remote" if self.training_launch_target.currentText() == "Remote workers" else "local" + + def _fine_tune_launch_target_value(self) -> str: + """Return whether fine-tuning should launch locally or remotely. + + Returns: + ``local`` or ``remote``. + """ + + if not hasattr(self, "fine_tune_launch_target"): + return "local" + if self.fine_tune_launch_target.currentText() == "RunPod cloud": + return "runpod" + return "remote" if self.fine_tune_launch_target.currentText() == "Remote workers" else "local" + + def _architecture_style_config(self) -> dict[str, Any]: + """Return ModelConfig keyword arguments for the selected block style. + + Returns: + Architecture style settings. + """ + + if self.architecture_style.currentText() == "Llama-like": + return { + "norm_type": "rmsnorm", + "position_encoding": "rope", + "mlp_type": "swiglu", + "rope_theta": self.rope_theta.value(), + } + return { + "norm_type": "layernorm", + "position_encoding": "learned", + "mlp_type": "gelu", + "rope_theta": self.rope_theta.value(), + } + + def _optimizer_value(self) -> str: + """Return the selected optimizer identifier. + + Returns: + Stable optimizer name used by the trainer. + """ + + return { + "AdamW": "adamw", + "Adam": "adam", + "Lion": "lion", + "Adafactor": "adafactor", + }.get(self.optimizer_name.currentText(), "adamw") + + def _scheduler_value(self) -> str: + """Return the selected scheduler identifier. + + Returns: + Stable scheduler name used by the trainer. + """ + + return { + "Warmup linear": "warmup_linear", + "Cosine decay": "cosine", + "Polynomial decay": "polynomial", + "One-cycle": "one_cycle", + "Constant": "constant", + }.get(self.scheduler_name.currentText(), "warmup_linear") + + def _precision_value(self) -> str: + """Return the selected numeric precision identifier. + + Returns: + Stable precision name used by the trainer. + """ + + return { + "FP16": "fp16", + "BF16": "bf16", + "FP32": "fp32", + }.get(self.precision.currentText(), "fp16") + + def _fine_tune_output_path(self) -> Path: + """Return the selected fine-tune output folder. + + Returns: + Folder where fine-tuned artifacts should be written. + """ + + text = self.fine_tune_output_dir.text().strip() if hasattr(self, "fine_tune_output_dir") else "" + if text: + path = Path(text) + elif self.current_project_file is not None: + path = self.current_project_file.parent / "fine_tunes" / "latest" + else: + path = Path(self.model_dir.text()) / "fine_tuned" + try: + if path.resolve() == Path(self.model_dir.text()).resolve(): + path = Path(self.model_dir.text()) / "fine_tuned" + except OSError: + pass + if hasattr(self, "fine_tune_output_dir"): + self.fine_tune_output_dir.setText(str(path)) + return path + + def _refresh_fine_tune_default_output(self, *_args: Any) -> None: + """Keep the fine-tune output folder stage-specific unless a custom folder was chosen.""" + + if not hasattr(self, "fine_tune_output_dir") or self.current_project_file is None: + return + project_dir = self.current_project_file.parent + fine_tunes_dir = project_dir / "fine_tunes" + stage = self._training_stage_value() + stage_folder = { + "instruction": "instruction_latest", + "conversation": "conversation_latest", + "code": "code_latest", + "domain": "domain_latest", + }.get(stage, "fine_tune_latest") + desired = fine_tunes_dir / stage_folder + current_text = self.fine_tune_output_dir.text().strip() + if not current_text: + self.fine_tune_output_dir.setText(str(desired)) + return + try: + current = Path(current_text) + current_resolved = current.resolve() + fine_tunes_resolved = fine_tunes_dir.resolve() + except OSError: + return + managed_names = { + "latest", + "fine_tune", + "fine_tuned", + "instruction", + "conversation", + "code", + "domain", + "instruction_latest", + "conversation_latest", + "code_latest", + "domain_latest", + "fine_tune_latest", + } + if current_resolved.parent == fine_tunes_resolved and current.name in managed_names: + self.fine_tune_output_dir.setText(str(desired)) + + def _prepare_fine_tune_run_folder(self, training_config: TrainingConfig) -> None: + """Create fine-tune folders and snapshot the base checkpoint. + + Args: + training_config: Fine-tune training configuration. + """ + + output_dir = Path(training_config.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + checkpoints_dir = output_dir / "checkpoints" + checkpoints_dir.mkdir(parents=True, exist_ok=True) + base_checkpoint = training_config.fine_tune_from_checkpoint + if base_checkpoint is None: + return + base_checkpoint = Path(base_checkpoint) + if not base_checkpoint.exists(): + return + try: + base_resolved = base_checkpoint.resolve() + output_resolved = output_dir.resolve() + if base_resolved == (output_resolved / base_checkpoint.name) or output_resolved in base_resolved.parents: + raise ValueError( + "Fine-tune base checkpoint must be outside the selected fine-tune output folder. " + "Choose the original pretrained model checkpoint instead." + ) + except RuntimeError as exc: + raise ValueError(f"Could not validate fine-tune base checkpoint path: {exc}") from exc + snapshot_dir = output_dir / "base_model" + snapshot_dir.mkdir(parents=True, exist_ok=True) + copied_checkpoint = snapshot_dir / base_checkpoint.name + if not copied_checkpoint.exists() or copied_checkpoint.stat().st_size != base_checkpoint.stat().st_size: + shutil.copy2(base_checkpoint, copied_checkpoint) + base_parent = base_checkpoint.parent + for file_name in ("tokenizer.json", "training_summary.json", "model_lineage.json"): + source = base_parent / file_name + if source.exists(): + target = snapshot_dir / file_name + if not target.exists() or target.stat().st_size != source.stat().st_size: + shutil.copy2(source, target) + manifest = { + "base_checkpoint": str(base_checkpoint), + "copied_checkpoint": str(copied_checkpoint), + "fine_tune_output": str(output_dir), + "created_at": datetime.now().isoformat(timespec="seconds"), + } + (snapshot_dir / "base_model_manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8") + self.fine_tune_log.append(f"Base model snapshot: {copied_checkpoint}") + + def _training_output_dir_for_mode(self, training_mode: Optional[str]) -> Path: + """Return the output folder for a training mode. + + Args: + training_mode: Training mode override. + + Returns: + Base model or fine-tune output folder. + """ + + return self._fine_tune_output_path() if training_mode == "fine_tune" else Path(self.model_dir.text()) + + diff --git a/interface/main_window_part13.py b/interface/main_window_part13.py new file mode 100644 index 0000000..41a78eb --- /dev/null +++ b/interface/main_window_part13.py @@ -0,0 +1,439 @@ +from __future__ import annotations + +# MainWindow implementation mixin. Runtime names are provided by app.py. +from typing import Any +from . import app as _app + +globals().update({name: value for name, value in vars(_app).items() if not name.startswith("__")}) + +class MainWindowPart13: + def _training_mode_value(self) -> str: + """Return the selected training mode identifier. + + Returns: + Stable training mode used by the trainer. + """ + + return { + "Pretrain from scratch": "pretrain", + "Fine-tune checkpoint": "fine_tune", + "Instruction fine-tune": "fine_tune", + "Conversation fine-tune": "fine_tune", + "Code fine-tune": "fine_tune", + }.get(self.training_mode.currentText(), "pretrain") + + def _training_stage_value(self) -> str: + """Return the higher-level training stage selected in the UI. + + Returns: + Training stage identifier. + """ + + return { + "Pretrain from scratch": "base", + "Fine-tune checkpoint": "domain", + "Instruction fine-tune": "instruction", + "Conversation fine-tune": "conversation", + "Code fine-tune": "code", + }.get(self.training_mode.currentText(), "base") + + def _peft_method_value(self) -> str: + """Return the selected PEFT method identifier. + + Returns: + Stable PEFT method used by the trainer. + """ + + return { + "Full fine-tune": "none", + "LoRA adapters": "lora", + }.get(self.peft_method.currentText(), "none") + + def _lora_target_value(self) -> str: + """Return selected LoRA target groups. + + Returns: + Comma-separated target group string. + """ + + return { + "Attention projections": "attention", + "MLP projections": "mlp", + "Attention + MLP": "attention,mlp", + }.get(self.lora_targets.currentText(), "attention") + + def _update_training_mode_controls(self) -> None: + """Enable fine-tune controls only when fine-tuning is selected.""" + + enabled = self._training_mode_value() == "fine_tune" + lora_enabled = enabled and self._peft_method_value() == "lora" + self.fine_tune_checkpoint.setEnabled(enabled) + self.peft_method.setEnabled(enabled) + self.fine_tune_check_button.setEnabled(enabled) + self.lora_rank.setEnabled(lora_enabled) + self.lora_alpha.setEnabled(lora_enabled) + self.lora_dropout.setEnabled(lora_enabled) + self.lora_targets.setEnabled(lora_enabled) + self.refresh_fine_tune_workflow() + + def _current_dataset_summary(self) -> dict[str, Any]: + """Read the active prepared dataset summary. + + Returns: + Dataset summary dictionary, or an empty dictionary. + """ + + summary_path = Path(self.train_data_dir.text()) / "dataset_summary.json" + if not summary_path.exists(): + summary_path = Path(self.dataset_dir.text()) / "dataset_summary.json" + if not summary_path.exists(): + return {} + try: + data = json.loads(summary_path.read_text(encoding="utf-8")) + return data if isinstance(data, dict) else {} + except Exception as exc: + LOGGER.warning("Could not read dataset summary %s: %s", summary_path, exc) + return {} + + def _fine_tune_dataset_stage_status(self) -> tuple[bool, str]: + """Check whether the prepared dataset matches the fine-tune type. + + Returns: + Tuple containing whether the workflow may proceed and a user-facing message. + """ + + expected_stage = self._training_stage_value() + summary = self._current_dataset_summary() + if not summary: + return False, "Dataset: not prepared. Prepare the fine-tune dataset first." + dataset_stage = str(summary.get("dataset_stage") or self._dataset_stage_value()) + tokens = int(summary.get("token_count", 0) or 0) + vocab = int(summary.get("tokenizer_vocab_size", 0) or 0) + stage_name = dataset_stage_label(dataset_stage) if dataset_stage in {"base", "instruction", "conversation", "code"} else dataset_stage + details = f"{stage_name}, {tokens:,} tokens, vocab {vocab:,}" + if expected_stage == "instruction" and dataset_stage != "instruction": + return False, f"Dataset mismatch: selected Instruction fine-tune, but prepared dataset is {details}." + if expected_stage == "conversation" and dataset_stage != "conversation": + return False, f"Dataset mismatch: selected Conversation fine-tune, but prepared dataset is {details}." + if expected_stage == "code" and dataset_stage != "code": + return False, f"Dataset mismatch: selected Code fine-tune, but prepared dataset is {details}." + if expected_stage == "domain" and dataset_stage == "base": + return True, f"Dataset warning: {details}. Base datasets usually belong to pretraining; continue only for domain adaptation." + return True, f"Dataset ready: {details}." + + def refresh_fine_tune_workflow(self) -> None: + """Refresh fine-tune workflow guidance in the Fine-Tuning tab.""" + + if not hasattr(self, "fine_tune_dataset_status"): + return + self._refresh_fine_tune_default_output() + ok, message = self._fine_tune_dataset_stage_status() + self.fine_tune_dataset_status.setText(message) + self.fine_tune_dataset_status.setProperty("state", "ok" if ok else "warning") + self.fine_tune_dataset_status.style().unpolish(self.fine_tune_dataset_status) + self.fine_tune_dataset_status.style().polish(self.fine_tune_dataset_status) + + def apply_recommended_fine_tune_settings(self) -> None: + """Apply conservative fine-tuning defaults for the selected workflow.""" + + stage = self._training_stage_value() + synced = self._sync_architecture_from_fine_tune_base() + self._set_combo_text(self.peft_method, "LoRA adapters") + self.lora_dropout.setValue(0.05) + self._set_combo_text(self.lora_targets, "Attention projections") + self.max_grad_norm.setValue(0.5) + self.weight_decay.setValue(0.05) + self._set_combo_by_data(self.scheduler_name, "cosine", { + "warmup_linear": "Warmup linear", + "cosine": "Cosine decay", + "polynomial": "Polynomial decay", + "one_cycle": "One-cycle", + "constant": "Constant", + }) + if stage == "conversation": + self.lora_rank.setValue(16) + self.lora_alpha.setValue(32.0) + self.learning_rate.setValue(0.00003) + self.epochs.setValue(max(1, min(self.epochs.value(), 2))) + elif stage == "code": + self.lora_rank.setValue(8) + self.lora_alpha.setValue(16.0) + self.lora_dropout.setValue(0.05) + self.learning_rate.setValue(0.00005) + self.max_grad_norm.setValue(0.5) + self.epochs.setValue(max(1, min(self.epochs.value(), 3))) + elif stage == "instruction": + self.lora_rank.setValue(8) + self.lora_alpha.setValue(16.0) + self.learning_rate.setValue(0.00005) + self.epochs.setValue(max(1, min(self.epochs.value(), 3))) + else: + self.lora_rank.setValue(8) + self.lora_alpha.setValue(16.0) + self.learning_rate.setValue(0.00005) + self._update_training_mode_controls() + message = "Recommended LoRA settings applied." + if synced: + message += "\nArchitecture was synced from the selected base checkpoint." + message += "\nUse Check Fine-tune before starting so checkpoint and tokenizer compatibility are verified." + self.fine_tune_preview.setText(message) + + def _sync_architecture_from_fine_tune_base(self) -> bool: + """Sync architecture controls from the selected fine-tune base checkpoint. + + Returns: + True when a checkpoint was read and architecture controls were updated. + """ + + if not hasattr(self, "fine_tune_checkpoint"): + return False + checkpoint_text = self.fine_tune_checkpoint.text().strip() + if not checkpoint_text: + return False + checkpoint_path = Path(checkpoint_text) + if not checkpoint_path.exists(): + return False + try: + checkpoint = torch.load(checkpoint_path, map_location="cpu") + except Exception as exc: + LOGGER.warning("Could not read fine-tune base checkpoint %s: %s", checkpoint_path, exc) + return False + model_config = checkpoint.get("model_config", {}) if isinstance(checkpoint, dict) else {} + if not isinstance(model_config, dict): + return False + mappings = { + "embedding_size": self.n_embd, + "head_count": self.n_head, + "layer_count": self.n_layer, + # NOT self.context_length -- that is the Dataset tab's tokenizer + # window-size setting (DatasetConfig.context_length), an + # unrelated dataset-preparation parameter. _current_model_config() + # reads self.train_context_length for ModelConfig.context_length, + # which is the field resume-compatibility actually checks. + "context_length": self.train_context_length, + } + for key, widget in mappings.items(): + if key in model_config: + try: + widget.setValue(int(model_config[key])) + except (TypeError, ValueError): + LOGGER.warning("Invalid %s in checkpoint %s: %r", key, checkpoint_path, model_config[key]) + if "dropout" in model_config: + try: + self.dropout.setValue(float(model_config["dropout"])) + except (TypeError, ValueError): + LOGGER.warning("Invalid dropout in checkpoint %s: %r", checkpoint_path, model_config["dropout"]) + if "rope_theta" in model_config: + try: + self.rope_theta.setValue(float(model_config["rope_theta"])) + except (TypeError, ValueError): + LOGGER.warning("Invalid rope_theta in checkpoint %s: %r", checkpoint_path, model_config["rope_theta"]) + if "bias" in model_config: + self.use_bias.setChecked(bool(model_config["bias"])) + norm_type = str(model_config.get("norm_type", "layernorm")).lower() + position_encoding = str(model_config.get("position_encoding", "learned")).lower() + mlp_type = str(model_config.get("mlp_type", "gelu")).lower() + if norm_type == "rmsnorm" or position_encoding == "rope" or mlp_type == "swiglu": + # Must match training_tab.py's actual combo item text exactly + # ("Llama-like") -- _set_combo_text() silently no-ops on a + # non-editable combo when the text doesn't match any item, so a + # wrong string here does not raise or log anything. It used to + # say "Modern LLM", which does not exist as an option: this + # left architecture_style un-synced while every other field + # (n_embd, n_head, n_layer, ...) synced correctly, guaranteeing + # a resume-compatibility mismatch on norm_type/position_encoding + # /mlp_type with no indication of why. + self._set_combo_text(self.architecture_style, "Llama-like") + else: + self._set_combo_text(self.architecture_style, "Classic GPT") + attention_type = str(model_config.get("attention_type", "mha")).lower() + self._set_combo_by_data( + self.attention_type, + attention_type, + { + "mha": "Multi-head", + "mqa": "Multi-query", + "gqa": "Grouped-query", + }, + ) + if "kv_head_count" in model_config: + try: + self.kv_head_count.setValue(int(model_config["kv_head_count"])) + except (TypeError, ValueError): + LOGGER.warning("Invalid kv_head_count in checkpoint %s: %r", checkpoint_path, model_config["kv_head_count"]) + backend = str(model_config.get("attention_backend", "sdpa")).lower() + self._set_combo_by_data( + self.attention_backend, + backend, + { + "sdpa": "SDPA / Flash when available", + "eager": "PyTorch eager", + }, + ) + if "attention_window" in model_config: + try: + self.attention_window.setValue(int(model_config["attention_window"])) + except (TypeError, ValueError): + LOGGER.warning("Invalid attention_window in checkpoint %s: %r", checkpoint_path, model_config["attention_window"]) + LOGGER.info("Fine-tune architecture synced from base checkpoint: %s", checkpoint_path) + return True + + def _attention_type_value(self) -> str: + """Return the selected attention layout identifier. + + Returns: + Stable attention type used by the model. + """ + + return { + "Multi-head": "mha", + "Grouped-query": "gqa", + "Multi-query": "mqa", + }.get(self.attention_type.currentText(), "mha") + + def _attention_backend_value(self) -> str: + """Return the selected attention backend identifier. + + Returns: + Stable attention backend used by the model. + """ + + return { + "SDPA / Flash when available": "sdpa", + "Manual": "manual", + }.get(self.attention_backend.currentText(), "sdpa") + + def apply_training_profile(self) -> None: + """Apply the selected optimizer/scheduler/regularization profile. + + Each branch below explicitly sets every field it conceptually owns + (optimizer, scheduler, LR/regularization, precision/memory knobs, + batch shape, and early-stopping patience), even fields that happen + to match the previous profile's value. This is deliberate: profiles + must be idempotent when switched between, or a field set by a + previously applied profile (e.g. activation_checkpointing=True from + Low-memory) can silently survive into a later profile that never + mentions it, producing a configuration no single profile actually + intended. + + Two categories of fields are deliberately NOT touched here: + - attention_type / kv_head_count: an architecture choice, not a + training-strategy choice. Low-memory sets these to + Grouped-query because that specific profile is about reducing + memory end-to-end; the other profiles leave whatever the user + has selected alone rather than silently reverting it. + - training_mode / peft_method / lora_* (Code fine-tune only): + these belong to the fine-tuning tab's widgets, not this tab's. + """ + + profile = self.training_profile.currentText() + if profile == "Low-memory": + self._set_combo_text(self.optimizer_name, "Adafactor") + self._set_combo_text(self.scheduler_name, "Cosine decay") + self.learning_rate.setValue(0.0002) + self.weight_decay.setValue(0.05) + self.min_lr_ratio.setValue(0.05) + self.polynomial_power.setValue(1.0) + self.max_grad_norm.setValue(1.0) + self._set_combo_text(self.precision, "BF16" if torch.cuda.is_available() else "FP32") + self.use_amp.setChecked(True) + self._set_combo_text(self.attention_type, "Grouped-query") + self.kv_head_count.setValue(max(1, self.n_head.value() // 2)) + self.activation_checkpointing.setChecked(True) + # The two knobs that most directly control peak memory: shrink + # the batch and make up the lost effective batch size with + # gradient accumulation, and avoid extra data-loader worker + # processes competing for memory. + self.batch_size.setValue(4) + self.gradient_accumulation.setValue(4) + self.data_loader_workers.setValue(0) + self.warmup_steps.setValue(100) + self.dropout.setValue(0.1) + self.early_stopping_patience.setValue(3) + elif profile == "Code fine-tune": + self._set_combo_text(self.optimizer_name, "AdamW") + self._set_combo_text(self.scheduler_name, "Cosine decay") + self.learning_rate.setValue(0.00005) + self.weight_decay.setValue(0.05) + self.min_lr_ratio.setValue(0.1) + self.polynomial_power.setValue(1.0) + self.max_grad_norm.setValue(0.5) + self._set_combo_text(self.precision, "FP16") + self.use_amp.setChecked(True) + self.activation_checkpointing.setChecked(False) + self.batch_size.setValue(16) + self.gradient_accumulation.setValue(1) + self.data_loader_workers.setValue(0) + self.warmup_steps.setValue(50) + self.dropout.setValue(0.05) + # Fine-tuning generally needs less patience than a full + # pretraining run before validation loss plateaus meaningfully. + self.early_stopping_patience.setValue(2) + self._set_combo_text(self.training_mode, "Fine-tune checkpoint") + self._set_combo_text(self.peft_method, "LoRA adapters") + self.lora_rank.setValue(8) + self.lora_alpha.setValue(16.0) + self.lora_dropout.setValue(0.05) + self._set_combo_text(self.lora_targets, "Attention projections") + elif profile == "Experimental Lion": + self._set_combo_text(self.optimizer_name, "Lion") + self._set_combo_text(self.scheduler_name, "One-cycle") + self.learning_rate.setValue(0.0001) + self.weight_decay.setValue(0.1) + self.min_lr_ratio.setValue(0.01) + self.polynomial_power.setValue(1.0) + self.max_grad_norm.setValue(1.0) + # Lion is reported to be more sensitive to fp16 under/overflow + # than AdamW; prefer bf16 where available, fp32 otherwise. + self._set_combo_text(self.precision, "BF16" if torch.cuda.is_available() else "FP32") + self.use_amp.setChecked(True) + self.activation_checkpointing.setChecked(False) + self.batch_size.setValue(16) + self.gradient_accumulation.setValue(1) + self.data_loader_workers.setValue(0) + self.warmup_steps.setValue(100) + self.dropout.setValue(0.1) + self.early_stopping_patience.setValue(3) + else: + self._set_combo_text(self.optimizer_name, "AdamW") + self._set_combo_text(self.scheduler_name, "Cosine decay") + self.learning_rate.setValue(0.0003) + self.weight_decay.setValue(0.1) + self.min_lr_ratio.setValue(0.1) + self.polynomial_power.setValue(1.0) + self.max_grad_norm.setValue(1.0) + self._set_combo_text(self.precision, "FP16") + self.use_amp.setChecked(True) + self.activation_checkpointing.setChecked(False) + self.batch_size.setValue(16) + self.gradient_accumulation.setValue(1) + self.data_loader_workers.setValue(0) + self.warmup_steps.setValue(100) + self.dropout.setValue(0.1) + self.early_stopping_patience.setValue(3) + self._update_training_mode_controls() + self.refresh_model_estimate() + self.training_log.append(f"Applied training profile: {profile}") + + def _tokenizer_strategy_reuses(self) -> bool: + """Return whether current tokenizer strategy ignores vocabulary controls. + + Returns: + True when an existing tokenizer is selected directly. + """ + + return self.tokenizer_strategy.currentText() in {"Reuse dataset tokenizer", "Import tokenizer.json"} + + def _update_tokenizer_strategy_controls(self) -> None: + """Enable only the tokenizer inputs relevant to the selected strategy.""" + + imports_tokenizer = self.tokenizer_strategy.currentText() == "Import tokenizer.json" + reuses_tokenizer = self._tokenizer_strategy_reuses() + if hasattr(self, "tokenizer_path_row"): + self.tokenizer_path_row.setEnabled(imports_tokenizer) + self.tokenizer_path.setEnabled(imports_tokenizer) + self.auto_vocab.setEnabled(not reuses_tokenizer) + self.manual_vocab_size.setEnabled(not reuses_tokenizer and not self.auto_vocab.isChecked()) + self.min_frequency.setEnabled(not reuses_tokenizer) + + diff --git a/interface/main_window_part14.py b/interface/main_window_part14.py new file mode 100644 index 0000000..21fff62 --- /dev/null +++ b/interface/main_window_part14.py @@ -0,0 +1,403 @@ +from __future__ import annotations + +# MainWindow implementation mixin. Runtime names are provided by app.py. +from typing import Any +from . import app as _app + +globals().update({name: value for name, value in vars(_app).items() if not name.startswith("__")}) + +class MainWindowPart14: + def _update_model_estimate_chips( + self, + estimate: dict[str, Any], + model_config: Optional[ModelConfig] = None, + training_config: Optional[TrainingConfig] = None, + train_tokens: int = 0, + ) -> None: + """Update model and VRAM estimate chips. + + Args: + estimate: Estimate dictionary from the training planning service. + model_config: Model architecture used for the estimate. + training_config: Training options used for the estimate. + train_tokens: Number of available training tokens. + """ + + params = int(estimate.get("parameters", 0)) + checkpoint_bytes = float(estimate.get("checkpoint_bytes", 0)) + vram_bytes = float(estimate.get("vram_bytes", 0)) + self.model_size_metric.setText(f"Model: {params / 1_000_000:.2f}M, ckpt {format_bytes(checkpoint_bytes)}") + self.vram_estimate_metric.setText(f"VRAM est: {format_bytes(vram_bytes)}") + parameter_breakdown = estimate.get("parameter_breakdown", {}) or {} + memory_breakdown = estimate.get("memory_breakdown", {}) or {} + embedding_params = int(parameter_breakdown.get("token_embedding", 0)) + int( + parameter_breakdown.get("position_embedding", 0) + ) + attention_params = int(parameter_breakdown.get("attention", 0)) + mlp_params = int(parameter_breakdown.get("mlp", 0)) + norm_params = int(parameter_breakdown.get("norms", 0)) + self.parameter_breakdown_metric.setText( + "Params: " + f"emb {self._compact_number(embedding_params)}, " + f"attn {self._compact_number(attention_params)}, " + f"mlp {self._compact_number(mlp_params)}" + ) + self._tip( + self.parameter_breakdown_metric, + ( + f"Embedding: {embedding_params:,}\n" + f"Attention: {attention_params:,}\n" + f"MLP: {mlp_params:,}\n" + f"Norms/output: {norm_params:,}\n" + f"Total: {params:,}" + ), + ) + weights = float(memory_breakdown.get("weights", 0)) + optimizer = float(memory_breakdown.get("optimizer", 0)) + activations = float(memory_breakdown.get("activations", 0)) + kv_cache = float(memory_breakdown.get("kv_cache", 0)) + self.memory_breakdown_metric.setText( + f"Memory: w {format_bytes(weights)}, opt {format_bytes(optimizer)}, act {format_bytes(activations)}" + ) + self._tip( + self.memory_breakdown_metric, + ( + f"Weights: {format_bytes(weights)}\n" + f"Optimizer state: {format_bytes(optimizer)}\n" + f"Activations: {format_bytes(activations)}\n" + f"KV cache estimate: {format_bytes(kv_cache)}\n" + f"Total training estimate: {format_bytes(vram_bytes)}" + ), + ) + self._update_architecture_advisor(estimate, model_config, training_config, train_tokens) + + def _update_architecture_advisor( + self, + estimate: dict[str, Any], + model_config: Optional[ModelConfig], + training_config: Optional[TrainingConfig], + train_tokens: int, + ) -> None: + """Update the compact architecture advisor chip. + + Args: + estimate: Estimate dictionary from the training planning service. + model_config: Model architecture used for the estimate. + training_config: Training options used for the estimate. + train_tokens: Number of available training tokens. + """ + + params = max(int(estimate.get("parameters", 0) or 0), 1) + tokens_per_param = float(train_tokens) / float(params) if train_tokens > 0 else 0.0 + vram_bytes = float(estimate.get("vram_bytes", 0) or 0) + notes: list[str] = [] + if tokens_per_param <= 0: + label = "Advisor: prepare data" + notes.append("Prepare a dataset to compare token budget against model size.") + elif tokens_per_param < 20: + label = "Advisor: data-light" + notes.append( + f"Token budget is about {tokens_per_param:.1f} tokens per parameter. More data or fewer epochs may reduce overfitting." + ) + elif tokens_per_param > 150: + label = "Advisor: data-rich" + notes.append( + f"Token budget is about {tokens_per_param:.1f} tokens per parameter. The model may be small for this much data." + ) + else: + label = "Advisor: balanced" + notes.append(f"Token budget is about {tokens_per_param:.1f} tokens per parameter.") + if model_config is not None: + if model_config.context_length >= 2048 and model_config.embedding_size <= 256: + notes.append("Long context with a small embedding can be memory-heavy without adding much capacity.") + if model_config.attention_type in {"grouped_query", "multi_query"}: + notes.append("Grouped/multi-query attention reduces KV memory and is useful for longer contexts.") + if model_config.mlp_type == "swiglu" and model_config.norm_type == "rmsnorm": + notes.append("Llama-like blocks improve modern compatibility but must match checkpoints when resuming.") + if training_config is not None and training_config.device == "cuda" and vram_bytes > 3.5 * 1024**3: + notes.append("Estimated VRAM is high for 4 GB GPUs. Try lower batch, context, embedding, or layers.") + if label == "Advisor: balanced": + label = "Advisor: memory check" + self.architecture_advisor_metric.setText(label) + self._tip(self.architecture_advisor_metric, "\n".join(notes)) + + @staticmethod + def _compact_number(value: int) -> str: + """Format a large count for tight metric chips. + + Args: + value: Count to format. + + Returns: + Compact display string. + """ + + magnitude = abs(value) + if magnitude >= 1_000_000_000: + return f"{value / 1_000_000_000:.1f}B" + if magnitude >= 1_000_000: + return f"{value / 1_000_000:.1f}M" + if magnitude >= 1_000: + return f"{value / 1_000:.1f}K" + return str(value) + + def _current_model_config(self, vocab_size: int = 1) -> ModelConfig: + """Build a model config from the current AI tab settings. + + Args: + vocab_size: Tokenizer vocabulary size to use. + + Returns: + Current model configuration. + """ + + return ModelConfig( + vocab_size=vocab_size, + context_length=self.train_context_length.value(), + embedding_size=self.n_embd.value(), + head_count=self.n_head.value(), + layer_count=self.n_layer.value(), + dropout=self.dropout.value(), + bias=self.use_bias.isChecked(), + attention_type=self._attention_type_value(), + kv_head_count=self.kv_head_count.value(), + attention_backend=self._attention_backend_value(), + attention_window=self.attention_window.value(), + **self._architecture_style_config(), + ) + + def _current_training_config( + self, + resume_path: Optional[Path] = None, + training_mode: Optional[str] = None, + ) -> TrainingConfig: + """Build a training config from the current AI tab settings. + + Args: + resume_path: Optional specific checkpoint to resume from. + training_mode: Optional explicit training mode override. + + Returns: + Current training configuration. + """ + + return TrainingConfig( + output_dir=self._training_output_dir_for_mode(training_mode), + epochs=self.epochs.value(), + batch_size=self.batch_size.value(), + learning_rate=self.learning_rate.value(), + weight_decay=self.weight_decay.value(), + optimizer_name=self._optimizer_value(), + scheduler_name=self._scheduler_value(), + scheduler_min_lr_ratio=self.min_lr_ratio.value(), + polynomial_power=self.polynomial_power.value(), + gradient_accumulation=self.gradient_accumulation.value(), + sample_stride=self.sample_stride.value(), + warmup_steps=self.warmup_steps.value(), + eval_interval=self.eval_interval.value(), + max_eval_batches=self.max_eval_batches.value(), + save_interval=self.save_interval.value(), + data_loader_workers=self.data_loader_workers.value(), + max_grad_norm=self.max_grad_norm.value(), + activation_checkpointing=self.activation_checkpointing.isChecked(), + device=self.device.currentText(), + use_amp=self.use_amp.isChecked(), + precision=self._precision_value(), + seed=self.seed.value(), + training_mode=training_mode or self._training_mode_value(), + fine_tune_from_checkpoint=( + Path(self.fine_tune_checkpoint.text()) + if training_mode != "pretrain" and self.fine_tune_checkpoint.text().strip() + else None + ), + peft_method="none" if training_mode == "pretrain" else self._peft_method_value(), + lora_rank=self.lora_rank.value(), + lora_alpha=self.lora_alpha.value(), + lora_dropout=self.lora_dropout.value(), + lora_target_modules=self._lora_target_value(), + resume=self.resume_training.isChecked(), + resume_from_checkpoint=resume_path if self.resume_training.isChecked() else None, + require_compatible_resume=self.resume_safety.isChecked(), + early_stopping=self.early_stopping.isChecked(), + early_stopping_patience=self.early_stopping_patience.value(), + ) + + def _current_training_vocab_size(self, data_dir: Path) -> int: + """Return the tokenizer vocabulary size for the current training dataset. + + Args: + data_dir: Prepared dataset folder. + + Returns: + Vocabulary size, or zero if unavailable. + """ + + summary_path = data_dir / "dataset_summary.json" + if summary_path.exists(): + summary = json.loads(summary_path.read_text(encoding="utf-8")) + vocab_size = int(summary.get("tokenizer_vocab_size", 0) or 0) + if vocab_size > 0: + return vocab_size + tokenizer_path = data_dir / "tokenizer.json" + if tokenizer_path.exists(): + tokenizer_data = json.loads(tokenizer_path.read_text(encoding="utf-8")) + vocab = tokenizer_data.get("model", {}).get("vocab", {}) + if isinstance(vocab, dict): + return len(vocab) + return 0 + + def _checkpoint_vocab_size(self, checkpoint_path: Path) -> int: + """Return the tokenizer vocabulary size saved in a checkpoint. + + Args: + checkpoint_path: Checkpoint file to inspect. + + Returns: + Saved vocabulary size, or zero when unavailable. + """ + + try: + checkpoint = torch.load(checkpoint_path, map_location="cpu") + model_config = checkpoint.get("model_config", {}) + if isinstance(model_config, dict): + return int(model_config.get("vocab_size", 0) or 0) + except Exception as exc: + LOGGER.warning("Could not inspect checkpoint vocab size for %s: %s", checkpoint_path, exc) + return 0 + + @staticmethod + def _tokenizer_mismatch_help(checkpoint_vocab: int, dataset_vocab: int) -> str: + """Return user-facing help for tokenizer mismatch errors. + + Args: + checkpoint_vocab: Vocabulary size saved in the checkpoint. + dataset_vocab: Vocabulary size in the prepared dataset. + + Returns: + Help text. + """ + + return ( + f"Tokenizer mismatch: base checkpoint vocab is {checkpoint_vocab:,}, " + f"but prepared dataset vocab is {dataset_vocab:,}.\n" + "Fix: rebuild the fine-tune dataset using the exact tokenizer from the base model. " + "In Ingest, set Tokenizer policy to Import tokenizer.json and choose the tokenizer.json " + "beside the base checkpoint, then prepare the fine-tune dataset again." + ) + + def _training_run_artifacts(self, output_dir: Path) -> list[Path]: + """Return training-run artifacts in a model output folder. + + Args: + output_dir: Model output folder to inspect. + + Returns: + Existing training-run artifact paths. + """ + + candidates = [ + output_dir / "checkpoints", + output_dir / "final_model.pt", + output_dir / "final_adapter.pt", + output_dir / "training_summary.json", + output_dir / "training_history.json", + output_dir / "model_lineage.json", + ] + return [path for path in candidates if path.exists()] + + def _clear_training_run_artifacts(self, output_dir: Path) -> list[Path]: + """Delete resumable training artifacts from a model output folder. + + Args: + output_dir: Model output folder to clean. + + Returns: + Paths that were removed. + """ + + output_dir = output_dir.resolve() + removed: list[Path] = [] + candidates = self._training_run_artifacts(output_dir) + for path in candidates: + try: + resolved = path.resolve() + except FileNotFoundError: + resolved = path + if output_dir not in resolved.parents and resolved != output_dir: + LOGGER.warning("Skipped training cleanup outside model output folder: %s", path) + continue + if path.is_dir(): + shutil.rmtree(path) + removed.append(path) + elif path.exists(): + path.unlink() + removed.append(path) + return removed + + def _selected_resume_path(self) -> Optional[Path]: + """Return the selected or latest checkpoint path. + + Returns: + Checkpoint path, or ``None`` when no checkpoint exists. + """ + + if self.resume_checkpoint.text().strip(): + return Path(self.resume_checkpoint.text()) + return latest_checkpoint(Path(self.model_dir.text()) / "checkpoints") + + def preview_resume_compatibility(self) -> None: + """Preview whether the selected checkpoint can resume safely.""" + + if not self.resume_training.isChecked(): + self.resume_training_preview.setText("[INFO] Resume latest is off. Enable resume to continue from a checkpoint.") + return + resume_path = self._selected_resume_path() + if resume_path is None: + self.resume_training_preview.setText("[INFO] No checkpoint found in the current model folder.") + return + if not resume_path.exists(): + self.resume_training_preview.setText(f"[BLOCK] Checkpoint does not exist:\n{resume_path}") + return + try: + vocab_size = self._current_training_vocab_size(Path(self.train_data_dir.text())) + if vocab_size <= 0: + self.resume_training_preview.setText("[BLOCK] Could not determine current dataset tokenizer vocabulary size.") + return + model_config = self._current_model_config(vocab_size=vocab_size) + # Explicit override: this is the AI/Training tab's own "Check + # Resume" button, checking a pretrain checkpoint. Without this, + # training_mode falls back to reading the separate Fine-Tuning + # tab's mode combo (self.training_mode), which defaults to + # "Instruction fine-tune" on a fresh session -- resolving to + # "fine_tune" and making training_config.validate() below raise + # "fine_tune_from_checkpoint is required for fine_tune mode", + # a confusing error unrelated to what the user is checking. + training_config = self._current_training_config(resume_path, training_mode="pretrain") + model_config.validate() + training_config.validate() + report = check_resume_compatibility(resume_path, model_config, training_config) + errors = list(report.errors) + if training_config.require_compatible_resume: + if not report.can_load_optimizer_state: + errors.append("Safe resume requires matching optimizer state.") + if not report.can_load_scheduler_state: + errors.append("Safe resume requires matching scheduler state.") + if not report.can_load_scaler_state: + errors.append("Safe resume requires matching AMP scaler state.") + lines: list[str] = [] + if errors: + lines.append("[BLOCK] Resume is not safe with the current settings.") + elif report.warnings: + lines.append("[WARN] Resume is possible, but settings changed.") + else: + lines.append("[OK] Checkpoint can resume with the current settings.") + lines.extend(f"[OK] {line}" for line in report.info) + lines.extend(f"[WARN] {line}" for line in report.warnings) + lines.extend(f"[BLOCK] {line}" for line in errors) + if not training_config.require_compatible_resume and not errors: + lines.append("[INFO] Safe resume is off. Compatible weights will load; incompatible optimizer state may be skipped.") + self.resume_training_preview.setText("\n".join(lines)) + except Exception as exc: + self.resume_training_preview.setText(f"[BLOCK] Could not check resume compatibility:\n{exc}") + + diff --git a/interface/main_window_part15.py b/interface/main_window_part15.py new file mode 100644 index 0000000..b584ce4 --- /dev/null +++ b/interface/main_window_part15.py @@ -0,0 +1,205 @@ +from __future__ import annotations + +# MainWindow implementation mixin. Runtime names are provided by app.py. +from typing import Any +from . import app as _app + +globals().update({name: value for name, value in vars(_app).items() if not name.startswith("__")}) + +class MainWindowPart15: + def preview_fine_tune_compatibility(self) -> None: + """Preview whether the selected checkpoint can be used for fine-tuning.""" + + stage_ok, stage_message = self._fine_tune_dataset_stage_status() + if not stage_ok: + self.fine_tune_preview.setText(f"[BLOCK] {stage_message}") + return + base_path = Path(self.fine_tune_checkpoint.text()) if self.fine_tune_checkpoint.text().strip() else None + if base_path is None: + self.fine_tune_preview.setText("[BLOCK] Choose a base checkpoint for fine-tuning.") + return + if not base_path.exists(): + self.fine_tune_preview.setText(f"[BLOCK] Fine-tune base checkpoint does not exist:\n{base_path}") + return + try: + vocab_size = self._current_training_vocab_size(Path(self.train_data_dir.text())) + if vocab_size <= 0: + self.fine_tune_preview.setText("[BLOCK] Could not determine current dataset tokenizer vocabulary size.") + return + model_config = self._current_model_config(vocab_size=vocab_size) + training_config = self._current_training_config() + model_config.validate() + report = check_resume_compatibility(base_path, model_config, training_config) + lines: list[str] = [] + if report.errors: + lines.append("[BLOCK] Base checkpoint cannot be fine-tuned with the current model/dataset settings.") + else: + lines.append("[OK] Base checkpoint weights can be used for fine-tuning.") + lines.append(f"[OK] {stage_message}" if stage_ok else f"[BLOCK] {stage_message}") + lines.extend(self._fine_tune_lineage_advice(base_path)) + lines.extend(f"[OK] {line}" for line in report.info) + behavior_warnings = [ + warning for warning in report.warnings + if not warning.startswith("Optimizer changed:") and not warning.startswith("LR scheduler changed:") + ] + lines.extend(f"[WARN] {line}" for line in behavior_warnings) + lines.extend(f"[BLOCK] {line}" for line in report.errors) + checkpoint_vocab = self._checkpoint_vocab_size(base_path) + if checkpoint_vocab and checkpoint_vocab != vocab_size: + lines.append(f"[FIX] {self._tokenizer_mismatch_help(checkpoint_vocab, vocab_size)}") + if not report.errors: + lines.append("[INFO] Fine-tuning starts fresh optimizer, scheduler, and scaler state.") + self.fine_tune_preview.setText("\n".join(lines)) + except Exception as exc: + self.fine_tune_preview.setText(f"[BLOCK] Could not check fine-tune compatibility:\n{exc}") + + def _fine_tune_lineage_advice(self, base_path: Path) -> list[str]: + """Return guidance about the selected fine-tune base checkpoint. + + Args: + base_path: Selected checkpoint path. + + Returns: + Lines for the fine-tune compatibility report. + """ + + lines: list[str] = [] + try: + output_dir = self._fine_tune_output_path().resolve() + base_resolved = base_path.resolve() + if output_dir == base_resolved.parent or output_dir in base_resolved.parents: + return [ + "[BLOCK] Selected base checkpoint is inside the current fine-tune output folder.", + "[FIX] Choose the original pretrained model or a completed earlier fine-tune from another folder.", + ] + except OSError: + pass + lineage_path = base_path.parent / "model_lineage.json" + summary_path = base_path.parent / "training_summary.json" + lineage = read_json(lineage_path, default={}) or {} + summary = read_json(summary_path, default={}) or {} + training_mode = str(lineage.get("training_mode") or (summary.get("training_config") or {}).get("training_mode") or "") + stage = str((summary.get("model_lineage") or lineage).get("fine_tune_stage") or "") + if training_mode == "fine_tune": + stage_text = f" ({stage})" if stage else "" + lines.append(f"[INFO] Selected base is a previous fine-tuned checkpoint{stage_text}.") + lines.append("[INFO] This is correct for cumulative tuning, such as conversation -> instruction -> code.") + elif training_mode == "pretrain": + lines.append("[OK] Selected base is the pretrained model checkpoint.") + lines.append("[INFO] This is correct when starting a new independent fine-tune branch.") + else: + lines.append("[INFO] Could not read model lineage; compatibility check will still validate tensor shapes.") + project_base = self.current_project_file.parent / "models" / "final_model.pt" if self.current_project_file else None + if project_base and project_base.exists(): + try: + if base_path.resolve() != project_base.resolve() and training_mode != "fine_tune": + lines.append(f"[HINT] Project pretrained model is: {project_base}") + except OSError: + pass + return lines + + def _training_history_path(self) -> Path: + """Return the training history path for the selected model folder. + + Returns: + Path to ``training_history.json``. + """ + + output_dir = getattr(self, "active_training_output_dir", None) + if output_dir is None: + output_dir = Path(self.model_dir.text()) + return Path(output_dir) / "training_history.json" + + def _load_training_history(self) -> list[dict[str, Any]]: + """Load training run history. + + Returns: + List of training run entries. + """ + + path = self._training_history_path() + if not path.exists(): + return [] + try: + data = json.loads(path.read_text(encoding="utf-8")) + return data if isinstance(data, list) else [] + except Exception: + return [] + + def refresh_model_estimate(self) -> None: + """Refresh model size, rough VRAM, and run history widgets.""" + + model_config = self._current_model_config() + # Same reasoning as preview_resume_compatibility: this is the + # AI/Training tab's shared "Model Estimate" card, not the + # Fine-Tuning tab's; pass an explicit override rather than + # inheriting the Fine-Tuning tab's mode combo by fallback. + training_config = self._current_training_config(training_mode="pretrain") + data_dir = Path(self.train_data_dir.text()) + train_tokens = max(model_config.context_length * training_config.batch_size, 1) + try: + summary_path = data_dir / "dataset_summary.json" + if summary_path.exists(): + summary = json.loads(summary_path.read_text(encoding="utf-8")) + vocab_size = int(summary.get("tokenizer_vocab_size", 0) or 0) + train_tokens = int(summary.get("train_token_count", summary.get("token_count", train_tokens)) or train_tokens) + if vocab_size > 0: + model_config.vocab_size = vocab_size + elif (data_dir / "tokenizer.json").exists(): + tokenizer_data = json.loads((data_dir / "tokenizer.json").read_text(encoding="utf-8")) + vocab = tokenizer_data.get("model", {}).get("vocab", {}) + if vocab: + model_config.vocab_size = len(vocab) + except Exception as exc: + self.training_log.append(f"[WARN] Could not refresh dataset-based estimate: {exc}") + estimate = estimate_training_resources(model_config, training_config, train_tokens) + self.last_training_estimate = estimate + self._update_model_estimate_chips(estimate, model_config, training_config, train_tokens) + self.history_metric.setText(f"Runs: {len(self._load_training_history())}") + self.training_log.append( + "Model estimate refreshed: " + f"{int(estimate['parameters']):,} params, " + f"checkpoint {format_bytes(float(estimate['checkpoint_bytes']))}, " + f"VRAM {format_bytes(float(estimate['vram_bytes']))}." + ) + + def _append_training_history(self, result: Any) -> None: + """Persist a training run entry to ``training_history.json``. + + Args: + result: Training result object. + """ + + history_path = self._training_history_path() + history_path.parent.mkdir(parents=True, exist_ok=True) + history = self._load_training_history() + summary = {} + try: + if Path(result.summary_path).exists(): + summary = json.loads(Path(result.summary_path).read_text(encoding="utf-8")) + except Exception: + summary = {} + estimate = getattr(self, "last_training_estimate", {}) or {} + entry = { + "completed_at": datetime.now().isoformat(timespec="seconds"), + "checkpoint_path": str(result.checkpoint_path), + "summary_path": str(result.summary_path), + "stopped": bool(getattr(result, "stopped", False)), + "final_train_loss": result.final_train_loss, + "final_val_loss": result.final_val_loss, + "best_val_loss": summary.get("best_val_loss"), + "recommended_checkpoint_path": summary.get("recommended_checkpoint_path"), + "best_checkpoint_path": summary.get("best_checkpoint_path"), + "dataset_dir": self.train_data_dir.text(), + "dataset_version": (summary.get("model_lineage") or {}).get("dataset_version"), + "training_run_id": summary.get("training_run_id"), + "parameters": estimate.get("parameters") or summary.get("parameters"), + "model_config": summary.get("model_config"), + "training_config": summary.get("training_config"), + } + history.append(entry) + history_path.write_text(json.dumps(history[-200:], indent=2), encoding="utf-8") + self.history_metric.setText(f"Runs: {len(history[-200:])}") + (self.active_training_log or self.training_log).append(f"Training history updated: {history_path}") + + diff --git a/interface/main_window_part16.py b/interface/main_window_part16.py new file mode 100644 index 0000000..af37946 --- /dev/null +++ b/interface/main_window_part16.py @@ -0,0 +1,391 @@ +from __future__ import annotations + +# MainWindow implementation mixin. Runtime names are provided by app.py. +from typing import Any +from . import app as _app + +globals().update({name: value for name, value in vars(_app).items() if not name.startswith("__")}) + +class MainWindowPart16: + def _run_training_preflight(self, model_config: ModelConfig, training_config: TrainingConfig) -> bool: + """Run pre-training checklist and disk-space guard. + + Args: + model_config: Selected model architecture. + training_config: Selected training settings. + + Returns: + True when training may continue. + """ + + log = self.active_training_log or self.training_log + data_dir = Path(self.train_data_dir.text()) + output_dir = training_config.output_dir + errors: list[str] = [] + warnings: list[str] = [] + info: list[str] = [] + resettable_errors: list[str] = [] + missing: list[str] = [] + if not (data_dir / "tokenizer.json").exists(): + missing.append("tokenizer.json") + has_npy_tokens = (data_dir / "train_tokens.npy").exists() and (data_dir / "val_tokens.npy").exists() + has_json_tokens = (data_dir / "train_tokens.json").exists() and (data_dir / "val_tokens.json").exists() + if not has_npy_tokens and not has_json_tokens: + missing.append("train_tokens.(npy/json), val_tokens.(npy/json)") + if not data_dir.exists(): + errors.append(f"Dataset folder does not exist: {data_dir}") + elif missing: + errors.append(f"Dataset is not prepared. Missing: {', '.join(missing)}") + else: + info.append("Dataset artifacts found.") + + vocab_size = 0 + train_tokens = 0 + val_tokens = 0 + summary = {} + try: + summary_path = data_dir / "dataset_summary.json" + if summary_path.exists(): + summary = json.loads(summary_path.read_text(encoding="utf-8")) + vocab_size = int(summary.get("tokenizer_vocab_size", 0) or 0) + if vocab_size <= 0: + tokenizer_data = json.loads((data_dir / "tokenizer.json").read_text(encoding="utf-8")) + vocab_size = len(tokenizer_data.get("model", {}).get("vocab", {})) + train_tokens = int(summary.get("train_token_count", 0) or 0) + val_tokens = int(summary.get("val_token_count", 0) or 0) + if train_tokens <= 0 and (data_dir / "train_tokens.json").exists(): + train_tokens = len(json.loads((data_dir / "train_tokens.json").read_text(encoding="utf-8"))) + if val_tokens <= 0 and (data_dir / "val_tokens.json").exists(): + val_tokens = len(json.loads((data_dir / "val_tokens.json").read_text(encoding="utf-8"))) + except Exception as exc: + warnings.append(f"Could not fully inspect dataset metadata: {exc}") + + if vocab_size > 0: + model_config.vocab_size = vocab_size + info.append(f"Tokenizer vocab: {vocab_size:,}.") + elif not missing: + errors.append("Could not determine tokenizer vocabulary size.") + if train_tokens and train_tokens <= model_config.context_length: + errors.append("Training token count must be larger than context length.") + elif train_tokens: + info.append(f"Training tokens: {train_tokens:,}; validation tokens: {val_tokens:,}.") + if train_tokens < 50_000: + warnings.append("Training token count is very small; expect smoke-test quality.") + + try: + model_config.validate() + except Exception as exc: + errors.append(f"Model architecture is invalid: {exc}") + try: + training_config.validate() + except Exception as exc: + errors.append(f"Training options are invalid: {exc}") + if model_config.attention_backend == "sdpa": + if hasattr(torch.nn.functional, "scaled_dot_product_attention"): + if training_config.device == "cuda" and torch.cuda.is_available(): + flash_enabled = bool(getattr(torch.backends.cuda, "flash_sdp_enabled", lambda: False)()) + info.append("Attention backend: SDPA selected; Flash Attention may be used by PyTorch." if flash_enabled else "Attention backend: SDPA selected; CUDA flash kernel is not enabled.") + else: + info.append("Attention backend: SDPA selected; CPU/backend fallback will be used if needed.") + else: + warnings.append("SDPA attention selected, but this PyTorch build does not expose scaled_dot_product_attention.") + else: + warnings.append("Manual attention backend selected. This is useful for debugging but can be slower.") + if training_config.peft_method == "lora": + info.append( + "PEFT: LoRA adapters enabled. Intermediate checkpoints will save adapter weights; final_model.pt will be merged." + ) + + if training_config.device == "cuda" and not torch.cuda.is_available(): + errors.append("CUDA is selected, but PyTorch cannot use CUDA on this machine.") + elif training_config.device == "cuda": + info.append(f"CUDA ready: {torch.cuda.get_device_name(0)}.") + if training_config.data_loader_workers > 0: + info.append(f"CPU-assisted batch loading enabled with {training_config.data_loader_workers} worker(s).") + else: + warnings.append("CPU training is selected. This can be very slow.") + if sys.platform.startswith("win") and training_config.data_loader_workers > 4: + warnings.append("High CPU worker counts can duplicate dataset memory on Windows. Start with 2-4 workers and increase carefully.") + + active_resume_path: Optional[Path] = None + resume_path = training_config.resume_from_checkpoint if training_config.resume else None + if resume_path and not Path(resume_path).exists(): + errors.append(f"Selected resume checkpoint does not exist: {resume_path}") + elif training_config.resume: + if resume_path is None: + resume_path = latest_checkpoint(output_dir / "checkpoints") + if resume_path is None: + info.append("Resume latest is enabled, but no checkpoint exists yet.") + else: + active_resume_path = Path(resume_path) + try: + compatibility = check_resume_compatibility(active_resume_path, model_config, training_config) + info.extend(compatibility.info) + warnings.extend(compatibility.warnings) + errors.extend(compatibility.errors) + resettable_errors.extend(compatibility.errors) + if training_config.require_compatible_resume: + if not compatibility.can_load_optimizer_state: + message = "Safe resume requires matching optimizer state." + errors.append(message) + resettable_errors.append(message) + if not compatibility.can_load_scheduler_state: + message = "Safe resume requires matching scheduler state." + errors.append(message) + resettable_errors.append(message) + if not compatibility.can_load_scaler_state: + message = "Safe resume requires matching AMP scaler state." + errors.append(message) + resettable_errors.append(message) + except Exception as exc: + errors.append(f"Could not inspect resume checkpoint: {exc}") + if training_config.training_mode == "fine_tune" and active_resume_path is None: + base_path = training_config.fine_tune_from_checkpoint + if base_path is None: + errors.append("Fine-tune mode requires a base checkpoint.") + elif not Path(base_path).exists(): + errors.append(f"Fine-tune base checkpoint does not exist: {base_path}") + else: + try: + compatibility = check_resume_compatibility(Path(base_path), model_config, training_config) + info.append(f"Fine-tune base checkpoint: {Path(base_path).name}.") + warnings.extend( + warning for warning in compatibility.warnings + if not warning.startswith("Optimizer changed:") and not warning.startswith("LR scheduler changed:") + ) + errors.extend(compatibility.errors) + checkpoint_vocab = self._checkpoint_vocab_size(Path(base_path)) + if checkpoint_vocab and checkpoint_vocab != vocab_size: + errors.append(self._tokenizer_mismatch_help(checkpoint_vocab, vocab_size)) + if not compatibility.errors: + info.append("Fine-tune base weights are compatible. Optimizer state will start fresh.") + except Exception as exc: + errors.append(f"Could not inspect fine-tune base checkpoint: {exc}") + elif training_config.training_mode == "pretrain" and active_resume_path is None: + info.append("A fresh pretraining run will start from random weights.") + elif training_config.training_mode == "fine_tune" and active_resume_path is not None: + info.append("Existing run checkpoint found; training will resume that run instead of reloading the fine-tune base.") + + output_dir.mkdir(parents=True, exist_ok=True) + estimate = estimate_training_resources(model_config, training_config, train_tokens) + self.last_training_estimate = estimate + self._update_model_estimate_chips(estimate, model_config, training_config, train_tokens) + params = int(estimate["parameters"]) + checkpoint_bytes = float(estimate["checkpoint_bytes"]) + checkpoint_count = int(estimate["checkpoint_count"]) + estimated_storage = float(estimate["estimated_storage"]) + estimated_vram = float(estimate["vram_bytes"]) + free_bytes = shutil.disk_usage(output_dir).free + info.append(f"Estimated parameters: {params:,}.") + info.append(f"Estimated checkpoint size: {format_bytes(checkpoint_bytes)}.") + info.append(f"Estimated training VRAM: {format_bytes(estimated_vram)}.") + info.append(f"Estimated training storage need: {format_bytes(estimated_storage)}.") + info.append(f"Free space on model drive: {format_bytes(free_bytes)}.") + if training_config.device == "cuda" and torch.cuda.is_available(): + free_vram, total_vram = torch.cuda.mem_get_info() + info.append(f"GPU free/total VRAM: {format_bytes(free_vram)} / {format_bytes(total_vram)}.") + if estimated_vram > free_vram * 0.9: + warnings.append("Estimated VRAM is close to or above currently free GPU memory.") + if free_bytes < estimated_storage * 1.25: + errors.append("Not enough free disk space for estimated checkpoints and final model.") + elif free_bytes < estimated_storage * 2: + warnings.append("Free disk space is close to the estimated training storage need.") + if checkpoint_count > 50: + warnings.append("Save interval may create many checkpoints. Increase Save every or clean old checkpoints.") + + log.clear() + log.append("Training checklist") + for line in info: + log.append(f"[OK] {line}") + for line in warnings: + log.append(f"[WARN] {line}") + for line in errors: + log.append(f"[ERROR] {line}") + + if errors: + hard_errors = [error for error in errors if error not in resettable_errors] + if active_resume_path is not None and resettable_errors and not hard_errors: + message = ( + "The existing checkpoint was created with different model settings, so it cannot be resumed.\n\n" + "This is expected if you intentionally changed architecture, block style, tokenizer, " + "context length, attention layout, or other checkpoint-shaped settings.\n\n" + "You can start a fresh training run with the current settings. This will delete old " + "checkpoints and training summaries in the selected model output folder.\n\n" + f"Model folder:\n{output_dir}\n\n" + "Continue and start from scratch?" + ) + choice = QMessageBox.question( + self, + "Start From Scratch?", + message, + QMessageBox.Yes | QMessageBox.No, + QMessageBox.No, + ) + if choice == QMessageBox.Yes: + removed = self._clear_training_run_artifacts(output_dir) + training_config.resume = False + training_config.resume_from_checkpoint = None + self.resume_training.setChecked(False) + self.resume_checkpoint.clear() + log.append("") + log.append("Starting fresh run with current settings.") + for path in removed: + log.append(f"Removed old training artifact: {path}") + LOGGER.warning( + "User chose to discard incompatible resume checkpoint %s and start fresh in %s", + active_resume_path, + output_dir, + ) + if training_config.training_mode == "fine_tune": + base_path = training_config.fine_tune_from_checkpoint + if base_path is None or not Path(base_path).exists(): + log.append("[ERROR] Fine-tune mode requires an existing base checkpoint after reset.") + QMessageBox.warning(self, "Training blocked", "Fine-tune mode still needs a valid base checkpoint.") + return False + compatibility = check_resume_compatibility(Path(base_path), model_config, training_config) + if compatibility.errors: + for line in compatibility.errors: + log.append(f"[ERROR] {line}") + QMessageBox.warning(self, "Training blocked", "The base checkpoint is still incompatible with current settings.") + return False + log.append("[OK] Old run cleared; fine-tune base checkpoint is compatible.") + else: + log.append("[OK] Old run cleared; pretraining will start from random weights.") + self.project_state.setText("Training reset") + self.train_status.setText("Training: starting fresh") + return True + LOGGER.error("Training blocked by preflight checklist.") + for line in info: + LOGGER.info("Training preflight OK: %s", line) + for line in warnings: + LOGGER.warning("Training preflight warning: %s", line) + for line in errors: + LOGGER.error("Training preflight error: %s", line) + self.project_state.setText("Training blocked") + self.train_status.setText("Training: blocked") + QMessageBox.warning(self, "Training blocked", "Fix the checklist errors before starting training.") + return False + existing_artifacts = self._training_run_artifacts(output_dir) + if ( + training_config.training_mode == "pretrain" + and active_resume_path is None + and existing_artifacts + ): + artifact_text = "\n".join(f"- {path.name}" for path in existing_artifacts) + message = ( + "This model folder already contains training artifacts from a previous run.\n\n" + "If you changed architecture or low-memory settings and want a clean start, " + "the old checkpoints should be removed first.\n\n" + f"Model folder:\n{output_dir}\n\n" + f"Artifacts found:\n{artifact_text}\n\n" + "Delete these artifacts and start from scratch?" + ) + choice = QMessageBox.question( + self, + "Clean Previous Run?", + message, + QMessageBox.Yes | QMessageBox.No, + QMessageBox.No, + ) + if choice != QMessageBox.Yes: + self.project_state.setText("Training cancelled") + self.train_status.setText("Training: idle") + log.append("Training cancelled. Previous run artifacts were kept.") + return False + removed = self._clear_training_run_artifacts(output_dir) + training_config.resume = False + training_config.resume_from_checkpoint = None + self.resume_training.setChecked(False) + self.resume_checkpoint.clear() + log.append("") + log.append("Previous run artifacts removed. Training will start from scratch with current settings.") + for path in removed: + log.append(f"Removed old training artifact: {path}") + LOGGER.warning("User cleaned previous training artifacts in %s before starting from scratch.", output_dir) + if warnings: + message = "Training checklist has warnings. Continue anyway?\n\n" + "\n".join(f"- {warning}" for warning in warnings[:8]) + choice = QMessageBox.question(self, "Training warnings", message, QMessageBox.Yes | QMessageBox.No, QMessageBox.No) + if choice != QMessageBox.Yes: + self.project_state.setText("Training cancelled") + self.train_status.setText("Training: idle") + return False + return True + + def start_training(self) -> None: + """Collect training options and start model training.""" + + launch_target = self._training_launch_target_value() + if launch_target == "runpod": + self.launch_runpod_worker_for_current_training() + return + if launch_target == "remote": + self.publish_remote_training_job() + return + self.active_training_log = self.training_log + self.active_training_progress = self.training_progress + self.active_training_final_button_text = "Start Training" + resume_path = Path(self.resume_checkpoint.text()) if self.resume_checkpoint.text().strip() else None + dataset_dir = Path(self.train_data_dir.text()) + vocab_size = self._current_training_vocab_size(dataset_dir) + if vocab_size <= 0: + QMessageBox.warning(self, "Training blocked", "Could not determine tokenizer vocabulary size. Prepare the dataset first.") + return + model_config = self._current_model_config(vocab_size=vocab_size) + training_config = self._current_training_config(resume_path, training_mode="pretrain") + if not self._run_training_preflight(model_config, training_config): + return + self.active_training_output_dir = training_config.output_dir + self._init_telemetry_store(training_config.output_dir) + self.training_log.append("") + self.training_progress.setValue(0) + self.training_epoch_metric.setText("Epoch: -") + self.training_step_metric.setText("Step: -") + self.training_loss_metric.setText("Train loss: -") + self.training_val_metric.setText("Val loss: -") + self.training_health_metric.setText("Health: -") + self.training_health_points = [] + self.training_lr_metric.setText("LR: -") + self.training_speed_metric.setText("Speed: -") + self.training_grad_metric.setText("Grad: -") + self.training_vram_metric.setText("VRAM: -") + self.training_eta_metric.setText("ETA: -") + self.loss_chart.clear() + self.optimization_chart.clear() + self.stability_chart.clear() + self.throughput_chart.clear() + self.memory_chart.clear() + self.live_prediction_chart.update_distribution(0, None) + self.live_attention_chart.update_heatmap(0, None) + self.live_activation_chart.update_histogram(0, None) + self.live_gradient_chart.update_flow(self.n_layer.value(), None, 0) + self.live_progress.setValue(0) + self.live_epoch_metric.setText("Epoch: -") + self.live_step_metric.setText("Step: -") + self.live_tokens_metric.setText("Tokens/sec: -") + self.live_loss_metric.setText("Loss: -") + self.live_lr_metric.setText("LR: -") + self.live_data_metric.setText("Data: -") + self.live_sample_text.setText("Training text: -") + self.live_flow.set_state(self.n_layer.value(), self.n_head.value(), 0, None) + self._set_meter(self.live_cpu_bar, "CPU", self._system_cpu_value()) + self._set_meter(self.live_gpu_bar, "GPU memory", None) + self._set_meter(self.live_vram_bar, "VRAM reserved", None) + self._set_meter(self.live_ram_bar, "System RAM", self._system_ram_value()) + self.live_worker_status.setText(f"CPU workers: {self.data_loader_workers.value()}") + self.training_log.append("Training started...") + self.project_state.setText("Training") + self.train_status.setText("Training: running") + self._run_task( + run_training_job, + (dataset_dir, model_config, training_config), + self._training_finished, + self.training_log, + self.training_progress, + with_progress=True, + button=self.train_button, + stop_button=self.stop_training_button, + busy_text="Training", + task_kind="training", + ) + + diff --git a/interface/main_window_part17.py b/interface/main_window_part17.py new file mode 100644 index 0000000..8361f28 --- /dev/null +++ b/interface/main_window_part17.py @@ -0,0 +1,436 @@ +from __future__ import annotations + +# MainWindow implementation mixin. Runtime names are provided by app.py. +from typing import Any +from . import app as _app + +globals().update({name: value for name, value in vars(_app).items() if not name.startswith("__")}) + +class MainWindowPart17: + def start_fine_tuning(self) -> None: + """Collect fine-tuning options and start adaptation training.""" + + fine_tune_launch = self._fine_tune_launch_target_value() + if fine_tune_launch in {"remote", "runpod"}: + stage_ok, stage_message = self._fine_tune_dataset_stage_status() + self.refresh_fine_tune_workflow() + if not stage_ok: + self.fine_tune_log.append(stage_message) + QMessageBox.warning(self, "Fine-tune blocked", stage_message) + return + if fine_tune_launch == "runpod": + self.launch_runpod_worker_for_current_training(training_mode="fine_tune", stage=self._training_stage_value()) + self.fine_tune_log.append("RunPod fine-tune job launched. Watch Job Manager for worker assignment and progress.") + else: + self.publish_remote_training_job(training_mode="fine_tune", stage=self._training_stage_value()) + self.fine_tune_log.append("Remote fine-tune job queued. Watch Job Manager for worker assignment and progress.") + return + self.active_training_log = self.fine_tune_log + self.active_training_progress = self.fine_tune_progress + self.active_training_final_button_text = "Start Fine-Tune" + stage_ok, stage_message = self._fine_tune_dataset_stage_status() + self.refresh_fine_tune_workflow() + if not stage_ok: + self.fine_tune_log.append(stage_message) + QMessageBox.warning(self, "Fine-tune blocked", stage_message) + return + resume_path = Path(self.resume_checkpoint.text()) if self.resume_checkpoint.text().strip() else None + dataset_dir = Path(self.train_data_dir.text()) + vocab_size = self._current_training_vocab_size(dataset_dir) + if vocab_size <= 0: + QMessageBox.warning(self, "Fine-tune blocked", "Could not determine tokenizer vocabulary size. Prepare the fine-tuning dataset first.") + return + model_config = self._current_model_config(vocab_size=vocab_size) + training_config = self._current_training_config(resume_path, training_mode="fine_tune") + if not self._run_training_preflight(model_config, training_config): + return + self.active_training_output_dir = training_config.output_dir + self._prepare_fine_tune_run_folder(training_config) + self._init_telemetry_store(training_config.output_dir) + self.fine_tune_log.append("") + self.fine_tune_progress.setValue(0) + self.training_progress.setValue(0) + self.fine_tune_eta_metric.setText("ETA: -") + self.fine_tune_epoch_metric.setText("Epoch: -") + self.fine_tune_step_metric.setText("Step: -") + self.fine_tune_loss_metric.setText("Train loss: -") + self.fine_tune_val_metric.setText("Val loss: -") + self.fine_tune_lr_metric.setText("LR: -") + self.fine_tune_speed_metric.setText("Speed: -") + self.fine_tune_grad_metric.setText("Grad: -") + self.training_epoch_metric.setText("Epoch: -") + self.training_step_metric.setText("Step: -") + self.training_loss_metric.setText("Train loss: -") + self.training_val_metric.setText("Val loss: -") + self.training_health_metric.setText("Health: -") + self.training_health_points = [] + self.training_lr_metric.setText("LR: -") + self.training_speed_metric.setText("Speed: -") + self.training_grad_metric.setText("Grad: -") + self.training_vram_metric.setText("VRAM: -") + self.training_eta_metric.setText("ETA: -") + self.loss_chart.clear() + self.optimization_chart.clear() + self.stability_chart.clear() + self.throughput_chart.clear() + self.memory_chart.clear() + self.live_progress.setValue(0) + self.live_sample_text.setText("Training text: -") + self.live_flow.set_state(self.n_layer.value(), self.n_head.value(), 0, None) + self.fine_tune_log.append("Fine-tuning started...") + self.project_state.setText("Fine-tuning") + self.train_status.setText("Training: fine-tuning") + self._run_task( + run_fine_tuning_job, + (dataset_dir, model_config, training_config, self._training_stage_value()), + self._training_finished, + self.fine_tune_log, + self.fine_tune_progress, + with_progress=True, + button=self.fine_tune_button, + stop_button=self.stop_fine_tune_button, + busy_text="Fine-tuning", + task_kind="fine_tune", + ) + + @Slot(object) + def _training_finished(self, result: Any) -> None: + """Update UI after training finishes. + + Args: + result: Training result. + """ + + log = self.active_training_log or self.training_log + progress = self.active_training_progress or self.training_progress + progress.setValue(100) + if progress is not self.training_progress: + self.training_progress.setValue(100) + if hasattr(self, "live_progress"): + self.live_progress.setValue(100) + log.append(f"Saved model: {result.checkpoint_path}") + log.append(f"Final train loss: {result.final_train_loss:.4f}") + if result.final_val_loss is not None: + log.append(f"Final validation loss: {result.final_val_loss:.4f}") + training_summary: dict[str, Any] = {} + try: + training_summary = json.loads(Path(result.summary_path).read_text(encoding="utf-8")) + except Exception: + training_summary = {} + best_checkpoint = str(training_summary.get("recommended_checkpoint_path") or "") + best_val_loss = training_summary.get("best_val_loss") + if best_checkpoint: + if best_val_loss is not None: + log.append(f"Recommended checkpoint: {best_checkpoint} (best validation loss {float(best_val_loss):.4f})") + else: + log.append(f"Recommended checkpoint: {best_checkpoint}") + output_dir = self.active_training_output_dir or Path(result.checkpoint_path).parent + stage_key = self.active_task_kind if self.active_task_kind in {"training", "fine_tune"} else "training" + self.export_model_dir.setText(str(output_dir)) + try: + if stage_key != "fine_tune" and Path(output_dir).resolve() == Path(self.model_dir.text()).resolve(): + self.fine_tune_checkpoint.setText(str(result.checkpoint_path)) + except OSError: + pass + if getattr(result, "stopped", False): + self.project_state.setText("Training stopped") + self.train_status.setText("Training: stopped, checkpoint saved") + log.append("Training stopped safely. Resume from this checkpoint or the latest checkpoint.") + else: + self.project_state.setText("Training complete") + self.train_status.setText(f"Training: loss {result.final_train_loss:.4f}") + title = "Fine-tuning complete" if stage_key == "fine_tune" else "Model training complete" + if getattr(result, "stopped", False): + title = "Fine-tuning stopped" if stage_key == "fine_tune" else "Model training stopped" + completion_lines = [ + f"Checkpoint: {result.checkpoint_path}", + f"Summary: {result.summary_path}", + f"Final train loss: {result.final_train_loss:.4f}", + ] + if result.final_val_loss is not None: + completion_lines.append(f"Final validation loss: {result.final_val_loss:.4f}") + if best_checkpoint: + completion_lines.append(f"Recommended checkpoint: {best_checkpoint}") + if best_val_loss is not None: + completion_lines.append(f"Best validation loss: {float(best_val_loss):.4f}") + completion_lines.append(f"Output: {output_dir}") + self._notify_complete(stage_key, title, completion_lines) + self._append_training_history(result) + self._clear_button_busy(self.active_training_final_button_text) + self.active_training_log = None + self.active_training_progress = None + self.active_training_output_dir = None + + def run_benchmark(self) -> None: + """Run benchmark prompts against the current trained model.""" + + prompts = normalize_prompts(self.benchmark_prompts.toPlainText()) + self.benchmark_log.append(f"Running benchmark with {len(prompts)} prompt(s)...") + self.benchmark_progress.setValue(0) + self.project_state.setText("Benchmarking") + self._run_task( + evaluate_checkpoint, + ( + Path(self.model_dir.text()), + prompts, + None, + self.benchmark_tokens.value(), + self.benchmark_temperature.value(), + 50, + self.device.currentText(), + self.benchmark_kv_cache.isChecked(), + ), + self._benchmark_finished, + self.benchmark_log, + self.benchmark_progress, + with_progress=True, + button=self.run_benchmark_button, + stop_button=self.stop_benchmark_button, + busy_text="Benchmarking", + ) + + @Slot(object) + def _benchmark_finished(self, result: Any) -> None: + """Update UI after benchmark prompts finish. + + Args: + result: Benchmark result object. + """ + + self.benchmark_progress.setRange(0, 100) + self.benchmark_progress.setValue(100) + self.benchmark_log.append( + f"Benchmark complete: {result.prompt_count} prompt(s), {result.total_seconds:.2f}s, " + f"{result.total_generated_tokens} generated token(s), {result.tokens_per_second:.2f} tok/s." + ) + self.benchmark_log.append(f"Benchmark saved: {result.output_path}") + self.project_state.setText("Benchmark complete") + self._clear_button_busy("Run Benchmark") + + def toggle_llm_model(self) -> None: + """Load or unload the selected chat model depending on current state.""" + + if self.chat_session is not None: + self.unload_llm_model() + return + self.load_llm_model() + + def load_llm_model(self) -> None: + """Load a selected model backend for chat testing.""" + + backend = self._chat_backend_value() + path_text = self.microgpt_chat_path.text().strip() if backend == "microgpt" else self.gguf_path.text().strip() + if not path_text: + required = "MicroGPT model folder or checkpoint" if backend == "microgpt" else "GGUF model file" + QMessageBox.information(self, "Model required", f"Choose a {required} first.") + return + model_path = Path(path_text) + self.chat_progress.setValue(0) + self._render_chat_markdown("**Loading model...**") + self.chat_stats.setText("Loading model...") + self.project_state.setText("Loading chat model") + self.chat_status.setText("Chat: loading model") + loader = load_microgpt_chat_session if backend == "microgpt" else load_llama_chat_session + args = ( + (model_path, self.device.currentText()) + if backend == "microgpt" + else (model_path, self.llama_context.value(), self.llama_threads.value(), self.llama_gpu_layers.value()) + ) + self._run_task( + loader, + args, + self._llm_loaded, + self.chat_event_log, + self.chat_progress, + button=self.load_llm_button, + busy_text="Loading Model", + task_kind="chat", + ) + + @Slot(object) + def _llm_loaded(self, session: Any) -> None: + """Store a loaded GGUF chat session. + + Args: + session: Loaded ``LlamaChatSession``. + """ + + self.chat_session = session + self._clear_chat_messages() + self.chat_markdown = "" + self._add_chat_message( + "assistant", + f"Loaded model: `{session.model_path.name}`\n\n{session.runtime_summary}\n\nSend a message to begin.", + ) + self.chat_progress.setValue(100) + self.chat_stats.setText(session.runtime_summary) + self.project_state.setText("Chat model loaded") + self.chat_status.setText(f"Chat: {session.runtime_summary}") + self._clear_button_busy("Unload") + self._tip(self.load_llm_button, "Unload the currently loaded model from memory.") + + def unload_llm_model(self) -> None: + """Unload the active chat model and clear chat state.""" + + if self.thread is not None: + QMessageBox.information(self, "Task running", "Please wait for the current task to finish.") + return + if self.chat_session is not None and hasattr(self.chat_session, "reset"): + self.chat_session.reset() + self.chat_session = None + self._clear_chat_messages() + self.chat_markdown = "" + self._add_chat_message("assistant", "Model unloaded.\n\nLoad a model to start testing.") + self.chat_progress.setRange(0, 100) + self.chat_progress.setValue(0) + self.chat_stats.setText("Idle") + self.project_state.setText("Ready") + self.chat_status.setText("Chat: no model loaded") + self.load_llm_button.setText("Load Model") + self._update_chat_backend_controls() + + def send_chat_message(self) -> None: + """Send a prompt to the loaded chat model.""" + + if self.chat_session is None: + QMessageBox.information(self, "Load model", "Load a model before sending a message.") + return + prompt = self.chat_input.toPlainText().strip() + if not prompt: + return + self.pending_user_message = prompt + self.chat_input.clear() + self._add_chat_message("user", prompt, resend_prompt=prompt) + self.chat_stream_reply = "" + self._add_chat_message("assistant", "_Thinking..._", resend_prompt=prompt) + self.chat_progress.setRange(0, 0) + self.chat_stats.setText("Thinking...") + self.project_state.setText("Generating") + self.chat_status.setText("Chat: generating reply") + streamer = stream_microgpt_chat_reply if self._chat_backend_value() == "microgpt" else stream_chat_reply + self._run_task( + streamer, + ( + self.chat_session, + prompt, + self.system_prompt.toPlainText(), + self.chat_max_tokens.value(), + self.chat_temperature.value(), + self.chat_top_p.value(), + self.chat_repeat_penalty.value(), + self.reasoning_effort.currentText(), + self.thinking_enabled.isChecked(), + ), + self._chat_reply_finished, + self.chat_event_log, + self.chat_progress, + with_progress=True, + button=self.send_chat_button, + stop_button=self.stop_chat_button, + busy_text="Thinking", + ) + + @Slot(object) + def _chat_reply_finished(self, reply: Any) -> None: + """Render the model reply. + + Args: + reply: Assistant reply text and metrics. + """ + + result = reply if isinstance(reply, dict) else {"reply": str(reply)} + text = str(result.get("reply", "")).strip() + if text: + self.chat_stream_reply = text + else: + self.chat_stream_reply = self.chat_stream_reply or "_No reply returned._" + self._render_chat_markdown(self.chat_stream_reply) + self.chat_progress.setRange(0, 100) + self.chat_progress.setValue(100) + self._set_chat_stats( + float(result.get("elapsed_seconds", 0.0)), + int(result.get("token_count", 0)), + float(result.get("tokens_per_second", 0.0)), + ) + self.project_state.setText("Ready") + self.chat_status.setText("Chat: ready") + self._clear_button_busy("Send") + + def reset_chat(self) -> None: + """Clear the chat transcript and model conversation memory.""" + + if self.chat_session is not None: + self.chat_session.reset() + self._clear_chat_messages() + self.chat_markdown = "" + self.chat_stream_prefix = "" + self.chat_stream_reply = "" + self._add_chat_message("assistant", "Chat reset.") + self.chat_stats.setText("Idle") + self.chat_status.setText("Chat: ready") + + def _append_chat_markdown(self, role: str, content: str) -> None: + """Append one rendered chat message. + + Args: + role: Display role heading. + content: Markdown content. + """ + + block = f"### {role}\n{content.strip()}\n" + self.chat_markdown = f"{self.chat_markdown.rstrip()}\n\n{block}" if self.chat_markdown else block + self._add_chat_message("user" if role.lower() in {"you", "user"} else "assistant", content) + + def create_bundle(self) -> None: + """Create a portable model export bundle.""" + + self.export_log.append("Creating model bundle...") + self.export_progress.setValue(15) + try: + output = export_project_bundle(Path(self.export_model_dir.text()), Path(self.export_dir.text())) + except Exception as exc: + self.export_log.append(f"Error: {exc}") + self.export_progress.setValue(0) + return + self.export_progress.setValue(100) + self.export_log.append(f"Bundle created: {output}") + self.export_status.setText("Export: bundle created") + + def quantize_model(self) -> None: + """Create a quantized FP16 checkpoint when selected.""" + + mode = self.quant_mode.currentText() + if not mode.startswith("FP16"): + self.export_log.append("This GGUF quantization target is planned. FP16 checkpoint quantization is available now.") + return + checkpoint = Path(self.export_model_dir.text()) / "final_model.pt" + output = Path(self.export_dir.text()) / "final_model_fp16.pt" + self.export_log.append("Creating FP16 checkpoint...") + self.export_progress.setValue(20) + try: + result = quantize_checkpoint(checkpoint, output, mode="fp16") + except Exception as exc: + self.export_log.append(f"Error: {exc}") + self.export_progress.setValue(0) + return + self.export_progress.setValue(100) + self.export_log.append(f"Quantized checkpoint created: {result}") + self.export_status.setText("Export: FP16 checkpoint ready") + + def export_hf_package(self) -> None: + """Create an HF-style MicroGPT package.""" + + self.export_log.append("Creating HF-style MicroGPT package...") + self.export_progress.setValue(20) + try: + result = export_hf_microgpt_package(Path(self.export_model_dir.text())) + except Exception as exc: + self.export_log.append(f"Error: {exc}") + self.export_progress.setValue(0) + return + self.export_progress.setValue(100) + self.export_log.append(f"HF package created: {result}") + self.export_log.append("Note: this package is MicroGPT model_type, not a llama.cpp-supported Llama model.") + self.export_status.setText("Export: HF package ready") + + diff --git a/interface/main_window_part18.py b/interface/main_window_part18.py new file mode 100644 index 0000000..5662263 --- /dev/null +++ b/interface/main_window_part18.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +# MainWindow implementation mixin. Runtime names are provided by app.py. +from typing import Any +from . import app as _app + +globals().update({name: value for name, value in vars(_app).items() if not name.startswith("__")}) + +class MainWindowPart18: + def export_llama_adapter(self) -> None: + """Create a directly loadable Llama-family package when compatible.""" + + self.export_log.append("Creating Llama-compatible adapter package...") + self.export_progress.setValue(20) + try: + result = export_llama_adapter_package(Path(self.export_model_dir.text())) + except Exception as exc: + self.export_log.append(f"Error: {exc}") + self.export_progress.setValue(0) + return + self.export_progress.setValue(100) + self.export_log.append(f"Llama adapter package created: {result}") + self.export_status.setText("Export: Llama adapter ready") + + def convert_hf_to_gguf(self) -> None: + """Convert an HF-compatible model folder to GGUF through llama.cpp.""" + + model_dir_text = self.export_model_dir.text().strip() + llama_dir_text = self.llama_cpp_dir.text().strip() + output_text = self.gguf_output_path.text().strip() + if not model_dir_text: + QMessageBox.warning(self, "GGUF blocked", "Choose the model core folder first.") + return + if not (Path(model_dir_text) / "hf_model").exists(): + QMessageBox.warning( + self, + "GGUF blocked", + "GGUF conversion needs an HF model package first. Use Export HF Package, then convert a llama.cpp-supported model.", + ) + return + if not llama_dir_text: + QMessageBox.warning(self, "GGUF blocked", "Choose your local llama.cpp folder containing convert_hf_to_gguf.py.") + return + if not output_text: + QMessageBox.warning(self, "GGUF blocked", "Choose a GGUF output file path.") + return + self.export_log.append("Starting llama.cpp GGUF conversion...") + self.export_progress.setValue(0) + self._run_task( + export_gguf_with_llama_cpp, + ( + Path(model_dir_text), + Path(llama_dir_text), + Path(output_text), + self.gguf_outtype.currentText(), + ), + self._gguf_conversion_finished, + self.export_log, + self.export_progress, + button=self.gguf_convert_button, + busy_text="Converting GGUF", + ) + + @Slot(object) + def _gguf_conversion_finished(self, result: Any) -> None: + """Update UI after GGUF conversion finishes. + + Args: + result: GGUF output path. + """ + + self.export_progress.setValue(100) + self.export_log.append(f"GGUF created: {result}") + self.gguf_path.setText(str(result)) + self.export_status.setText("Export: GGUF ready") + self._clear_button_busy("Convert HF to GGUF") + + def _apply_preset(self, preset: str) -> None: + """Apply architecture values for a preset. + + Args: + preset: Selected preset name. + """ + + if preset == "Tiny": + self.n_embd.setValue(128) + self.n_head.setValue(4) + self.n_layer.setValue(4) + elif preset == "Small": + self.n_embd.setValue(512) + self.n_head.setValue(8) + self.n_layer.setValue(8) + + + diff --git a/interface/main_window_part2.py b/interface/main_window_part2.py new file mode 100644 index 0000000..bf0c19f --- /dev/null +++ b/interface/main_window_part2.py @@ -0,0 +1,438 @@ +from __future__ import annotations + +# MainWindow implementation mixin. Runtime names are provided by app.py. +from typing import Any +from . import app as _app + +globals().update({name: value for name, value in vars(_app).items() if not name.startswith("__")}) + +class MainWindowPart2: + def _build_training_tab(self) -> QWidget: + """Build the training configuration page. + + Returns: + Training page widget. + """ + + return build_training_tab(self) + + def _build_fine_tuning_tab(self) -> QWidget: + """Build the fine-tuning page. + + Returns: + Fine-tuning page widget. + """ + + return build_fine_tuning_tab(self) + + def _build_live_training_tab(self) -> QWidget: + """Build the live training tracker page. + + Returns: + Live training tracker page widget. + """ + + return build_live_training_tab(self) + + def _build_job_manager_tab(self) -> QWidget: + """Build the distributed job manager page. + + Returns: + Job manager page widget. + """ + + return build_job_manager_tab(self) + + def refresh_job_manager_tab(self) -> None: + """Refresh the job manager dashboard tables.""" + + if not hasattr(self, "job_worker_table"): + return + workers = self.job_manager.list_workers() + jobs = self.job_manager.list_jobs() + heartbeats = self.job_manager.state_store.latest_heartbeats() + worker_rows = [] + for worker in workers: + heartbeat = heartbeats.get(worker.worker_id, {}) + metrics = heartbeat.get("metrics") or {} + active_job = heartbeat.get("active_job_id") or self._active_job_for_worker(worker.worker_id) + capabilities = worker.capabilities or {} + cpu_ram_gpu = ( + f"CPU {capabilities.get('cpu_count', '-')}, " + f"RAM {capabilities.get('system_ram_gb', '-')} GB, " + f"VRAM {capabilities.get('total_vram_gb', '-')} GB" + ) + if metrics: + cpu_ram_gpu = f"{cpu_ram_gpu}, util {metrics.get('gpu_util', metrics.get('gpu_memory_percent', '-'))}" + worker_rows.append( + [ + worker.worker_id, + worker.status.value, + worker.backend.value, + worker.device, + worker.last_heartbeat_at or "-", + active_job or "-", + cpu_ram_gpu, + ", ".join(capabilities.get("labels") or []) or "-", + ] + ) + set_table_rows(self.job_worker_table, worker_rows) + + job_rows = [] + for managed in jobs: + job = managed.spec + metrics = managed.latest_metrics + stage_label = str(job.metadata.get("training_stage") or job.metadata.get("training_mode") or job.training.training_mode) + job_rows.append( + [ + job.job_id, + stage_label, + job.status.value, + managed.assigned_worker_id or "-", + job.runtime.backend.value, + self._metric_pair(metrics.epoch if metrics else None, metrics.total_epochs if metrics else None), + self._metric_pair(metrics.step if metrics else None, metrics.total_steps if metrics else None), + str(job.training.batch_size), + str(job.model.config.layer_count), + self._metric_float(metrics.train_loss if metrics else None), + self._metric_float(metrics.tokens_per_second if metrics else None, suffix=" tok/s"), + managed.updated_at, + ] + ) + set_table_rows(self.job_table, job_rows) + active_count = sum(1 for item in jobs if item.spec.status.value in {"assigned", "running", "paused", "stopping"}) + queued_count = sum(1 for item in jobs if item.spec.status.value == "queued") + self.job_worker_count_label.setText(f"Workers: {len(workers)}") + self.job_active_count_label.setText(f"Active jobs: {active_count}") + self.job_queue_count_label.setText(f"Queued jobs: {queued_count}") + self.job_db_label.setText(f"State DB: {self.job_manager.state_store.db_path}") + self.job_manager_progress.setValue(100) + + def pause_all_managed_jobs(self) -> None: + """Pause all managed jobs.""" + + count = self.job_manager.pause_all_jobs() + self.job_manager_log.append(f"Pause requested for {count} job(s).") + self.refresh_job_manager_tab() + + def resume_all_managed_jobs(self) -> None: + """Resume all paused managed jobs.""" + + count = self.job_manager.resume_all_jobs() + self.job_manager_log.append(f"Resumed {count} job(s).") + self.refresh_job_manager_tab() + + def stop_all_managed_jobs(self) -> None: + """Stop all managed jobs.""" + + count = self.job_manager.stop_all_jobs() + self.job_manager_log.append(f"Stop requested for {count} job(s).") + self.refresh_job_manager_tab() + + def mark_stale_workers_offline(self) -> None: + """Mark stale remote workers offline.""" + + workers = self.job_manager.mark_stale_workers_offline() + if workers: + self.job_manager_log.append(f"Marked offline: {', '.join(workers)}") + else: + self.job_manager_log.append("No stale remote workers found.") + self.refresh_job_manager_tab() + + def start_coordinator_server(self) -> None: + """Start the coordinator API used by remote workers.""" + + if self.coordinator_server is not None: + self.job_manager_log.append("Coordinator API is already running.") + return + host = self.coordinator_host.text().strip() or "0.0.0.0" + port = self.coordinator_port.value() + artifact_root = Path(self.coordinator_artifact_root.text().strip()).expanduser() + artifact_root.mkdir(parents=True, exist_ok=True) + try: + self.coordinator_server = CoordinatorApiServer( + manager=self.job_manager, + host=host, + port=port, + artifact_root=artifact_root, + ) + self.coordinator_thread = Thread(target=self.coordinator_server.serve_forever, daemon=True) + self.coordinator_thread.start() + except Exception as exc: + self.coordinator_server = None + self.coordinator_thread = None + QMessageBox.warning(self, "Coordinator failed", f"Could not start coordinator API:\n{exc}") + return + public_url = self.coordinator_public_url.text().strip() or f"http://127.0.0.1:{port}" + self.coordinator_public_url.setText(public_url.rstrip("/")) + self.coordinator_status_label.setText(f"Coordinator: running at {public_url.rstrip('/')}") + self.coordinator_start_button.setEnabled(False) + self.coordinator_stop_button.setEnabled(True) + self.project_state.setText("Coordinator running") + self.job_manager_log.append(f"Coordinator API started on {host}:{port}.") + self.job_manager_log.append(f"Artifact sync root: {artifact_root}") + + def stop_coordinator_server(self) -> None: + """Stop the coordinator API.""" + + if self.coordinator_server is None: + return + self.coordinator_server.shutdown() + if self.coordinator_thread is not None: + self.coordinator_thread.join(timeout=3) + self.coordinator_server = None + self.coordinator_thread = None + self.coordinator_status_label.setText("Coordinator: stopped") + self.coordinator_start_button.setEnabled(True) + self.coordinator_stop_button.setEnabled(False) + self.project_state.setText("Coordinator stopped") + self.job_manager_log.append("Coordinator API stopped.") + + def _runpod_config_path(self) -> Path: + """Return the active RunPod config path. + + Returns: + Project-local RunPod config path when a project is open. + """ + + project_dir = self.current_project_file.parent if self.current_project_file is not None else None + return default_runpod_config_path(project_dir) + + def load_runpod_settings(self) -> None: + """Load RunPod settings into the Job Manager UI.""" + + if not hasattr(self, "runpod_api_key"): + return + config_path = self._runpod_config_path() + try: + config = load_runpod_config(config_path) + except Exception as exc: + LOGGER.error("Could not load RunPod config: %s", exc) + self.runpod_status_label.setText(f"RunPod config error: {exc}") + return + self.runpod_api_key.setText(config.api_key) + self._set_combo_text(self.runpod_gpu_type, config.gpu_type_id) + self._set_combo_text(self.runpod_cloud_type, config.cloud_type) + self.runpod_image.setText(config.image_name) + self.runpod_container_disk.setValue(config.container_disk_gb) + self.runpod_volume_gb.setValue(config.volume_gb) + self.runpod_min_ram.setValue(config.min_ram_per_gpu) + self.runpod_min_vcpu.setValue(config.min_vcpu_per_gpu) + self.runpod_spot.setChecked(config.interruptible) + self.runpod_auto_terminate.setChecked(config.auto_terminate) + status = "configured" if config.api_key.strip() else "API key needed" + self.runpod_status_label.setText(f"RunPod: {status} ({config_path})") + + def save_runpod_settings(self) -> None: + """Save RunPod settings from the Job Manager UI.""" + + config = self._runpod_config_from_ui() + config_path = self._runpod_config_path() + save_runpod_config(config_path, config) + self.runpod_status_label.setText(f"RunPod settings saved: {config_path}") + self.job_manager_log.append(f"RunPod settings saved: {config_path}") + LOGGER.info("RunPod settings saved: %s", config_path) + + def _runpod_config_from_ui(self) -> RunPodConfig: + """Collect RunPod settings from the UI. + + Returns: + RunPod configuration. + """ + + return RunPodConfig( + api_key=self.runpod_api_key.text().strip(), + image_name=self.runpod_image.text().strip(), + gpu_type_id=self.runpod_gpu_type.currentText().strip(), + gpu_count=1, + cloud_type=self.runpod_cloud_type.currentText().strip(), + interruptible=self.runpod_spot.isChecked(), + container_disk_gb=self.runpod_container_disk.value(), + volume_gb=self.runpod_volume_gb.value(), + min_vcpu_per_gpu=self.runpod_min_vcpu.value(), + min_ram_per_gpu=self.runpod_min_ram.value(), + auto_terminate=self.runpod_auto_terminate.isChecked(), + worker_labels="runpod,gpu", + ) + + def launch_runpod_worker_for_current_training(self, training_mode: str = "pretrain", stage: str = "base") -> None: + """Publish the current training job and launch a RunPod worker Pod. + + Args: + training_mode: Training mode for the queued job. + stage: Dataset/training stage label. + """ + + if isinstance(training_mode, bool): + training_mode = "pretrain" + stage = "base" + try: + config = self._runpod_config_from_ui() + save_runpod_config(self._runpod_config_path(), config) + coordinator_url = self.coordinator_public_url.text().strip().rstrip("/") + if not public_url_is_cloud_reachable(coordinator_url): + raise ValueError( + "RunPod needs a public Worker URL. Start a tunnel or set Worker URL to a public address, " + "not localhost/127.0.0.1." + ) + if self.coordinator_server is None: + self.start_coordinator_server() + if self.coordinator_server is None: + return + job, bundle_path = self._publish_remote_training_job_spec( + training_mode=training_mode, + stage=stage, + backend_label="runpod", + ) + artifact_root = Path(self.coordinator_artifact_root.text().strip()).expanduser() + bootstrap_path = create_runpod_worker_bundle(Path(__file__).resolve().parents[2], artifact_root) + bootstrap_url = f"{coordinator_url}/artifacts/{bootstrap_path.name}" + worker_id = f"runpod-{job.job_id}" + pod_name = f"micro-llm-{self._safe_project_name(self.search_box.text().strip() or 'project')}-{job.job_id[-8:]}" + result = RunPodClient(config.api_key).create_worker_pod( + config=config, + pod_name=pod_name, + worker_id=worker_id, + coordinator_url=coordinator_url, + bootstrap_url=bootstrap_url, + ) + managed = self.job_manager.get_job(job.job_id) + managed.spec.metadata["runpod_pod_id"] = result.pod_id + managed.spec.metadata["runpod_worker_id"] = result.worker_id + managed.spec.metadata["runpod_cost_per_hour"] = result.cost_per_hour + self.job_manager._persist_job(job.job_id) + except Exception as exc: + LOGGER.exception("RunPod launch failed") + QMessageBox.warning(self, "RunPod launch failed", str(exc)) + if hasattr(self, "runpod_status_label"): + self.runpod_status_label.setText(f"RunPod launch failed: {exc}") + return + self.runpod_status_label.setText( + f"RunPod pod {result.pod_id} launched for {job.job_id} ({result.gpu_name}, {result.cost_per_hour}/hr)" + ) + self.job_manager_log.append(f"RunPod pod launched: {result.pod_id}") + self.job_manager_log.append(f"RunPod worker: {result.worker_id}") + self.job_manager_log.append(f"RunPod GPU: {result.gpu_name}, cost/hr: {result.cost_per_hour}") + self.job_manager_log.append(f"Worker bootstrap: {result.bootstrap_url}") + self.project_state.setText("RunPod worker launched") + self.refresh_job_manager_tab() + + def publish_remote_training_job(self, training_mode: str = "pretrain", stage: str = "base") -> None: + """Bundle the current training setup and queue it for remote workers. + + Args: + training_mode: Trainer mode to publish, either ``pretrain`` or ``fine_tune``. + stage: Higher-level stage label for job manager display. + """ + + if isinstance(training_mode, bool): + training_mode = "pretrain" + stage = "base" + if self.coordinator_server is None: + self.start_coordinator_server() + if self.coordinator_server is None: + return + try: + job, bundle_path = self._publish_remote_training_job_spec(training_mode=training_mode, stage=stage) + except Exception as exc: + QMessageBox.warning(self, "Publish failed", f"Could not publish remote job:\n{exc}") + return + self.job_manager_log.append(f"Published remote job: {job.job_id}") + self.job_manager_log.append(f"Input bundle: {bundle_path}") + self.job_manager_log.append(f"Worker download URL: {job.metadata.get('artifact_bundle_url')}") + self.project_state.setText("Remote job queued") + self.refresh_job_manager_tab() + + def _publish_remote_training_job_spec( + self, + training_mode: str = "pretrain", + stage: str = "base", + backend_label: str = "remote", + ) -> tuple[TrainingJobSpec, Path]: + """Bundle and queue the current remote training job. + + Args: + training_mode: Trainer mode to publish. + stage: Higher-level stage label. + backend_label: Human-readable backend label stored in metadata. + + Returns: + Queued job and bundle path. + """ + + job = self._current_remote_training_job(training_mode=training_mode, stage=stage) + job.metadata["launch_backend"] = backend_label + artifact_root = Path(self.coordinator_artifact_root.text().strip()).expanduser() + base_url = f"{self.coordinator_public_url.text().strip().rstrip('/')}/artifacts" + bundle_path = create_job_artifact_bundle(job, artifact_root=artifact_root, base_url=base_url) + self.job_manager.submit(job) + return job, bundle_path + + def _current_remote_training_job(self, training_mode: str = "pretrain", stage: str = "base") -> TrainingJobSpec: + """Build a remote-worker job from current training controls. + + Args: + training_mode: Trainer mode to publish. + stage: Higher-level stage label for job manager display. + + Returns: + Complete training job spec ready to bundle and queue. + + Raises: + FileNotFoundError: If the prepared dataset is missing. + ValueError: If model or training options are invalid. + """ + + dataset_dir = Path(self.train_data_dir.text().strip()) + if not dataset_dir.exists(): + raise FileNotFoundError(f"Prepared dataset folder does not exist: {dataset_dir}") + if not self._dataset_artifacts_exist(dataset_dir): + raise FileNotFoundError( + "Prepared dataset is missing tokenizer or token files. " + "Expected tokenizer.json plus train/val tokens in .npy or .json." + ) + vocab_size = self._current_training_vocab_size(dataset_dir) + if vocab_size <= 0: + raise ValueError("Could not determine tokenizer vocabulary size from the prepared dataset.") + resume_path = Path(self.resume_checkpoint.text()) if self.resume_checkpoint.text().strip() else None + if resume_path is None and self.resume_training.isChecked(): + resume_path = latest_checkpoint(self._training_output_dir_for_mode(training_mode) / "checkpoints") + model_config = self._current_model_config(vocab_size=vocab_size) + training_config = self._current_training_config(resume_path, training_mode=training_mode) + model_config.validate() + training_config.validate() + job = TrainingJobSpec.local( + dataset_dir, + model_config, + training_config, + metadata={ + "project_name": self.search_box.text().strip(), + "submitted_from": "desktop_ui", + "coordinator_url": self.coordinator_public_url.text().strip().rstrip("/"), + "training_mode": training_mode, + "training_stage": stage, + }, + ) + job.runtime = RuntimeSpec( + backend=BackendKind.REMOTE_CLIENT, + device=training_config.device, + tags=[training_config.device, "remote"], + ) + return job + + def _active_job_for_worker(self, worker_id: str) -> str: + """Return the active job ID for a worker. + + Args: + worker_id: Worker identifier. + + Returns: + Active job ID or empty string. + """ + + for managed in self.job_manager.list_jobs(): + if managed.assigned_worker_id == worker_id and managed.spec.status.value in {"assigned", "running", "paused", "stopping"}: + return managed.spec.job_id + return "" + + diff --git a/interface/main_window_part3.py b/interface/main_window_part3.py new file mode 100644 index 0000000..b822274 --- /dev/null +++ b/interface/main_window_part3.py @@ -0,0 +1,334 @@ +from __future__ import annotations + +# MainWindow implementation mixin. Runtime names are provided by app.py. +from typing import Any +from . import app as _app + +globals().update({name: value for name, value in vars(_app).items() if not name.startswith("__")}) + +class MainWindowPart3: + @staticmethod + def _metric_pair(value: Optional[int], total: Optional[int]) -> str: + """Format a metric pair. + + Args: + value: Current value. + total: Total value. + + Returns: + Display text. + """ + + if value is None: + return "-" + if total is None: + return str(value) + return f"{value}/{total}" + + @staticmethod + def _metric_float(value: Optional[float], suffix: str = "") -> str: + """Format a floating-point metric. + + Args: + value: Metric value. + suffix: Optional suffix. + + Returns: + Display text. + """ + + if value is None: + return "-" + return f"{value:.4g}{suffix}" + + def _init_telemetry_store(self, model_dir: Path) -> None: + """Create or reset the SQLite telemetry store for a training run. + + Args: + model_dir: Model output directory. + """ + + self.telemetry_db_path = initialize_store(model_dir) + self.telemetry_run_id = datetime.now().strftime("%Y%m%d_%H%M%S") + self.telemetry_latest_id = 0 + self.telemetry_latest_index = 0 + self.live_time_slider.setRange(0, 0) + self.live_time_slider.setValue(0) + self.live_timeline_label.setText("Timeline: live") + self.live_scrub_active = False + + def _record_live_metric(self, event: dict[str, Any]) -> None: + """Persist one live training metric event to SQLite. + + Args: + event: Training progress event. + """ + + if self.telemetry_db_path is None or not self.telemetry_run_id or event.get("step") is None: + return + self.telemetry_latest_id = insert_metric(self.telemetry_db_path, self.telemetry_run_id, event) + self.telemetry_latest_index += 1 + self.live_time_slider.blockSignals(True) + self.live_time_slider.setRange(0, self.telemetry_latest_index) + if not self.live_scrub_active: + self.live_time_slider.setValue(self.telemetry_latest_index) + self.live_timeline_label.setText("Timeline: live") + self.live_time_slider.blockSignals(False) + + def _load_existing_telemetry(self, model_dir: Path) -> None: + """Load the latest saved telemetry run for an opened project. + + Args: + model_dir: Model output directory that may contain ``training_telemetry.sqlite``. + """ + + db_path = telemetry_db_path(model_dir) + self.telemetry_db_path = db_path if db_path.exists() else None + self.telemetry_run_id = "" + self.telemetry_latest_id = 0 + self.telemetry_latest_index = 0 + self.live_scrub_active = False + self.live_time_slider.blockSignals(True) + self.live_time_slider.setRange(0, 0) + self.live_time_slider.setValue(0) + self.live_time_slider.blockSignals(False) + self.live_timeline_label.setText("Timeline: no saved telemetry") + self.live_sample_text.setText("Training text: -") + if self.telemetry_db_path is None: + return + try: + run_row = latest_run(self.telemetry_db_path) + if run_row is None: + self.live_timeline_label.setText("Timeline: no samples") + return + self.telemetry_run_id = str(run_row["run_id"]) + self.telemetry_latest_index = int(run_row["sample_count"] or 0) + self.telemetry_latest_id = int(run_row["latest_id"] or 0) + except sqlite3.Error as exc: + self.live_timeline_label.setText("Timeline: could not load") + self.training_log.append(f"Telemetry load warning: {exc}") + return + self.live_time_slider.blockSignals(True) + self.live_time_slider.setRange(0, self.telemetry_latest_index) + self.live_time_slider.setValue(self.telemetry_latest_index) + self.live_time_slider.blockSignals(False) + if self.telemetry_latest_index: + rows = self._timeline_rows_until(self.telemetry_latest_index) + if rows: + self._apply_timeline_rows(rows) + + def _timeline_rows_until(self, sample_index: int) -> list[sqlite3.Row]: + """Load telemetry rows up to a selected sample index. + + Args: + sample_index: Maximum number of samples to load for the active run. + + Returns: + Ordered telemetry rows for the active run. + """ + + if self.telemetry_db_path is None or not self.telemetry_run_id or sample_index <= 0: + return [] + return rows_until(self.telemetry_db_path, self.telemetry_run_id, sample_index) + + def _begin_live_scrub(self) -> None: + """Pause live auto-follow while the timeline slider is being dragged.""" + + self.live_scrub_active = True + + def _end_live_scrub(self) -> None: + """Apply the selected timeline snapshot after slider drag.""" + + self._scrub_live_timeline(self.live_time_slider.value()) + + def _jump_live_timeline_to_latest(self) -> None: + """Return timeline display to the latest live point.""" + + self.live_scrub_active = False + self.live_time_slider.setValue(self.telemetry_latest_index) + self._scrub_live_timeline(self.telemetry_latest_index) + self.live_timeline_label.setText("Timeline: live") + + def _scrub_live_timeline(self, sample_index: int) -> None: + """Replay charts and live visual widgets to a selected telemetry point. + + Args: + sample_index: Timeline sample selected by the slider. + """ + + rows = self._timeline_rows_until(sample_index) + if not rows: + return + self._apply_timeline_rows(rows) + + def _apply_timeline_rows(self, rows: list[sqlite3.Row]) -> None: + """Apply historical telemetry rows to charts and live widgets. + + Args: + rows: Ordered SQLite telemetry rows. + """ + + def series(name: str) -> list[tuple[int, float]]: + return [(int(row["step"]), float(row[name])) for row in rows if row[name] is not None] + + latest = rows[-1] + self.loss_chart.set_points(series("train_loss"), series("val_loss")) + self.optimization_chart.set_points(series("learning_rate"), series("grad_norm")) + self.stability_chart.set_points(series("weight_norm"), series("update_ratio")) + self.throughput_chart.set_points(series("tokens_per_second"), series("samples_per_second")) + self.memory_chart.set_points(series("vram_allocated_gb"), series("vram_reserved_gb")) + snapshot = {key: latest[key] for key in latest.keys()} + sample_text = str(snapshot.get("sample_text") or "").strip() + if sample_text: + self.live_sample_text.setText(f"Training text: {self._compact_preview_text(sample_text, 220)}") + else: + self.live_sample_text.setText("Training text: -") + self._update_live_training_metrics( + int(latest["step"]), + snapshot, + snapshot.get("train_loss"), + snapshot.get("learning_rate"), + snapshot.get("grad_norm"), + snapshot.get("update_ratio"), + snapshot.get("tokens_per_second"), + snapshot.get("samples_per_second"), + snapshot.get("vram_allocated_gb"), + snapshot.get("vram_reserved_gb"), + snapshot.get("gpu_memory_percent"), + snapshot.get("system_cpu_percent"), + snapshot.get("system_ram_percent"), + snapshot.get("data_loader_workers"), + ) + timestamp = datetime.fromtimestamp(float(latest["recorded_at"])).strftime("%H:%M:%S") + self.live_timeline_label.setText(f"Timeline: step {int(latest['step']):,} @ {timestamp}") + + @staticmethod + def _compact_preview_text(text: str, limit: int = 220) -> str: + """Normalize a training preview into a compact single line. + + Args: + text: Raw decoded preview text. + limit: Maximum number of displayed characters. + + Returns: + Single-line text preview. + """ + + compact = re.sub(r"\s+", " ", text).strip() + if len(compact) <= limit: + return compact + return compact[: max(0, limit - 3)].rstrip() + "..." + + def _build_export_tab(self) -> QWidget: + """Build the export page. + + Returns: + Export page widget. + """ + + return build_export_tab(self) + + def _build_benchmark_tab(self) -> QWidget: + """Build the benchmark prompt page. + + Returns: + Benchmark page widget. + """ + + return build_benchmark_tab(self) + + def _build_chat_tab(self) -> QWidget: + """Build the model test chat page. + + Returns: + Chat page widget. + """ + + return build_chat_tab(self) + + def _panel(self) -> QWidget: + """Create a base page panel. + + Returns: + Panel widget. + """ + + page = QWidget() + page.setObjectName("Panel") + return page + + def _page_title(self, text: str) -> QLabel: + """Create a page title label. + + Args: + text: Title text. + + Returns: + Label configured as a page title. + """ + + label = QLabel(text) + label.setObjectName("PageTitle") + return label + + def _metric_chip(self, text: str, tooltip: str) -> QLabel: + """Create a compact metric display label. + + Args: + text: Initial metric text. + tooltip: User-facing explanation. + + Returns: + Configured metric label. + """ + + label = QLabel(text) + label.setObjectName("MetricChip") + label.setMinimumWidth(150) + label.setMinimumHeight(28) + label.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed) + self._tip(label, tooltip) + return label + + def _hardware_meter(self, name: str) -> QProgressBar: + """Create a slider-like hardware utilization meter. + + Args: + name: Display name for the meter. + + Returns: + Configured progress bar. + """ + + meter = QProgressBar() + meter.setObjectName("HardwareMeter") + meter.setRange(0, 100) + meter.setValue(0) + meter.setTextVisible(False) + meter.setFixedHeight(8) + meter.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed) + self._tip(meter, f"Live {name} utilization.") + return meter + + def _set_meter(self, meter: QProgressBar, name: str, value: Optional[float]) -> None: + """Update a hardware utilization meter. + + Args: + meter: Meter to update. + name: Display name for the meter. + value: Utilization percentage. + """ + + if value is None: + meter.setValue(0) + label = self.hardware_meter_labels.get(id(meter)) + if label is not None: + label.setText(f"{name}: -") + return + bounded = max(0.0, min(100.0, float(value))) + meter.setValue(int(round(bounded))) + label = self.hardware_meter_labels.get(id(meter)) + if label is not None: + label.setText(f"{name}: {bounded:.1f}%") + + diff --git a/interface/main_window_part4.py b/interface/main_window_part4.py new file mode 100644 index 0000000..abcd56c --- /dev/null +++ b/interface/main_window_part4.py @@ -0,0 +1,438 @@ +from __future__ import annotations + +# MainWindow implementation mixin. Runtime names are provided by app.py. +from typing import Any +from . import app as _app + +globals().update({name: value for name, value in vars(_app).items() if not name.startswith("__")}) + +class MainWindowPart4: + def _update_dataset_quality_report(self, summary: dict[str, Any]) -> None: + """Update dataset quality chips from a summary dictionary. + + Args: + summary: Dataset summary fields. + """ + + document_count = int(summary.get("document_count", 0) or 0) + token_count = int(summary.get("token_count", 0) or 0) + train_window_count = int(summary.get("train_window_count", 0) or 0) + val_window_count = int(summary.get("val_window_count", 0) or 0) + character_count = int(summary.get("character_count", 0) or 0) + vocab_size = int(summary.get("tokenizer_vocab_size", summary.get("vocab_size", 0)) or 0) + code_count = int(summary.get("code_sample_count", 0) or 0) + prose_count = int(summary.get("prose_sample_count", 0) or 0) + conversation_count = int(summary.get("conversation_sample_count", 0) or 0) + cached_count = int(summary.get("cached_file_count", 0) or 0) + processed_count = int(summary.get("processed_file_count", 0) or 0) + skipped_count = int(summary.get("skipped_file_count", 0) or 0) + failed_count = int(summary.get("failed_file_count", 0) or 0) + warning = str(summary.get("warning") or "none") + sequence_stats = summary.get("sequence_token_stats", {}) or {} + quality_score = float(summary.get("quality_score", 0.0) or 0.0) + quality_stars = float(summary.get("quality_stars", 0.0) or 0.0) + quality_label = str(summary.get("quality_label") or "") + corpus_block_count = int(summary.get("corpus_block_count", 0) or 0) + unique_block_count = int(summary.get("unique_block_count", 0) or 0) + duplicate_block_count = int(summary.get("duplicate_block_count", 0) or 0) + duplicate_block_ratio = float(summary.get("duplicate_block_ratio", 0.0) or 0.0) + if not quality_score and (token_count or train_window_count or vocab_size): + quality_score, quality_stars, quality_label = self._estimate_dataset_rating( + token_count, + vocab_size, + train_window_count, + val_window_count, + document_count, + code_count, + prose_count, + conversation_count, + skipped_count, + failed_count, + warning, + sequence_stats, + ) + self.dataset_quality_samples.setText(f"Documents: {document_count:,}") + self.dataset_quality_tokens.setText(f"Tokens: {token_count:,}") + if train_window_count or val_window_count: + self.dataset_quality_windows.setText(f"Windows: {train_window_count:,}/{val_window_count:,}") + else: + self.dataset_quality_windows.setText("Windows: -") + self.dataset_quality_vocab.setText(f"Vocab: {vocab_size:,}" if vocab_size else "Vocab: -") + self.dataset_quality_rating.setText( + f"Rating: {self._star_text(quality_stars)} {quality_stars:.1f}/5" + if quality_stars + else "Rating: -" + ) + self.dataset_quality_code.setText(f"Code/prose/chat: {code_count:,}/{prose_count:,}/{conversation_count:,}") + self.dataset_quality_balance.setText("Balance: prepared") + self.dataset_quality_readiness.setText("Readiness: preview needed") + self.dataset_quality_cache.setText(f"Files: {processed_count:,} ok, {cached_count:,} cached, {skipped_count:,} skipped, {failed_count:,} failed") + if corpus_block_count: + self.dataset_quality_duplicates.setText(f"Duplicates: {duplicate_block_ratio * 100:.1f}%") + self._tip( + self.dataset_quality_duplicates, + ( + f"{duplicate_block_count:,} repeated blocks out of {corpus_block_count:,}; " + f"{unique_block_count:,} unique blocks." + ), + ) + else: + self.dataset_quality_duplicates.setText("Duplicates: -") + self.dataset_quality_warning.setText(f"Warnings: {warning}") + self._tip(self.dataset_quality_samples, f"{character_count:,} source characters across prepared documents.") + if quality_stars: + self._tip( + self.dataset_quality_rating, + f"{quality_label or 'Rated'} dataset: {quality_score:.1f}/100. Higher scores usually mean more usable tokens, richer vocabulary, more windows, and fewer extraction issues.", + ) + self._tip( + self.dataset_quality_windows, + f"{train_window_count:,} training and {val_window_count:,} validation sliding windows.", + ) + self._update_dataset_stat_charts(summary, code_count, prose_count, conversation_count, sequence_stats) + if hasattr(self, "dataset_advisor") and (train_window_count or val_window_count): + advice = [ + "Documents are source items. Windows are the actual context slices used by training.", + f"This dataset can provide about {train_window_count:,} training windows and {val_window_count:,} validation windows.", + ] + if sequence_stats: + advice.append( + "Approx token distribution per source: " + f"min {int(sequence_stats.get('min', 0) or 0):,}, " + f"avg {float(sequence_stats.get('average', 0.0) or 0.0):,.0f}, " + f"median {float(sequence_stats.get('median', 0.0) or 0.0):,.0f}, " + f"max {int(sequence_stats.get('max', 0) or 0):,}." + ) + if document_count < 100 and train_window_count >= 10_000: + advice.append( + "A low document count can still be useful when each document is long, because the trainer samples many overlapping windows." + ) + if train_window_count < 1_000: + advice.append("Add more text or lower context length if training looks repetitive.") + if corpus_block_count: + advice.append( + f"Block diversity: {unique_block_count:,}/{corpus_block_count:,} unique blocks " + f"({duplicate_block_ratio * 100:.1f}% repeated)." + ) + if quality_stars: + advice.append(f"Dataset rating: {quality_stars:.1f}/5 stars ({quality_label or 'rated'}, score {quality_score:.1f}/100).") + for reason in list(summary.get("quality_reasons", []) or [])[:4]: + advice.append(f"- {reason}") + self.dataset_advisor.setPlainText("\n".join(advice)) + + def _star_text(self, stars: float) -> str: + """Return a compact five-star display string. + + Args: + stars: Rating from zero to five. + + Returns: + Unicode star display with rounded whole stars. + """ + + whole = max(0, min(5, int(round(float(stars))))) + return "*" * whole + "-" * (5 - whole) + + def _estimate_dataset_rating( + self, + token_count: int, + vocab_size: int, + train_window_count: int, + val_window_count: int, + document_count: int, + code_count: int, + prose_count: int, + conversation_count: int, + skipped_count: int, + failed_count: int, + warning: str, + sequence_stats: dict[str, Any], + ) -> tuple[float, float, str]: + """Estimate a dataset rating for older summaries that lack saved quality fields. + + Args: + token_count: Total prepared token count. + vocab_size: Tokenizer vocabulary size. + train_window_count: Number of training windows. + val_window_count: Number of validation windows. + document_count: Number of source documents. + code_count: Code sample count. + prose_count: Prose sample count. + conversation_count: Conversation/instruction sample count. + skipped_count: Skipped source file count. + failed_count: Failed source file count. + warning: Dataset warning text. + sequence_stats: Approximate source sequence statistics. + + Returns: + Score, stars, and label. + """ + + def ratio(value: float, target: float) -> float: + return max(0.0, min(1.0, float(value) / float(target))) if target > 0 else 0.0 + + families = sum(1 for count in (code_count, prose_count, conversation_count) if count > 0) + score = ( + 30.0 * ratio(token_count, 1_000_000) + + 20.0 * ratio(train_window_count, 50_000) + + 18.0 * ratio(vocab_size, 8_000) + + 12.0 * ratio(document_count, 1_000) + + 8.0 * ratio(val_window_count, 2_000) + + 7.0 * ratio(families, 3) + + 5.0 * ratio(float(sequence_stats.get("average", 0.0) or 0.0), 256) + ) + score -= min(20.0, failed_count * 3.0 + skipped_count * 0.5) + if warning and warning != "none": + score -= 5.0 + score = max(0.0, min(100.0, score)) + stars = round(score / 20.0 * 2.0) / 2.0 + if score >= 85: + label = "Excellent" + elif score >= 70: + label = "Good" + elif score >= 50: + label = "Usable" + elif score >= 30: + label = "Weak" + else: + label = "Very weak" + return score, stars, label + + def _update_dataset_stat_charts( + self, + summary: dict[str, Any], + code_count: int, + prose_count: int, + conversation_count: int, + sequence_stats: dict[str, Any], + ) -> None: + """Update dataset statistics charts. + + Args: + summary: Dataset summary fields. + code_count: Number of code samples. + prose_count: Number of prose samples. + conversation_count: Number of conversation or instruction samples. + sequence_stats: Approximate token distribution statistics. + """ + + if not hasattr(self, "dataset_mix_chart"): + return + mixture_report = summary.get("mixture_report", {}) or {} + family_rows = list((mixture_report.get("families", {}) or {}).values()) + labels: list[str] = [] + values: list[float] = [] + for row in family_rows: + actual = float(row.get("actual_percent", 0.0) or 0.0) + selected = int(row.get("selected_documents", 0) or 0) + if actual > 0.0 or selected > 0: + labels.append(str(row.get("label") or "source")) + values.append(actual) + if not labels: + total = max(code_count + prose_count + conversation_count, 1) + labels = ["Code", "Prose", "Conversation"] + values = [ + code_count * 100.0 / total, + prose_count * 100.0 / total, + conversation_count * 100.0 / total, + ] + self.dataset_mix_chart.set_values(labels, values, "%") + if sequence_stats: + self.dataset_sequence_chart.set_values( + ["Min", "Average", "Median", "Max"], + [ + float(sequence_stats.get("min", 0) or 0), + float(sequence_stats.get("average", 0.0) or 0.0), + float(sequence_stats.get("median", 0.0) or 0.0), + float(sequence_stats.get("max", 0) or 0), + ], + ) + else: + self.dataset_sequence_chart.clear() + + def _reset_dataset_quality_report(self) -> None: + """Reset dataset quality chips to their empty state.""" + + self.dataset_quality_samples.setText("Documents: -") + self.dataset_quality_tokens.setText("Tokens: -") + self.dataset_quality_windows.setText("Windows: -") + self.dataset_quality_vocab.setText("Vocab: -") + self.dataset_quality_rating.setText("Rating: -") + self.dataset_quality_code.setText("Code/prose: -") + self.dataset_quality_balance.setText("Balance: -") + self.dataset_quality_readiness.setText("Readiness: -") + self.dataset_quality_cache.setText("Cache: -") + self.dataset_quality_duplicates.setText("Duplicates: -") + self.dataset_quality_extraction.setText("Extraction: -") + self.dataset_quality_warning.setText("Warnings: none") + if hasattr(self, "dataset_mix_chart"): + self.dataset_mix_chart.clear() + if hasattr(self, "dataset_sequence_chart"): + self.dataset_sequence_chart.clear() + if hasattr(self, "dataset_advisor"): + self.dataset_advisor.setPlainText("Run Preview Dataset to get cleanup suggestions.") + + def _card(self, title: str, content_layout: Union[QVBoxLayout, QFormLayout, QGridLayout, QHBoxLayout]) -> QWidget: + """Create a neon module card. + + Args: + title: Card heading. + content_layout: Layout to place inside the card. + + Returns: + Card widget. + """ + + card = QWidget() + card.setObjectName("Card") + layout = QVBoxLayout(card) + layout.setContentsMargins(14, 12, 14, 12) + layout.setSpacing(8) + title_label = QLabel(title) + title_label.setObjectName("SectionLabel") + layout.addWidget(title_label) + layout.addLayout(content_layout) + return card + + def _spin(self, minimum: int, maximum: int, value: int) -> QSpinBox: + """Create a bounded integer input. + + Args: + minimum: Minimum value. + maximum: Maximum value. + value: Initial value. + + Returns: + Configured spin box. + """ + + spin = QSpinBox() + spin.setRange(minimum, maximum) + spin.setValue(value) + spin.setMaximumWidth(220) + return spin + + def _double_spin(self, minimum: float, maximum: float, value: float, step: float, decimals: int) -> QDoubleSpinBox: + """Create a bounded float input. + + Args: + minimum: Minimum value. + maximum: Maximum value. + value: Initial value. + step: Increment step. + decimals: Number of displayed decimal places. + + Returns: + Configured double spin box. + """ + + spin = QDoubleSpinBox() + spin.setRange(minimum, maximum) + spin.setDecimals(decimals) + spin.setSingleStep(step) + spin.setValue(value) + spin.setMaximumWidth(220) + return spin + + def _path_row(self, field: QLineEdit, directory: bool = True, file_filter: str = "Checkpoints (*.pt)") -> QWidget: + """Create a path field with a browse button. + + Args: + field: Path input widget. + directory: Whether the browse dialog selects folders. + file_filter: File dialog filter used when ``directory`` is false. + + Returns: + Row widget containing the path input and button. + """ + + row = QWidget() + row.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed) + layout = QHBoxLayout(row) + layout.setContentsMargins(0, 0, 0, 0) + layout.setSpacing(8) + browse = QPushButton("Browse") + browse.setFixedWidth(88) + self._tip(browse, "Open a file/folder picker for this path.") + field.setMinimumWidth(180) + field.setSizePolicy(QSizePolicy.Ignored, QSizePolicy.Fixed) + browse.clicked.connect(lambda: self._browse(field, directory, file_filter)) + layout.addWidget(field, 1) + layout.addWidget(browse) + return row + + def _multi_file_path_row(self, field: QLineEdit, file_filter: str = "All files (*)") -> QWidget: + """Create a path field with a multi-file browse button. + + Args: + field: Path input widget. Multiple paths are separated with semicolons. + file_filter: File dialog filter. + + Returns: + Row widget containing the path input and browse button. + """ + + row = QWidget() + row.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed) + layout = QHBoxLayout(row) + layout.setContentsMargins(0, 0, 0, 0) + layout.setSpacing(8) + browse = QPushButton("Browse") + browse.setFixedWidth(88) + self._tip(browse, "Choose one or more JSON/JSONL files. You can also paste a folder path.") + field.setMinimumWidth(180) + field.setSizePolicy(QSizePolicy.Ignored, QSizePolicy.Fixed) + browse.clicked.connect(lambda: self._browse_multiple_files(field, file_filter)) + layout.addWidget(field, 1) + layout.addWidget(browse) + return row + + def _configure_form(self, form: QFormLayout) -> None: + """Apply common form spacing and growth policy. + + Args: + form: Form layout to configure. + """ + + form.setLabelAlignment(Qt.AlignRight | Qt.AlignVCenter) + form.setFormAlignment(Qt.AlignLeft | Qt.AlignTop) + form.setFieldGrowthPolicy(QFormLayout.AllNonFixedFieldsGrow) + form.setHorizontalSpacing(12) + form.setVerticalSpacing(7) + + def _configure_device_options(self) -> None: + """Populate training device choices without duplicate CPU entries.""" + + self.device.clear() + if torch.cuda.is_available(): + device_name = torch.cuda.get_device_name(0) + self.device.addItem("cuda") + self.device.addItem("cpu") + self.device_info.setText(f"CUDA ready: {device_name}") + self.use_amp_default = True + else: + self.device.addItem("cpu") + cuda_build = getattr(torch.backends, "cuda", None) + built_with_cuda = bool(cuda_build and torch.backends.cuda.is_built()) + if built_with_cuda: + detail = "CUDA build found, but no usable NVIDIA GPU/driver was detected." + else: + detail = "CUDA is not available in this PyTorch install." + self.device_info.setText(detail) + self.use_amp_default = False + + def _thin_progress(self) -> QProgressBar: + """Create a thin bottom progress bar. + + Returns: + Configured progress bar. + """ + + progress = QProgressBar() + progress.setRange(0, 100) + progress.setTextVisible(False) + progress.setFixedHeight(4) + progress.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed) + self._tip(progress, "Progress indicator for the current page operation.") + return progress + diff --git a/interface/main_window_part5.py b/interface/main_window_part5.py new file mode 100644 index 0000000..5a8fe33 --- /dev/null +++ b/interface/main_window_part5.py @@ -0,0 +1,420 @@ +from __future__ import annotations + +# MainWindow implementation mixin. Runtime names are provided by app.py. +from typing import Any +from . import app as _app + +globals().update({name: value for name, value in vars(_app).items() if not name.startswith("__")}) + +class MainWindowPart5: + def _tip(self, widget: QWidget, text: str) -> None: + """Attach tooltip and status tip text. + + Args: + widget: Widget receiving the tip. + text: Tooltip text. + """ + + widget.setToolTip(text) + widget.setStatusTip(text) + + def _render_chat_markdown(self, markdown_text: str) -> None: + """Render chat Markdown with highlighted fenced code blocks when possible. + + Args: + markdown_text: Markdown transcript to render. + """ + + if not hasattr(self, "current_assistant_message") or self.current_assistant_message is None: + return + self.current_assistant_message.set_content(markdown_text) + + def _add_chat_message( + self, + role: str, + content: str, + metrics: str = "", + resend_prompt: Optional[str] = None, + ) -> QTextBrowser: + """Add one chat bubble. + + Args: + role: Message role, either ``user`` or ``assistant``. + content: Markdown message content. + metrics: Optional metric text shown under assistant replies. + resend_prompt: Prompt to resend from the bubble. + + Returns: + Text browser used by the bubble. + """ + + should_follow = self._is_chat_near_bottom() + max_width = max(320, int(self.chat_scroll.viewport().width() * 0.78)) if hasattr(self, "chat_scroll") else 900 + message = ChatMessageWidget( + role, + content, + markdown_to_html, + self._resend_chat_message, + metrics=metrics, + resend_prompt=resend_prompt, + max_width=max_width, + ) + self.chat_messages.insertWidget(max(self.chat_messages.count() - 1, 0), message) + if should_follow: + message.scroll_later(lambda: self.chat_scroll.verticalScrollBar().setValue(self.chat_scroll.verticalScrollBar().maximum())) + if role == "assistant": + self.current_assistant_message = message + self.current_assistant_browser = message.browser + self.current_assistant_meta = message.meta_label + return message.browser + + def _is_chat_near_bottom(self) -> bool: + """Return whether the chat scroll is close enough to follow streaming. + + Returns: + True when the view should auto-scroll. + """ + + if not hasattr(self, "chat_scroll"): + return True + bar = self.chat_scroll.verticalScrollBar() + return bar.maximum() - bar.value() < 48 + + def _clear_chat_messages(self) -> None: + """Remove all message bubbles.""" + + while self.chat_messages.count() > 1: + item = self.chat_messages.takeAt(0) + widget = item.widget() + if widget is not None: + widget.deleteLater() + self.current_assistant_message = None + self.current_assistant_browser = None + self.current_assistant_meta = None + + def _resend_chat_message(self, prompt: str) -> None: + """Resend text from a message bubble. + + Args: + prompt: Prompt text to send. + """ + + self.chat_input.setPlainText(prompt) + self.send_chat_message() + + def _set_chat_stats(self, elapsed_seconds: float, token_count: int, tokens_per_second: float) -> None: + """Update live chat generation metrics. + + Args: + elapsed_seconds: Elapsed generation time. + token_count: Generated token count. + tokens_per_second: Approximate token speed. + """ + + text = f"Time: {elapsed_seconds:.2f}s | Tokens: {token_count:,} | Speed: {tokens_per_second:.2f} tok/s" + self.chat_stats.setText(text) + if self.current_assistant_meta is not None: + self.current_assistant_meta.setText(text) + self.current_assistant_meta.setVisible(True) + + def _chat_backend_value(self) -> str: + """Return the selected chat model backend. + + Returns: + Stable chat backend identifier. + """ + + if not hasattr(self, "chat_model_backend"): + return "gguf" + return "microgpt" if self.chat_model_backend.currentText() == "MicroGPT checkpoint" else "gguf" + + def _update_chat_backend_controls(self) -> None: + """Show controls relevant to the selected chat backend.""" + + if not hasattr(self, "chat_model_backend"): + return + native = self._chat_backend_value() == "microgpt" + self.gguf_path_row.setVisible(not native) + self.microgpt_path_row.setVisible(native) + self.llama_gpu_layers.setEnabled(not native) + self.llama_threads.setEnabled(not native) + self.llama_context.setEnabled(not native) + if native: + self._tip(self.load_llm_button, "Load the native MicroGPT checkpoint into memory once for repeated chat messages.") + else: + self._tip(self.load_llm_button, "Load the GGUF model into memory once for repeated chat messages.") + + def _app_icon(self) -> QIcon: + """Create the application icon. + + Returns: + Application icon. + """ + + return self._static_app_icon() + + @staticmethod + def _app_logo_path() -> Path: + """Return the bundled logo path. + + Returns: + Logo path. + """ + + candidates = [ + Path(__file__).resolve().parents[1] / "drunken_bot_logo_small.png", + Path(__file__).resolve().parents[2] / "drunken_bot_logo_small.png", + ] + if hasattr(sys, "_MEIPASS"): + candidates.insert(0, Path(sys._MEIPASS) / "drunken_bot_logo_small.png") + app_root = os.environ.get("DRUNKENBOT_APP_ROOT") + if app_root: + candidates.insert(0, Path(app_root) / "drunken_bot_logo_small.png") + return next((path for path in candidates if path.exists()), candidates[0]) + + @staticmethod + def _app_logo_pixmap(size: int = 64) -> QPixmap: + """Load the bundled logo as a pixmap. + + Args: + size: Maximum square size. + + Returns: + Logo pixmap, or null pixmap when the file is missing. + """ + + pixmap = QPixmap(str(MainWindowPart5._app_logo_path())) + if pixmap.isNull(): + return pixmap + return pixmap.scaled(size, size, Qt.KeepAspectRatio, Qt.SmoothTransformation) + + @staticmethod + def _static_app_icon() -> QIcon: + """Create the static app icon. + + Returns: + Application icon. + """ + + logo_path = MainWindowPart5._app_logo_path() + if logo_path.exists(): + icon = QIcon(str(logo_path)) + if not icon.isNull(): + return icon + pixmap = QPixmap(64, 64) + pixmap.fill(Qt.transparent) + painter = QPainter(pixmap) + try: + painter.setRenderHint(QPainter.Antialiasing) + painter.setBrush(QBrush(QColor("#1f1f1f"))) + painter.setPen(QPen(QColor("#f5b041"), 3)) + painter.drawRoundedRect(4, 4, 56, 56, 12, 12) + bolt = QPolygon([ + QPoint(36, 8), + QPoint(17, 35), + QPoint(31, 35), + QPoint(25, 56), + QPoint(48, 25), + QPoint(33, 25), + ]) + painter.setPen(QPen(QColor("#ffd27a"), 2)) + painter.setBrush(QBrush(QColor("#f5b041"))) + painter.drawPolygon(bolt) + finally: + painter.end() + return QIcon(pixmap) + + @staticmethod + def _windows_icon_path() -> Path: + """Return the Windows icon file path. + + Returns: + Path to the generated ``.ico`` file. + """ + + return Path(__file__).with_name("drunkenbot_llm_ide.ico") + + @staticmethod + def _ensure_windows_icon_file() -> Optional[Path]: + """Ensure the generated Windows ``.ico`` file exists. + + Returns: + Icon path on Windows, otherwise ``None``. + """ + + if sys.platform != "win32": + return None + icon_path = MainWindowPart5._windows_icon_path() + if icon_path.exists(): + return icon_path + icon = MainWindowPart5._static_app_icon() + pixmap = icon.pixmap(256, 256) + if pixmap.isNull() or not pixmap.save(str(icon_path), "ICO"): + return None + return icon_path + + def apply_windows_taskbar_icon(self) -> None: + """Apply the app icon to the native Windows window handle.""" + + if sys.platform != "win32": + return + icon_path = self._ensure_windows_icon_file() + if icon_path is None: + return + + hwnd = int(self.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) + self._windows_icon_handles.append(hicon_big) + if hicon_small: + user32.SendMessageW(hwnd, wm_seticon, icon_small, hicon_small) + self._windows_icon_handles.append(hicon_small) + + def _browse(self, field: QLineEdit, directory: bool, file_filter: str = "Checkpoints (*.pt)") -> None: + """Open a file or folder picker for a path field. + + Args: + field: Path input to update. + directory: Whether to select a folder instead of a file. + file_filter: File dialog filter used for files. + """ + + start_dir = self._browse_start_dir(field, directory) + if directory: + value = QFileDialog.getExistingDirectory(self, "Choose folder", start_dir) + else: + value, _ = QFileDialog.getOpenFileName(self, "Choose file", start_dir, file_filter) + if value: + field.setText(value) + + def _browse_multiple_files(self, field: QLineEdit, file_filter: str) -> None: + """Open a multi-file picker and write selected paths to a field. + + Args: + field: Path field to update. + file_filter: File dialog filter. + """ + + values, _ = QFileDialog.getOpenFileNames(self, "Choose files", self._browse_start_dir(field, False), file_filter) + if values: + field.setText("; ".join(values)) + + def _browse_start_dir(self, field: QLineEdit, directory: bool) -> str: + """Return the best initial folder for a browse dialog. + + Args: + field: Path field being browsed. + directory: Whether the dialog selects a folder. + + Returns: + Existing field path, active project folder, or current folder. + """ + + text = field.text().strip() + if text: + path = Path(text) + if path.exists(): + if path.is_dir(): + return str(path) + return str(path.parent) + parent = path if directory else path.parent + if parent.exists(): + return str(parent) + if self.current_project_file is not None: + return str(self.current_project_file.parent) + return str(Path.cwd()) + + def save_project(self) -> None: + """Save the current project settings into a named project folder.""" + + project_name = self.search_box.text().strip() or "MicroLLMProject" + safe_name = self._safe_project_name(project_name) + if self.current_project_file is None: + base_dir = QFileDialog.getExistingDirectory(self, "Choose parent folder for project", self._project_dialog_start_dir()) + if not base_dir: + return + project_dir = Path(base_dir) / safe_name + project_file = project_dir / "project.json" + else: + project_file = self.current_project_file + project_dir = project_file.parent + project_dir.mkdir(parents=True, exist_ok=True) + self._ensure_project_workspace(project_dir) + if self.current_project_file is None: + self._apply_project_workspace_paths(project_dir) + project_file.write_text(json.dumps(self._project_state_dict(project_name, project_dir), indent=2), encoding="utf-8") + self.current_project_file = project_file + _register_recent_project(project_file) + self._apply_project_runtime_environment(project_dir) + self._refresh_notification_manager(project_dir) + if hasattr(self, "runpod_api_key"): + self.load_runpod_settings() + self.project_state.setText("Project saved") + LOGGER.info("Project saved: %s", project_file) + if self.current_project_file == project_file: + self.dataset_log.append(f"Project saved: {project_file}") + self.dataset_log.append(f"Project workspace: {project_dir}") + self.dataset_log.append(f"Notifier config: {project_dir / 'notifier_config.json'}") + + def new_project(self) -> None: + """Start a fresh project and clear the active project file binding.""" + + if self.thread is not None: + QMessageBox.information(self, "Task running", "Please stop or wait for the current task before creating a new project.") + return + if self.current_project_file is not None or self.search_box.text().strip(): + choice = QMessageBox.question( + self, + "New project", + "Start a new project? Unsaved changes in the current project will not be saved automatically.", + QMessageBox.Yes | QMessageBox.No, + QMessageBox.No, + ) + if choice != QMessageBox.Yes: + return + + base_dir = QFileDialog.getExistingDirectory( + self, + "Choose folder where the new project will be created", + self._project_dialog_start_dir(), + ) + if not base_dir: + self.project_state.setText("Project creation cancelled") + return + if self.chat_session is not None and hasattr(self.chat_session, "reset"): + self.chat_session.reset() + self.chat_session = None + project_name = self.search_box.text().strip() or "MicroLLMProject" + try: + self._create_project_at(project_name, Path(base_dir)) + except Exception as exc: + QMessageBox.warning(self, "New project failed", f"Could not create project:\n{exc}") + + def open_project(self) -> None: + """Open a saved project file and restore UI settings.""" + + project_file, _ = QFileDialog.getOpenFileName( + self, + "Open Micro LLM project", + self._project_dialog_start_dir(), + "Micro LLM project (project.json *.json);;All files (*)", + ) + if not project_file: + return + try: + self._open_project_file(Path(project_file)) + except Exception as exc: + QMessageBox.warning(self, "Open failed", f"Could not open project:\n{exc}") + return diff --git a/interface/main_window_part6.py b/interface/main_window_part6.py new file mode 100644 index 0000000..d064143 --- /dev/null +++ b/interface/main_window_part6.py @@ -0,0 +1,376 @@ +from __future__ import annotations + +# MainWindow implementation mixin. Runtime names are provided by app.py. +from typing import Any +from . import app as _app + +globals().update({name: value for name, value in vars(_app).items() if not name.startswith("__")}) + +class MainWindowPart6: + def _create_project_at(self, project_name: str, base_dir: Path) -> Path: + """Create and activate a new project at the selected folder. + + Args: + project_name: User-facing project name. + base_dir: Parent folder for the new project. + + Returns: + Path to the created project.json file. + """ + + if self.chat_session is not None and hasattr(self.chat_session, "reset"): + self.chat_session.reset() + self.chat_session = None + project_dir = base_dir / self._safe_project_name(project_name) + project_file = project_dir / "project.json" + project_dir.mkdir(parents=True, exist_ok=True) + self._ensure_project_workspace(project_dir) + copied_count = self._ensure_project_training_data(project_dir) + self.current_project_file = project_file + self._apply_project_state(self._default_project_state()) + self.search_box.setText(project_name) + self._apply_project_workspace_paths(project_dir) + self._reset_dataset_blueprint_source(project_dir / "training_data") + self._apply_project_runtime_environment(project_dir) + self._refresh_notification_manager(project_dir) + if hasattr(self, "runpod_api_key"): + self.load_runpod_settings() + self._reset_project_runtime_state() + project_file.write_text(json.dumps(self._project_state_dict(project_name, project_dir), indent=2), encoding="utf-8") + _register_recent_project(project_file) + self.project_state.setText("New project") + LOGGER.info("New project created: %s", project_file) + self.dataset_log.append(f"Started a new project: {project_file}") + self.dataset_log.append(f"Project workspace: {project_dir}") + self.dataset_log.append( + "Bundled training data is no longer included; use Dataset Sources to download or select a dataset." + ) + self.dataset_log.append(f"Notifier config: {project_dir / 'notifier_config.json'}") + return project_file + + def _reset_dataset_blueprint_source(self, data_root: Path) -> None: + """Reset the current Dataset Sources tree without replacing its widget.""" + self.blueprint_data_root = Path(data_root) + if hasattr(self, "external_dataset_dir"): + self.external_dataset_dir.setText(str(data_root)) + if hasattr(self, "dataset_plan_source_label"): + self.dataset_plan_source_label.setText(f"Source: {data_root}") + if hasattr(self, "default_data_tree"): + self.default_data_tree.clear() + self.default_data_actions.clear() + self.default_data_category_items.clear() + self.default_data_tree.addTopLevelItem( + QTreeWidgetItem(["No project data files were found.", "", ""]) + ) + + def _open_project_file(self, project_file: Path) -> None: + """Open and activate a project file. + + Args: + project_file: Path to ``project.json``. + """ + + data = json.loads(project_file.read_text(encoding="utf-8")) + self.current_project_file = project_file + _register_recent_project(project_file) + self._ensure_project_workspace(self.current_project_file.parent) + dataset_state = data.get("dataset", {}) if isinstance(data, dict) else {} + saved_default_data_paths = dataset_state.get("default_data_paths") + self._refresh_dataset_blueprint_source( + self.current_project_file.parent / "training_data", + saved_paths=(list(saved_default_data_paths) if saved_default_data_paths is not None else None), + saved_plan=dict(dataset_state.get("domain_plan", {})), + preset=str(dataset_state.get("domain_plan_preset", "Balanced Tiny LLM")), + ) + self._apply_project_state(data) + self._apply_project_runtime_environment(self.current_project_file.parent) + self._refresh_notification_manager(self.current_project_file.parent) + if hasattr(self, "runpod_api_key"): + self.load_runpod_settings() + if self.model_dir.text().strip(): + self._load_existing_telemetry(Path(self.model_dir.text())) + self.project_state.setText("Project opened") + LOGGER.info("Project opened: %s", project_file) + self.dataset_log.append(f"Opened project: {project_file}") + self.dataset_log.append(f"Notifier config: {self.current_project_file.parent / 'notifier_config.json'}") + self.refresh_model_estimate() + + def _project_dialog_start_dir(self) -> str: + """Return the best initial folder for project dialogs. + + Returns: + Active project folder, its parent, or the current folder. + """ + + if self.current_project_file is not None: + return str(self.current_project_file.parent) + text = self.dataset_dir.text().strip() if hasattr(self, "dataset_dir") else "" + if text: + path = Path(text) + for candidate in (path, path.parent): + if candidate.exists(): + return str(candidate) + return str(Path.cwd()) + + def _ensure_project_workspace(self, project_dir: Path) -> None: + """Create standard folders inside a project. + + Args: + project_dir: Project root folder. + """ + + for name in ("datasets", "models", "fine_tunes", "exports", "training_data", "cache", "temp"): + (project_dir / name).mkdir(parents=True, exist_ok=True) + ensure_notifier_config(project_dir / "notifier_config.json") + ensure_runpod_config(project_dir / "runpod_config.json") + + def _ensure_project_training_data(self, project_dir: Path) -> int: + """Create the project training-data folder without bundling corpus files.""" + target_root = project_dir / "training_data" + target_root.mkdir(parents=True, exist_ok=True) + return 0 + + def _refresh_dataset_blueprint_source( + self, + data_root: Path, + saved_paths: Optional[list[Any]] = None, + saved_plan: Optional[dict[str, Any]] = None, + preset: str = "Balanced Tiny LLM", + ) -> None: + """Rebuild the Dataset Blueprint tab from a source data folder. + + Args: + data_root: Project-local training data folder. + saved_paths: Optional selected file paths to restore. + saved_plan: Optional saved domain recipe. + preset: Saved recipe preset. + """ + + if not hasattr(self, "pages"): + self.blueprint_data_root = Path(data_root) + return + self.blueprint_data_root = Path(data_root) + if hasattr(self, "default_data_tree"): + selected = saved_paths if saved_paths is not None else self._selected_default_data_paths() + populate_default_data_tree(self, self.blueprint_data_root) + self._set_selected_default_data_paths(selected) + if saved_plan is not None: + self._set_dataset_plan(saved_plan, preset) + self.dataset_plan_source_label.setText(f"Source: {self.blueprint_data_root}") + return + current_index = self.pages.currentIndex() + old_page = self.pages.widget(0) + old_page.hide() + QApplication.processEvents() + new_page = self._build_dataset_plan_tab() + self.pages.removeWidget(old_page) + old_page.setParent(None) + old_page.deleteLater() + self.pages.insertWidget(0, new_page) + if saved_plan is not None: + self._set_dataset_plan(saved_plan, preset) + if saved_paths is not None: + self._set_selected_default_data_paths(saved_paths) + elif self.current_project_file is not None: + self._set_selected_default_data_paths(None) + self.pages.setCurrentIndex(current_index) + + def download_latest_external_dataset(self) -> None: + """Download the latest managed dataset into the selected install folder.""" + destination = Path(self.external_dataset_dir.text()).expanduser() + try: + manifest = load_manifest() + except Exception as exc: + self.external_dataset_version.setText(f"Could not load dataset options: {exc}") + return + dialog = QDialog(self) + dialog.setWindowTitle("Select dataset components") + dialog.setMinimumWidth(480) + dialog_layout = QVBoxLayout(dialog) + dialog_layout.addWidget(QLabel(f"Dataset version {manifest.version}")) + version_selector = QComboBox() + version_selector.addItem(manifest.version, DEFAULT_MANIFEST_URL) + installed_version_file = destination / "version.txt" + if installed_version_file.is_file(): + installed_version = installed_version_file.read_text(encoding="utf-8").strip() + if installed_version and installed_version != manifest.version: + version_selector.addItem( + installed_version, + f"https://github.com/drunkenbot-ai/dataset/releases/download/dataset-v{installed_version}/manifest.json", + ) + dialog_layout.addWidget(QLabel("Dataset version")) + dialog_layout.addWidget(version_selector) + component_tree = QTreeWidget() + component_tree.setHeaderLabels(["Component", "Files"]) + for category in manifest.categories: + if category.file_count <= 0: + continue + item = QTreeWidgetItem([category.name, str(category.file_count)]) + item.setFlags(item.flags() | Qt.ItemIsUserCheckable) + existing = destination / category.name + has_existing = existing.exists() and any(existing.rglob("*")) + item.setCheckState(0, Qt.Checked if has_existing else Qt.Unchecked) + component_tree.addTopLevelItem(item) + dialog_layout.addWidget(component_tree) + buttons = QHBoxLayout() + download_button = QPushButton("Download selected") + cancel_button = QPushButton("Cancel") + buttons.addStretch(1) + buttons.addWidget(cancel_button) + buttons.addWidget(download_button) + dialog_layout.addLayout(buttons) + cancel_button.clicked.connect(dialog.reject) + download_button.clicked.connect(dialog.accept) + if dialog.exec() != QDialog.Accepted: + return + categories = [ + component_tree.topLevelItem(index).text(0) + for index in range(component_tree.topLevelItemCount()) + if component_tree.topLevelItem(index).checkState(0) == Qt.Checked + ] + if not categories: + self.external_dataset_version.setText("Select at least one dataset component.") + return + + self.dataset_log.append(f"Downloading latest external dataset to {destination}...") + self.external_dataset_version.setText("Downloading latest dataset...") + self.dataset_plan_progress.setVisible(True) + self._run_task( + partial(download_latest_dataset, manifest_url=version_selector.currentData()), + (destination, categories), + self._external_dataset_download_finished, + self.dataset_log, + self.dataset_plan_progress, + with_progress=True, + button=self.external_dataset_download_button, + busy_text="Downloading dataset", + task_kind="dataset_download", + ) + + def _refresh_external_dataset_status(self) -> None: + """Restore the installed dataset version from the selected folder.""" + if not hasattr(self, "external_dataset_dir"): + return + version_file = Path(self.external_dataset_dir.text()).expanduser() / "version.txt" + if not version_file.is_file(): + self.external_dataset_version.setText("Installed version: not installed") + self.external_dataset_download_button.setEnabled(True) + return + try: + version = version_file.read_text(encoding="utf-8").strip() + except OSError: + version = "" + if version: + has_dataset_files = any( + path.is_file() and path.name not in {"version.txt", "manifest.json"} + for path in Path(self.external_dataset_dir.text()).expanduser().rglob("*") + ) + if not has_dataset_files: + self.external_dataset_version.setText("Installed version: not installed") + self.external_dataset_download_button.setEnabled(True) + return + self.external_dataset_version.setText(f"Installed version: {version}") + try: + latest = load_manifest() + except Exception as exc: + LOGGER.warning("Could not check for a newer dataset release: %s", exc) + self.external_dataset_download_button.setEnabled(True) + return + installed_root = Path(self.external_dataset_dir.text()).expanduser() + missing_components = [ + category.name + for category in latest.categories + if category.file_count > 0 + and not ( + (installed_root / category.name).exists() + and any((installed_root / category.name).rglob("*")) + ) + ] + if is_newer_version(latest.version, version): + self.external_dataset_version.setText( + f"Installed version: {version} (update available: {latest.version})" + ) + self.external_dataset_download_button.setEnabled(True) + elif missing_components: + self.external_dataset_version.setText( + f"Installed version: {version} ({len(missing_components)} components missing)" + ) + self.external_dataset_download_button.setEnabled(True) + else: + self.external_dataset_download_button.setEnabled(False) + else: + self.external_dataset_version.setText("Installed version: not installed") + self.external_dataset_download_button.setEnabled(True) + + @Slot(object) + def _external_dataset_download_finished(self, manifest: Any) -> None: + """Apply the downloaded dataset as the active source vault.""" + self.dataset_log.append(f"Installed external dataset version {manifest.version}.") + self.external_dataset_version.setText(f"Installed version: {manifest.version}") + self.external_dataset_download_button.setEnabled(False) + self.dataset_plan_progress.setVisible(False) + self._refresh_dataset_blueprint_source( + Path(self.external_dataset_dir.text()), + saved_plan=self._dataset_plan_from_ui(), + preset=( + self.dataset_plan_preset.currentText() + if hasattr(self, "dataset_plan_preset") + else "Balanced Tiny LLM" + ), + ) + self.input_dir.setText(self.external_dataset_dir.text()) + + def _refresh_notification_manager(self, project_dir: Optional[Path] = None) -> None: + """Load notification settings for the current project. + + Args: + project_dir: Optional project root folder. + """ + + if project_dir is None and self.current_project_file is not None: + project_dir = self.current_project_file.parent + config_path = default_notifier_config_path(project_dir) + self.notification_manager = NotificationManager(config_path) + LOGGER.info("Notifier config active: %s", config_path) + + def _apply_project_workspace_paths(self, project_dir: Path) -> None: + """Point project output fields at the standard project folders. + + Args: + project_dir: Project root folder. + """ + + dataset_dir = project_dir / "datasets" + model_dir = project_dir / "models" + fine_tune_dir = project_dir / "fine_tunes" + export_dir = project_dir / "exports" + training_data_dir = project_dir / "training_data" + self.dataset_dir.setText(str(dataset_dir)) + self.train_data_dir.setText(str(dataset_dir)) + self.model_dir.setText(str(model_dir)) + self.fine_tune_checkpoint.setText(str(model_dir / "final_model.pt")) + self.fine_tune_output_dir.setText(str(fine_tune_dir / "latest")) + self.export_model_dir.setText(str(model_dir)) + self.export_dir.setText(str(export_dir)) + self.gguf_output_path.setText(str(export_dir / "model.gguf")) + if not self.input_dir.text().strip(): + self.input_dir.setText(str(training_data_dir)) + + def _apply_project_runtime_environment(self, project_dir: Path) -> None: + """Prefer project-local cache/temp folders for runtime work. + + Args: + project_dir: Project root folder. + """ + + cache_dir = project_dir / "cache" + temp_dir = project_dir / "temp" + cache_dir.mkdir(parents=True, exist_ok=True) + temp_dir.mkdir(parents=True, exist_ok=True) + for key in ("TMPDIR", "TEMP", "TMP"): + os.environ[key] = str(temp_dir) + for key in ("TORCH_HOME", "HF_HOME", "TRANSFORMERS_CACHE", "PYTORCH_KERNEL_CACHE"): + os.environ[key] = str(cache_dir / key.lower()) + Path(os.environ[key]).mkdir(parents=True, exist_ok=True) + + diff --git a/interface/main_window_part7.py b/interface/main_window_part7.py new file mode 100644 index 0000000..7cbbe1b --- /dev/null +++ b/interface/main_window_part7.py @@ -0,0 +1,384 @@ +from __future__ import annotations + +# MainWindow implementation mixin. Runtime names are provided by app.py. +from typing import Any +from . import app as _app + +globals().update({name: value for name, value in vars(_app).items() if not name.startswith("__")}) + +class MainWindowPart7: + def _default_project_state(self) -> dict[str, Any]: + """Build the default state used for a newly created project. + + Returns: + JSON-style project state with fresh paths and default settings. + """ + + runs_dir = Path.cwd() / "runs" + dataset_dir = runs_dir / "dataset" + model_dir = runs_dir / "model" + fine_tune_dir = runs_dir / "fine_tune" + export_dir = runs_dir / "export" + return { + "schema": "drunkenbot_ide_project", + "version": 1, + "project_name": "", + "project_dir": "", + "paths": { + "source_vault": "", + "dataset_core": str(dataset_dir), + "training_dataset": str(dataset_dir), + "model_output": str(model_dir), + "export_model_core": str(model_dir), + "export_output": str(export_dir), + "llama_cpp_dir": "", + "gguf_output_path": str(export_dir / "model.gguf"), + "gguf_model": "", + "microgpt_chat_model": "", + "tokenizer_import": "", + "resume_checkpoint": "", + "fine_tune_checkpoint": "", + "fine_tune_output": str(fine_tune_dir), + }, + "dataset": { + "domain_plan_preset": "Balanced Tiny LLM", + "domain_plan": dataset_plan_defaults(), + "default_data_paths": [str(path) for path, _category in iter_default_data_files()], + "auto_vocab": True, + "manual_vocab_size": 8000, + "include_conversation_datasets": False, + "dataset_stage": "base", + "conversation_datasets": [], + "conversation_sample_limit": 20000, + "conversation_dataset_path": "", + "instruction_dataset_path": "", + "mixture_weights": {}, + "min_frequency": 2, + "context_length": 128, + "validation_split": 0.1, + "lowercase": False, + "max_workers": 4, + "prepare_mode": "incremental", + "tokenizer_strategy": "auto", + "code_training_mode": True, + "include_prose": True, + "include_source_code": True, + "extract_code_blocks": True, + "preserve_indentation": True, + "instruction_samples": True, + "reasoning_sample_mode": "scaffold", + }, + "training": { + "preset": "Tiny", + "architecture_style": "Classic GPT", + "launch_target": "local", + "training_mode": "pretrain", + "training_stage": "base", + "peft_method": "none", + "lora_rank": 8, + "lora_alpha": 16.0, + "lora_dropout": 0.05, + "lora_target_modules": "attention", + "n_embd": 128, + "n_head": 4, + "n_layer": 4, + "context_length": 128, + "dropout": 0.1, + "training_profile": "Stable LLM", + "epochs": 5, + "batch_size": 16, + "learning_rate": 0.0003, + "weight_decay": 0.1, + "gradient_accumulation": 1, + "warmup_steps": 100, + "eval_interval": 100, + "max_eval_batches": 50, + "save_interval": 500, + "data_loader_workers": 0, + "max_grad_norm": 1.0, + "activation_checkpointing": False, + "seed": 1337, + "device": self.device.currentText(), + "use_amp": self.use_amp_default, + "resume": True, + "require_compatible_resume": True, + "benchmark_prompts": "\n\n".join(DEFAULT_BENCHMARK_PROMPTS), + "benchmark_tokens": 128, + "benchmark_temperature": 0.7, + "benchmark_kv_cache": True, + }, + "export": { + "quantization": "FP16 checkpoint", + "gguf_outtype": "f16", + }, + "chat": { + "model_backend": "gguf", + "context": 2048, + "cpu_threads": 4, + "gpu_layers": -1, + "thinking_enabled": True, + "reasoning_effort": "Balanced", + "max_tokens": 512, + "temperature": 0.7, + "top_p": 0.9, + "repeat_penalty": 1.1, + "system_prompt": "", + }, + "distributed": { + "host": "0.0.0.0", + "port": 8765, + "artifact_root": str(Path.home() / ".drunkenbot_ide" / "artifacts"), + "public_url": "http://127.0.0.1:8765", + }, + "artifacts": {}, + } + + def _reset_project_runtime_state(self) -> None: + """Clear logs, progress, charts, and status labels for a new project.""" + + self.dataset_log.clear() + self.training_log.clear() + self.fine_tune_log.clear() + self.benchmark_log.clear() + self.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" + ) + for progress in ( + self.dataset_progress, + self.training_progress, + self.fine_tune_progress, + self.benchmark_progress, + self.export_progress, + self.chat_progress, + ): + progress.setRange(0, 100) + progress.setValue(0) + self.dataset_status.setText("Dataset: not prepared") + self.train_status.setText("Training: idle") + self.export_status.setText("Export: waiting") + self.chat_status.setText("Chat: no model loaded") + self.prepare_button.setText("Prepare Dataset") + self.train_button.setText("Start Training") + self.fine_tune_button.setText("Start Fine-Tune") + self.stop_dataset_button.setEnabled(False) + self.stop_training_button.setEnabled(False) + self.stop_fine_tune_button.setEnabled(False) + self.stop_benchmark_button.setEnabled(False) + self.stop_chat_button.setEnabled(False) + self.load_llm_button.setText("Load Model") + self._update_chat_backend_controls() + self._reset_dataset_quality_report() + self.training_epoch_metric.setText("Epoch: -") + self.training_step_metric.setText("Step: -") + self.training_loss_metric.setText("Train loss: -") + self.training_val_metric.setText("Val loss: -") + self.training_health_metric.setText("Health: -") + self.training_health_points = [] + self.training_lr_metric.setText("LR: -") + self.training_speed_metric.setText("Speed: -") + self.training_grad_metric.setText("Grad: -") + self.training_vram_metric.setText("VRAM: -") + self.training_eta_metric.setText("ETA: -") + self.model_size_metric.setText("Model: -") + self.vram_estimate_metric.setText("VRAM est: -") + self.parameter_breakdown_metric.setText("Params: -") + self.memory_breakdown_metric.setText("Memory: -") + self.architecture_advisor_metric.setText("Advisor: -") + self.history_metric.setText(f"Runs: {len(self._load_training_history())}") + self.loss_chart.clear() + self.optimization_chart.clear() + self.stability_chart.clear() + self.throughput_chart.clear() + self.memory_chart.clear() + self.live_prediction_chart.update_distribution(0, None) + self.live_attention_chart.update_heatmap(0, None) + self.live_activation_chart.update_histogram(0, None) + self.live_gradient_chart.update_flow(self.n_layer.value(), None, 0) + self.live_sample_text.setText("Training text: -") + self.telemetry_db_path = None + self.telemetry_run_id = "" + self.telemetry_latest_id = 0 + self.telemetry_latest_index = 0 + self.live_scrub_active = False + self.live_time_slider.blockSignals(True) + self.live_time_slider.setRange(0, 0) + self.live_time_slider.setValue(0) + self.live_time_slider.blockSignals(False) + self.live_timeline_label.setText("Timeline: no saved telemetry") + self._set_meter(self.live_cpu_bar, "CPU", self._system_cpu_value()) + self._set_meter(self.live_gpu_bar, "GPU memory", None) + self._set_meter(self.live_vram_bar, "VRAM reserved", None) + self._set_meter(self.live_ram_bar, "System RAM", self._system_ram_value()) + self.live_worker_status.setText(f"CPU workers: {self.data_loader_workers.value()}") + self._clear_chat_messages() + self.chat_markdown = "" + self.chat_stream_prefix = "" + self.chat_stream_reply = "" + self.chat_stats.setText("Idle") + self._add_chat_message("assistant", "Load a GGUF or MicroGPT model to start testing.") + + def _project_state_dict(self, project_name: str, project_dir: Path) -> dict[str, Any]: + """Collect all UI state that defines a Micro LLM project. + + Args: + project_name: User-facing project name. + project_dir: Folder where the project file will live. + + Returns: + JSON-serializable project state. + """ + + dataset_dir = Path(self.dataset_dir.text()) if self.dataset_dir.text().strip() else None + model_dir = Path(self.model_dir.text()) if self.model_dir.text().strip() else None + export_dir = Path(self.export_dir.text()) if self.export_dir.text().strip() else None + now_iso = datetime.now().isoformat(timespec="seconds") + created_at = now_iso + existing_project_file = project_dir / "project.json" + if existing_project_file.exists(): + try: + existing_data = json.loads(existing_project_file.read_text(encoding="utf-8")) + except Exception: + existing_data = {} + if isinstance(existing_data, dict): + # Preserve the original creation timestamp across saves. + # "saved_at" below is overwritten every save, so it cannot be + # used as a creation date; fall back to it only for projects + # saved before this field existed. + created_at = str(existing_data.get("created_at") or existing_data.get("saved_at") or now_iso) + return { + "schema": "drunkenbot_ide_project", + "version": 1, + "project_name": project_name, + "project_dir": str(project_dir), + "created_at": created_at, + "saved_at": now_iso, + "paths": { + "source_vault": self.input_dir.text(), + "dataset_core": self.dataset_dir.text(), + "training_dataset": self.train_data_dir.text(), + "model_output": self.model_dir.text(), + "export_model_core": self.export_model_dir.text(), + "export_output": self.export_dir.text(), + "llama_cpp_dir": self.llama_cpp_dir.text(), + "gguf_output_path": self.gguf_output_path.text(), + "gguf_model": self.gguf_path.text(), + "microgpt_chat_model": self.microgpt_chat_path.text(), + "tokenizer_import": self.tokenizer_path.text(), + "resume_checkpoint": self.resume_checkpoint.text(), + "fine_tune_checkpoint": self.fine_tune_checkpoint.text(), + "fine_tune_output": self.fine_tune_output_dir.text(), + }, + "dataset": { + "domain_plan_preset": self.dataset_plan_preset.currentText() if hasattr(self, "dataset_plan_preset") else "Balanced Tiny LLM", + "domain_plan": self._dataset_plan_from_ui(), + "default_data_paths": [str(path) for path in self._selected_default_data_paths()], + "external_dataset_dir": self.external_dataset_dir.text() if hasattr(self, "external_dataset_dir") else "", + "auto_vocab": self.auto_vocab.isChecked(), + "manual_vocab_size": self.manual_vocab_size.value(), + "include_conversation_datasets": self.include_conversation_datasets.isChecked(), + "dataset_stage": self._dataset_stage_value(), + "conversation_datasets": self._selected_conversation_datasets(), + "conversation_sample_limit": self.conversation_sample_limit.value(), + "mixture_weights": self._mixture_weights_from_ui(), + "min_frequency": self.min_frequency.value(), + "context_length": self.context_length.value(), + "validation_split": self.validation_split.value(), + "lowercase": False, + "max_workers": self.max_workers.value(), + "prepare_mode": self._prepare_mode_value(), + "tokenizer_strategy": self._tokenizer_strategy_value(), + "code_training_mode": self.code_training_mode.isChecked(), + "include_prose": self.include_prose.isChecked(), + "include_source_code": self.include_source_code.isChecked(), + "extract_code_blocks": self.extract_code_blocks.isChecked(), + "preserve_indentation": self.preserve_indentation.isChecked(), + "instruction_samples": self.instruction_samples.isChecked(), + "reasoning_sample_mode": self._reasoning_sample_mode_value(), + }, + "training": { + "preset": self.preset.currentText(), + "architecture_style": self.architecture_style.currentText(), + "launch_target": self._training_launch_target_value(), + "fine_tune_launch_target": self._fine_tune_launch_target_value(), + "training_stage": self._training_stage_value(), + "n_embd": self.n_embd.value(), + "n_head": self.n_head.value(), + "attention_type": self._attention_type_value(), + "kv_head_count": self.kv_head_count.value(), + "attention_backend": self._attention_backend_value(), + "attention_window": self.attention_window.value(), + "training_mode": self._training_mode_value(), + "peft_method": self._peft_method_value(), + "lora_rank": self.lora_rank.value(), + "lora_alpha": self.lora_alpha.value(), + "lora_dropout": self.lora_dropout.value(), + "lora_target_modules": self._lora_target_value(), + "n_layer": self.n_layer.value(), + "context_length": self.train_context_length.value(), + "dropout": self.dropout.value(), + "training_profile": self.training_profile.currentText(), + "epochs": self.epochs.value(), + "batch_size": self.batch_size.value(), + "learning_rate": self.learning_rate.value(), + "weight_decay": self.weight_decay.value(), + "optimizer_name": self._optimizer_value(), + "scheduler_name": self._scheduler_value(), + "scheduler_min_lr_ratio": self.min_lr_ratio.value(), + "polynomial_power": self.polynomial_power.value(), + "gradient_accumulation": self.gradient_accumulation.value(), + "sample_stride": self.sample_stride.value(), + "warmup_steps": self.warmup_steps.value(), + "eval_interval": self.eval_interval.value(), + "max_eval_batches": self.max_eval_batches.value(), + "save_interval": self.save_interval.value(), + "data_loader_workers": self.data_loader_workers.value(), + "max_grad_norm": self.max_grad_norm.value(), + "activation_checkpointing": self.activation_checkpointing.isChecked(), + "seed": self.seed.value(), + "device": self.device.currentText(), + "use_amp": self.use_amp.isChecked(), + "precision": self._precision_value(), + "resume": self.resume_training.isChecked(), + "require_compatible_resume": self.resume_safety.isChecked(), + "early_stopping": self.early_stopping.isChecked(), + "benchmark_prompts": self.benchmark_prompts.toPlainText(), + "benchmark_tokens": self.benchmark_tokens.value(), + "benchmark_temperature": self.benchmark_temperature.value(), + "benchmark_kv_cache": self.benchmark_kv_cache.isChecked(), + }, + "export": { + "quantization": self.quant_mode.currentText(), + "gguf_outtype": self.gguf_outtype.currentText(), + }, + "chat": { + "model_backend": self._chat_backend_value(), + "context": self.llama_context.value(), + "cpu_threads": self.llama_threads.value(), + "gpu_layers": self.llama_gpu_layers.value(), + "thinking_enabled": self.thinking_enabled.isChecked(), + "reasoning_effort": self.reasoning_effort.currentText(), + "max_tokens": self.chat_max_tokens.value(), + "temperature": self.chat_temperature.value(), + "top_p": self.chat_top_p.value(), + "repeat_penalty": self.chat_repeat_penalty.value(), + "system_prompt": self.system_prompt.toPlainText(), + }, + "distributed": { + "host": self.coordinator_host.text(), + "port": self.coordinator_port.value(), + "artifact_root": self.coordinator_artifact_root.text(), + "public_url": self.coordinator_public_url.text(), + }, + "artifacts": { + "dataset_summary": self._read_json_if_exists(dataset_dir / "dataset_summary.json") if dataset_dir else None, + "training_summary": self._read_json_if_exists(model_dir / "training_summary.json") if model_dir else None, + "export_summary": self._read_json_if_exists(export_dir / "export_summary.json") if export_dir else None, + }, + } + + diff --git a/interface/main_window_part8.py b/interface/main_window_part8.py new file mode 100644 index 0000000..6764939 --- /dev/null +++ b/interface/main_window_part8.py @@ -0,0 +1,429 @@ +from __future__ import annotations + +# MainWindow implementation mixin. Runtime names are provided by app.py. +from typing import Any +from . import app as _app + +globals().update({name: value for name, value in vars(_app).items() if not name.startswith("__")}) + +class MainWindowPart8: + def _apply_project_state(self, data: dict[str, Any]) -> None: + """Restore UI state from a saved project dictionary. + + Args: + data: Project state loaded from JSON. + """ + + self.search_box.setText(str(data.get("project_name", ""))) + paths = data.get("paths", {}) + dataset = data.get("dataset", {}) + training = data.get("training", {}) + export = data.get("export", {}) + chat = data.get("chat", {}) + distributed = data.get("distributed", {}) + + self.input_dir.setText(str(paths.get("source_vault", ""))) + self.dataset_dir.setText(str(paths.get("dataset_core", ""))) + self.train_data_dir.setText(str(paths.get("training_dataset", ""))) + self.model_dir.setText(str(paths.get("model_output", ""))) + self.export_model_dir.setText(str(paths.get("export_model_core", ""))) + self.export_dir.setText(str(paths.get("export_output", ""))) + self.llama_cpp_dir.setText(str(paths.get("llama_cpp_dir", ""))) + self.gguf_output_path.setText(str(paths.get("gguf_output_path", ""))) + self.gguf_path.setText(str(paths.get("gguf_model", ""))) + self.microgpt_chat_path.setText(str(paths.get("microgpt_chat_model", ""))) + self.tokenizer_path.setText(str(paths.get("tokenizer_import", ""))) + self.resume_checkpoint.setText(str(paths.get("resume_checkpoint", ""))) + self.fine_tune_checkpoint.setText(str(paths.get("fine_tune_checkpoint", ""))) + self.fine_tune_output_dir.setText(str(paths.get("fine_tune_output", ""))) + + self._set_dataset_plan( + dict(dataset.get("domain_plan", {})), + str(dataset.get("domain_plan_preset", "Balanced Tiny LLM")), + ) + saved_default_data_paths = dataset.get("default_data_paths") + self._set_selected_default_data_paths( + list(saved_default_data_paths) if saved_default_data_paths is not None else None + ) + if hasattr(self, "external_dataset_dir") and dataset.get("external_dataset_dir"): + self.external_dataset_dir.setText(str(dataset["external_dataset_dir"])) + self._refresh_external_dataset_status() + self.auto_vocab.setChecked(bool(dataset.get("auto_vocab", True))) + self.manual_vocab_size.setValue(int(dataset.get("manual_vocab_size", self.manual_vocab_size.value()))) + include_conversation = bool(dataset.get("include_conversation_datasets", False)) + self._set_dataset_stage(str(dataset.get("dataset_stage", "base"))) + self.include_conversation_datasets.setChecked(include_conversation) + self._set_selected_conversation_datasets(list(dataset.get("conversation_datasets", []))) + self.conversation_sample_limit.setValue(int(dataset.get("conversation_sample_limit", self.conversation_sample_limit.value()))) + self._set_mixture_weights(dict(dataset.get("mixture_weights", {}))) + self.min_frequency.setValue(int(dataset.get("min_frequency", self.min_frequency.value()))) + self.context_length.setValue(int(dataset.get("context_length", self.context_length.value()))) + self.validation_split.setValue(float(dataset.get("validation_split", self.validation_split.value()))) + self.max_workers.setValue(int(dataset.get("max_workers", self.max_workers.value()))) + self._set_combo_by_data(self.prepare_mode, str(dataset.get("prepare_mode", "incremental")), { + "incremental": "Incremental update", + "full_rebuild": "Full rebuild", + "force_reprocess": "Force reprocess", + }) + self._set_combo_by_data(self.tokenizer_strategy, str(dataset.get("tokenizer_strategy", "auto")), { + "auto": "Auto", + "train_new": "Train new tokenizer", + "reuse_dataset": "Reuse dataset tokenizer", + "import_tokenizer": "Import tokenizer.json", + }) + self.code_training_mode.setChecked(bool(dataset.get("code_training_mode", True))) + self.include_prose.setChecked(bool(dataset.get("include_prose", True))) + self.include_source_code.setChecked(bool(dataset.get("include_source_code", True))) + self.extract_code_blocks.setChecked(bool(dataset.get("extract_code_blocks", True))) + self.preserve_indentation.setChecked(bool(dataset.get("preserve_indentation", True))) + self.instruction_samples.setChecked(bool(dataset.get("instruction_samples", True))) + self._set_combo_by_data(self.reasoning_sample_mode, str(dataset.get("reasoning_sample_mode", "scaffold")), { + "scaffold": "Reasoning scaffold", + "detailed": "Detailed code reasoning", + "none": "No reasoning wrapper", + }) + + self._set_combo_text(self.preset, str(training.get("preset", self.preset.currentText()))) + self._set_combo_text(self.architecture_style, str(training.get("architecture_style", self.architecture_style.currentText()))) + self._set_combo_by_data(self.training_launch_target, str(training.get("launch_target", "local")), { + "local": "Local machine", + "remote": "Remote workers", + "runpod": "RunPod cloud", + }) + if hasattr(self, "fine_tune_launch_target"): + self._set_combo_by_data(self.fine_tune_launch_target, str(training.get("fine_tune_launch_target", "local")), { + "local": "Local machine", + "remote": "Remote workers", + "runpod": "RunPod cloud", + }) + self.n_embd.setValue(int(training.get("n_embd", self.n_embd.value()))) + self.n_head.setValue(int(training.get("n_head", self.n_head.value()))) + self._set_combo_by_data(self.attention_type, str(training.get("attention_type", "mha")), { + "mha": "Multi-head", + "gqa": "Grouped-query", + "mqa": "Multi-query", + }) + self.kv_head_count.setValue(int(training.get("kv_head_count", self.kv_head_count.value()))) + self._set_combo_by_data(self.attention_backend, str(training.get("attention_backend", "sdpa")), { + "sdpa": "SDPA / Flash when available", + "manual": "Manual", + }) + self.attention_window.setValue(int(training.get("attention_window", self.attention_window.value()))) + self._set_combo_by_data(self.training_mode, str(training.get("training_mode", "pretrain")), { + "pretrain": "Pretrain from scratch", + "fine_tune": "Fine-tune checkpoint", + "instruction_fine_tune": "Instruction fine-tune", + "conversation_fine_tune": "Conversation fine-tune", + "code_fine_tune": "Code fine-tune", + }) + training_stage = str(training.get("training_stage", "")) + if training_stage == "instruction": + self._set_combo_text(self.training_mode, "Instruction fine-tune") + elif training_stage == "conversation": + self._set_combo_text(self.training_mode, "Conversation fine-tune") + elif training_stage == "code": + self._set_combo_text(self.training_mode, "Code fine-tune") + self._set_combo_by_data(self.peft_method, str(training.get("peft_method", "none")), { + "none": "Full fine-tune", + "lora": "LoRA adapters", + }) + self.lora_rank.setValue(int(training.get("lora_rank", self.lora_rank.value()))) + self.lora_alpha.setValue(float(training.get("lora_alpha", self.lora_alpha.value()))) + self.lora_dropout.setValue(float(training.get("lora_dropout", self.lora_dropout.value()))) + self._set_combo_by_data(self.lora_targets, str(training.get("lora_target_modules", "attention")), { + "attention": "Attention projections", + "mlp": "MLP projections", + "attention,mlp": "Attention + MLP", + }) + self.n_layer.setValue(int(training.get("n_layer", self.n_layer.value()))) + self.train_context_length.setValue(int(training.get("context_length", self.train_context_length.value()))) + self.dropout.setValue(float(training.get("dropout", self.dropout.value()))) + self._set_combo_text(self.training_profile, str(training.get("training_profile", self.training_profile.currentText()))) + self.epochs.setValue(int(training.get("epochs", self.epochs.value()))) + self.batch_size.setValue(int(training.get("batch_size", self.batch_size.value()))) + self.learning_rate.setValue(float(training.get("learning_rate", self.learning_rate.value()))) + self.weight_decay.setValue(float(training.get("weight_decay", self.weight_decay.value()))) + self._set_combo_by_data(self.optimizer_name, str(training.get("optimizer_name", "adamw")), { + "adamw": "AdamW", + "adam": "Adam", + "lion": "Lion", + "adafactor": "Adafactor", + }) + self._set_combo_by_data(self.scheduler_name, str(training.get("scheduler_name", "warmup_linear")), { + "warmup_linear": "Warmup linear", + "cosine": "Cosine decay", + "polynomial": "Polynomial decay", + "one_cycle": "One-cycle", + "constant": "Constant", + }) + self.min_lr_ratio.setValue(float(training.get("scheduler_min_lr_ratio", self.min_lr_ratio.value()))) + self.polynomial_power.setValue(float(training.get("polynomial_power", self.polynomial_power.value()))) + self.gradient_accumulation.setValue(int(training.get("gradient_accumulation", self.gradient_accumulation.value()))) + self.sample_stride.setValue(int(training.get("sample_stride", self.sample_stride.value()))) + self.warmup_steps.setValue(int(training.get("warmup_steps", self.warmup_steps.value()))) + self.eval_interval.setValue(int(training.get("eval_interval", self.eval_interval.value()))) + self.max_eval_batches.setValue(int(training.get("max_eval_batches", self.max_eval_batches.value()))) + self.save_interval.setValue(int(training.get("save_interval", self.save_interval.value()))) + self.data_loader_workers.setValue(int(training.get("data_loader_workers", self.data_loader_workers.value()))) + self.max_grad_norm.setValue(float(training.get("max_grad_norm", self.max_grad_norm.value()))) + self.activation_checkpointing.setChecked(bool(training.get("activation_checkpointing", False))) + self.seed.setValue(int(training.get("seed", self.seed.value()))) + self._set_combo_text(self.device, str(training.get("device", self.device.currentText()))) + self.use_amp.setChecked(bool(training.get("use_amp", self.use_amp.isChecked()))) + self._set_combo_by_data(self.precision, str(training.get("precision", "fp16")), { + "fp16": "FP16", + "bf16": "BF16", + "fp32": "FP32", + }) + self.resume_training.setChecked(bool(training.get("resume", self.resume_training.isChecked()))) + self.resume_safety.setChecked(bool(training.get("require_compatible_resume", True))) + self.early_stopping.setChecked(bool(training.get("early_stopping", True))) + self.benchmark_prompts.setPlainText(str(training.get("benchmark_prompts", self.benchmark_prompts.toPlainText()))) + self.benchmark_tokens.setValue(int(training.get("benchmark_tokens", self.benchmark_tokens.value()))) + self.benchmark_temperature.setValue(float(training.get("benchmark_temperature", self.benchmark_temperature.value()))) + self.benchmark_kv_cache.setChecked(bool(training.get("benchmark_kv_cache", True))) + + self._set_combo_text(self.quant_mode, str(export.get("quantization", self.quant_mode.currentText()))) + self._set_combo_text(self.gguf_outtype, str(export.get("gguf_outtype", self.gguf_outtype.currentText()))) + self.llama_context.setValue(int(chat.get("context", self.llama_context.value()))) + self._set_combo_by_data(self.chat_model_backend, str(chat.get("model_backend", "gguf")), { + "gguf": "GGUF / llama.cpp", + "microgpt": "MicroGPT checkpoint", + }) + self.llama_threads.setValue(int(chat.get("cpu_threads", self.llama_threads.value()))) + self.llama_gpu_layers.setValue(int(chat.get("gpu_layers", self.llama_gpu_layers.value()))) + self.thinking_enabled.setChecked(bool(chat.get("thinking_enabled", True))) + self._set_combo_text(self.reasoning_effort, str(chat.get("reasoning_effort", self.reasoning_effort.currentText()))) + self.reasoning_effort.setEnabled(self.thinking_enabled.isChecked()) + self.chat_max_tokens.setValue(int(chat.get("max_tokens", self.chat_max_tokens.value()))) + self.chat_temperature.setValue(float(chat.get("temperature", self.chat_temperature.value()))) + self.chat_top_p.setValue(float(chat.get("top_p", self.chat_top_p.value()))) + self.chat_repeat_penalty.setValue(float(chat.get("repeat_penalty", self.chat_repeat_penalty.value()))) + self.system_prompt.setPlainText(str(chat.get("system_prompt", ""))) + if hasattr(self, "coordinator_host"): + self.coordinator_host.setText(str(distributed.get("host", self.coordinator_host.text()))) + self.coordinator_port.setValue(int(distributed.get("port", self.coordinator_port.value()))) + self.coordinator_artifact_root.setText(str(distributed.get("artifact_root", self.coordinator_artifact_root.text()))) + self.coordinator_public_url.setText(str(distributed.get("public_url", self.coordinator_public_url.text()))) + self._update_tokenizer_strategy_controls() + self._update_training_mode_controls() + self._restore_artifact_status(data.get("artifacts", {})) + self.refresh_fine_tune_workflow() + + def _restore_artifact_status(self, artifacts: dict[str, Any]) -> None: + """Refresh top-bar and button state from saved or existing artifacts. + + Args: + artifacts: Saved artifact summary dictionary. + """ + + dataset_dir = Path(self.dataset_dir.text()) if self.dataset_dir.text().strip() else None + if dataset_dir and self._dataset_artifacts_exist(dataset_dir): + summary = self._read_json_if_exists(dataset_dir / "dataset_summary.json") or artifacts.get("dataset_summary") or {} + document_count = int(summary.get("document_count", 0) or 0) + token_count = int(summary.get("token_count", 0) or 0) + code_count = int(summary.get("code_sample_count", 0) or 0) + prose_count = int(summary.get("prose_sample_count", 0) or 0) + conversation_count = int(summary.get("conversation_sample_count", 0) or 0) + vocab_size = int(summary.get("tokenizer_vocab_size", 0) or 0) + self._update_dataset_quality_report(summary) + self.prepare_button.setText("DataSet Prepared") + self.dataset_progress.setValue(100) + if vocab_size: + self.auto_vocab_label.setText(f"{vocab_size:,}") + if code_count or prose_count or conversation_count: + self.dataset_status.setText( + f"Dataset: {code_count:,} code, {prose_count:,} prose, {conversation_count:,} chat, {token_count:,} tokens" + ) + elif document_count or token_count: + self.dataset_status.setText(f"Dataset: {document_count:,} files, {token_count:,} tokens") + else: + self.dataset_status.setText("Dataset: prepared") + version = summary.get("dataset_version", {}) + if isinstance(version, dict) and version.get("version_id"): + self.dataset_log.append(f"Dataset version: {version['version_id']}") + self.train_data_dir.setText(str(dataset_dir)) + self.dataset_log.append(f"Dataset already prepared: {dataset_dir}") + else: + self.prepare_button.setText("Prepare Dataset") + self.dataset_progress.setValue(0) + self.dataset_status.setText("Dataset: not prepared") + self.auto_vocab_label.setText("Auto after reading files") + self._reset_dataset_quality_report() + + model_dir = Path(self.model_dir.text()) if self.model_dir.text().strip() else None + if model_dir and (model_dir / "final_model.pt").exists(): + summary = self._read_json_if_exists(model_dir / "training_summary.json") or artifacts.get("training_summary") or {} + loss = summary.get("final_train_loss") + self.train_status.setText(f"Training: loss {float(loss):.4f}" if loss is not None else "Training: model ready") + self.export_model_dir.setText(str(model_dir)) + + export_dir = Path(self.export_dir.text()) if self.export_dir.text().strip() else None + if export_dir and export_dir.exists() and any(export_dir.iterdir()): + self.export_status.setText("Export: artifacts found") + + @staticmethod + def _dataset_artifacts_exist(dataset_dir: Path) -> bool: + """Return whether a dataset folder has the required prepared files. + + Args: + dataset_dir: Dataset folder. + + Returns: + True if required dataset artifacts exist. + """ + + if not dataset_dir.exists(): + return False + if not (dataset_dir / "tokenizer.json").exists(): + return False + has_npy_tokens = (dataset_dir / "train_tokens.npy").exists() and (dataset_dir / "val_tokens.npy").exists() + has_json_tokens = (dataset_dir / "train_tokens.json").exists() and (dataset_dir / "val_tokens.json").exists() + return has_npy_tokens or has_json_tokens + + @staticmethod + def _safe_project_name(project_name: str) -> str: + """Return a filesystem-safe project folder name. + + Args: + project_name: Raw user project name. + + Returns: + Safe folder name. + """ + + return re.sub(r"[^A-Za-z0-9_.-]+", "_", project_name).strip("._") or "MicroLLMProject" + + @staticmethod + def _read_json_if_exists(path: Path) -> Optional[Any]: + """Read a JSON file when it exists. + + Args: + path: JSON file path. + + Returns: + Parsed JSON or ``None``. + """ + + if not path.exists(): + return None + try: + return json.loads(path.read_text(encoding="utf-8")) + except Exception: + return None + + @staticmethod + def _set_combo_text(combo: QComboBox, text: str) -> None: + """Set combo text when the value exists. + + Args: + combo: Combo box to update. + text: Display text to select. + """ + + index = combo.findText(text) + if index >= 0: + combo.setCurrentIndex(index) + elif combo.isEditable(): + combo.setEditText(text) + + def _set_combo_by_data(self, combo: QComboBox, value: str, labels: dict[str, str]) -> None: + """Set a combo by internal saved value. + + Args: + combo: Combo box to update. + value: Internal saved value. + labels: Mapping from saved value to display label. + """ + + self._set_combo_text(combo, labels.get(value, value)) + + def _run_task( + self, + fn, + args, + on_finished, + log: QTextEdit, + progress_bar: QProgressBar, + with_progress: bool = False, + button: Optional[QPushButton] = None, + stop_button: Optional[QPushButton] = None, + busy_text: str = "Working", + task_kind: str = "", + isolate_process: bool = False, + ) -> None: + """Run a long task on a background thread. + + Args: + fn: Callable to execute. + args: Positional arguments for the callable. + on_finished: Slot called with the task result. + log: Log widget receiving progress messages. + progress_bar: Progress bar receiving percent updates. + with_progress: Whether to pass a progress callback to the task. + button: Optional button to disable while running. + stop_button: Optional stop button to enable while running. + busy_text: Button text shown while running. + task_kind: Optional notification stage key. + isolate_process: Run the task inside a child process. + """ + + if self.thread is not None: + QMessageBox.information(self, "Task running", "Please wait for the current task to finish.") + return + + LOGGER.info("Starting background task: %s", getattr(fn, "__name__", str(fn))) + self.active_task_kind = task_kind + if button: + self._set_button_busy(button, busy_text) + if stop_button: + stop_button.setEnabled(True) + self.active_stop_button = stop_button + + self.stop_event = Event() + self.progress_queue = Queue() + self.active_log = log + self.active_progress_bar = progress_bar + self.thread = QThread(self) + worker_class = ProcessTaskWorker if isolate_process else TaskWorker + self.worker = worker_class( + fn, + *args, + progress_queue=self.progress_queue, + with_progress=with_progress, + stop_event=self.stop_event, + ) + self.result_bridge = WorkerSignalBridge(self) + self.worker.moveToThread(self.thread) + self.thread.started.connect(self.worker.run) + self.worker.finished.connect(self.result_bridge.finished) + self.result_bridge.finished.connect(on_finished) + self.worker.finished.connect(self.worker.deleteLater) + self.worker.finished.connect(self.thread.quit) + self.worker.failed.connect(self.result_bridge.failed) + self.result_bridge.failed.connect(self._task_failed_from_worker) + self.worker.failed.connect(self.worker.deleteLater) + self.worker.failed.connect(self.thread.quit) + self.thread.finished.connect(self.thread.deleteLater) + self.thread.finished.connect(self._thread_finished) + self.progress_timer.start(100) + self.thread.start() + + @Slot(str) + def _task_failed_from_worker(self, message: str) -> None: + """Handle a worker failure on the UI thread. + + Args: + message: Error message emitted by the worker. + """ + + if self.active_log is None or self.active_progress_bar is None: + return + LOGGER.error("Background task failed: %s", message) + if self.active_task_kind == "chat": + self.chat_status.setText(f"Chat: load failed - {message}") + elif self.active_task_kind == "dataset_download": + self.external_dataset_version.setText(f"Download failed: {message}") + self.dataset_plan_progress.setVisible(False) + self._task_failed(message, self.active_log, self.active_progress_bar) + diff --git a/interface/main_window_part9.py b/interface/main_window_part9.py new file mode 100644 index 0000000..8d8c3d6 --- /dev/null +++ b/interface/main_window_part9.py @@ -0,0 +1,430 @@ +from __future__ import annotations + +# MainWindow implementation mixin. Runtime names are provided by app.py. +from typing import Any +from . import app as _app + +globals().update({name: value for name, value in vars(_app).items() if not name.startswith("__")}) + +class MainWindowPart9: + def stop_active_task(self) -> None: + """Request a graceful stop for the active background task.""" + + if self.stop_event is None: + return + LOGGER.info("Stop requested for active background task") + self.stop_event.set() + self._notify_failure("Stop requested", "The task is stopping at the next safe point.") + if self.active_log is not None: + self.active_log.append("Stop requested. Finishing the current safe point...") + if self.active_stop_button is not None: + self.active_stop_button.setEnabled(False) + if torch.cuda.is_available(): + try: + torch.cuda.empty_cache() + except Exception: + LOGGER.exception( + "Failed to empty CUDA cache in _thread_finished") + + @Slot() + def request_shutdown_from_signal(self) -> None: + """Handle Ctrl+C from a terminal without leaving Qt threads wedged.""" + + self.interrupt_count += 1 + if self.interrupt_count > 1: + os._exit(130) + if self.stop_event is not None: + self.stop_event.set() + if self.active_log is not None: + self.active_log.append("Interrupt received. Requesting stop...") + self.project_state.setText("Stopping") + if self.thread is None: + QApplication.quit() + return + QTimer.singleShot(3000, lambda: os._exit(130) if self.thread is not None else QApplication.quit()) + + def closeEvent(self, event: Any) -> None: + """Clean up background services before the window closes. + + Args: + event: Qt close event. + """ + + if self.thread is not None: + if self.stop_event is not None: + self.stop_event.set() + if self.active_log is not None: + self.active_log.append("Close requested. Stopping active task first...") + self.project_state.setText("Stopping") + LOGGER.info("Close requested while background task is running; waiting for task shutdown") + event.ignore() + QTimer.singleShot(500, self.close) + return + if self.coordinator_server is not None: + self.stop_coordinator_server() + super().closeEvent(event) + + def _handle_progress(self, event: object, log: QTextEdit, progress_bar: QProgressBar) -> None: + """Apply one progress event to UI widgets. + + Args: + event: Progress dictionary or message. + log: Log widget to append messages to. + progress_bar: Progress bar to update. + """ + + if isinstance(event, dict): + if event.get("type") == "chat_delta": + self._apply_chat_delta(event) + return + message = event.get("message") + percent = event.get("percent") + if log in (self.training_log, getattr(self, "fine_tune_log", None)): + self._update_training_metrics(event, update_fine_tune=log is getattr(self, "fine_tune_log", None)) + if message: + log.append(str(message)) + if log in (self.training_log, getattr(self, "fine_tune_log", None)) and hasattr(self, "live_log"): + self.live_log.append(str(message)) + if percent is not None: + progress_bar.setValue(max(0, min(100, int(percent)))) + if log in (self.training_log, getattr(self, "fine_tune_log", None)) and hasattr(self, "live_progress"): + self.live_progress.setValue(max(0, min(100, int(percent)))) + else: + log.append(str(event)) + + def _notify_progress(self, event: dict[str, Any]) -> None: + """Send throttled external progress notifications for long tasks. + + Args: + event: Progress event emitted by a worker. + """ + + if not self.active_task_kind or self.notification_manager is None: + return + if self.active_task_kind not in {"dataset", "training", "fine_tune"}: + return + title = { + "dataset": "Dataset preparation", + "training": "Model training", + "fine_tune": "Fine-tuning", + }[self.active_task_kind] + percent = event.get("percent") + self.notification_manager.notify_progress( + self.active_task_kind, + title, + self._notification_lines_from_event(event), + int(percent) if percent is not None else None, + ) + + def _notify_complete(self, stage_key: str, title: str, lines: list[str]) -> None: + """Send an external completion notification when configured. + + Args: + stage_key: Notification stage key. + title: User-facing title. + lines: Plain-text summary lines. + """ + + if self.notification_manager is not None: + self.notification_manager.notify_complete(stage_key, title, lines) + + def _notify_failure(self, title: str, message: str) -> None: + """Send an external failure or stop notification for the active task. + + Args: + title: User-facing title. + message: Failure details. + """ + + if self.active_task_kind and self.notification_manager is not None: + self.notification_manager.notify_failure(self.active_task_kind, title, message) + + def _notification_lines_from_event(self, event: dict[str, Any]) -> list[str]: + """Build compact notification text from a worker progress event. + + Args: + event: Progress event emitted by a worker. + + Returns: + Body lines for the notification message. + """ + + lines: list[str] = [] + if event.get("message"): + lines.append(str(event["message"])) + if "epoch" in event and "total_epochs" in event: + lines.append(f"Epoch: {event['epoch']}/{event['total_epochs']}") + if "step" in event and "total_steps" in event: + lines.append(f"Step: {event['step']}/{event['total_steps']}") + train_loss = self._finite_metric(event.get("train_loss")) + if train_loss is not None: + lines.append(f"Train loss: {float(train_loss):.4f}") + val_loss = self._finite_metric(event.get("val_loss")) + if val_loss is not None: + lines.append(f"Validation loss: {float(val_loss):.4f}") + learning_rate = self._finite_metric(event.get("learning_rate")) + if learning_rate is not None: + lines.append(f"Learning rate: {float(learning_rate):.2e}") + tokens_per_second = self._finite_metric(event.get("tokens_per_second")) + if tokens_per_second is not None: + lines.append(f"Speed: {float(tokens_per_second):.0f} tokens/sec") + eta_seconds = self._finite_metric(event.get("eta_seconds")) + if eta_seconds is not None: + lines.append(f"ETA: {self._format_duration(float(eta_seconds))}") + vram_allocated = self._finite_metric(event.get("vram_allocated_gb")) + vram_reserved = self._finite_metric(event.get("vram_reserved_gb")) + if vram_allocated is not None or vram_reserved is not None: + allocated = "-" if vram_allocated is None else f"{float(vram_allocated):.2f} GB" + reserved = "-" if vram_reserved is None else f"{float(vram_reserved):.2f} GB" + lines.append(f"VRAM: {allocated} allocated, {reserved} reserved") + return lines[:10] + + def _update_training_metrics(self, event: dict[str, Any], update_fine_tune: bool = False) -> None: + """Update training metric chips from a progress event. + + Args: + event: Progress event emitted by the training backend. + update_fine_tune: Whether to mirror metrics into the Fine-Tuning tab chips. + """ + + if "epoch" in event and "total_epochs" in event: + self.training_epoch_metric.setText(f"Epoch: {event['epoch']}/{event['total_epochs']}") + if update_fine_tune and hasattr(self, "fine_tune_epoch_metric"): + self.fine_tune_epoch_metric.setText(f"Epoch: {event['epoch']}/{event['total_epochs']}") + if "step" in event and "total_steps" in event: + self.training_step_metric.setText(f"Step: {event['step']}/{event['total_steps']}") + if update_fine_tune and hasattr(self, "fine_tune_step_metric"): + self.fine_tune_step_metric.setText(f"Step: {event['step']}/{event['total_steps']}") + train_loss = self._finite_metric(event.get("train_loss")) + if train_loss is not None: + self.training_loss_metric.setText(f"Train loss: {float(train_loss):.4f}") + if update_fine_tune and hasattr(self, "fine_tune_loss_metric"): + self.fine_tune_loss_metric.setText(f"Train loss: {float(train_loss):.4f}") + val_loss = self._finite_metric(event.get("val_loss")) + if val_loss is not None: + self.training_val_metric.setText(f"Val loss: {float(val_loss):.4f}") + if update_fine_tune and hasattr(self, "fine_tune_val_metric"): + self.fine_tune_val_metric.setText(f"Val loss: {float(val_loss):.4f}") + step = event.get("step") + if step is not None and (train_loss is not None or val_loss is not None): + step_int_for_loss = int(step) + self.loss_chart.add_metrics(step_int_for_loss, train_loss, val_loss) + self._update_training_health(step_int_for_loss, train_loss, val_loss) + if step is None: + return + step_int = int(step) + self._record_live_metric(event) + learning_rate = self._finite_metric(event.get("learning_rate")) + grad_norm = self._finite_metric(event.get("grad_norm")) + weight_norm = self._finite_metric(event.get("weight_norm")) + update_ratio = self._finite_metric(event.get("update_ratio")) + tokens_per_second = self._finite_metric(event.get("tokens_per_second")) + samples_per_second = self._finite_metric(event.get("samples_per_second")) + vram_allocated = self._finite_metric(event.get("vram_allocated_gb")) + vram_reserved = self._finite_metric(event.get("vram_reserved_gb")) + gpu_memory = self._finite_metric(event.get("gpu_memory_percent")) + system_cpu = self._finite_metric(event.get("system_cpu_percent")) + system_ram = self._finite_metric(event.get("system_ram_percent")) + data_workers = event.get("data_loader_workers") + eta_seconds = self._finite_metric(event.get("eta_seconds")) + if learning_rate is not None: + self.training_lr_metric.setText(f"LR: {float(learning_rate):.2e}") + if update_fine_tune and hasattr(self, "fine_tune_lr_metric"): + self.fine_tune_lr_metric.setText(f"LR: {float(learning_rate):.2e}") + if grad_norm is not None: + self.training_grad_metric.setText(f"Grad: {float(grad_norm):.3f}") + if update_fine_tune and hasattr(self, "fine_tune_grad_metric"): + self.fine_tune_grad_metric.setText(f"Grad: {float(grad_norm):.3f}") + if tokens_per_second is not None: + self.training_speed_metric.setText(f"Speed: {float(tokens_per_second):.0f} tok/s") + if update_fine_tune and hasattr(self, "fine_tune_speed_metric"): + self.fine_tune_speed_metric.setText(f"Speed: {float(tokens_per_second):.0f} tok/s") + if vram_allocated is not None: + self.training_vram_metric.setText(f"VRAM: {float(vram_allocated):.2f} GB") + if eta_seconds is not None: + self.training_eta_metric.setText(f"ETA: {self._format_duration(float(eta_seconds))}") + if update_fine_tune and hasattr(self, "fine_tune_eta_metric"): + self.fine_tune_eta_metric.setText(f"ETA: {self._format_duration(float(eta_seconds))}") + if learning_rate is not None or grad_norm is not None: + self.optimization_chart.add_values(step_int, learning_rate, grad_norm) + if weight_norm is not None or update_ratio is not None: + self.stability_chart.add_values(step_int, weight_norm, update_ratio) + if tokens_per_second is not None or samples_per_second is not None: + self.throughput_chart.add_values(step_int, tokens_per_second, samples_per_second) + if vram_allocated is not None or vram_reserved is not None: + self.memory_chart.add_values(step_int, vram_allocated, vram_reserved) + if hasattr(self, "live_epoch_metric"): + self._update_live_training_metrics( + step_int, + event, + train_loss, + learning_rate, + grad_norm, + update_ratio, + tokens_per_second, + samples_per_second, + vram_allocated, + vram_reserved, + gpu_memory, + system_cpu, + system_ram, + data_workers, + ) + + def _update_training_health( + self, + step: int, + train_loss: Optional[float], + val_loss: Optional[float], + ) -> None: + """Update the training health advisor from recent loss values. + + Args: + step: Current optimizer step. + train_loss: Latest training loss. + val_loss: Latest validation loss. + """ + + self.training_health_points.append((step, train_loss, val_loss)) + self.training_health_points = self.training_health_points[-12:] + latest_train = next((item[1] for item in reversed(self.training_health_points) if item[1] is not None), None) + latest_val = next((item[2] for item in reversed(self.training_health_points) if item[2] is not None), None) + val_points = [(item[0], item[2]) for item in self.training_health_points if item[2] is not None] + if latest_train is None and latest_val is None: + label = "Health: collecting" + tip = "Waiting for train and validation loss." + elif latest_train is not None and latest_val is not None and latest_train < 0.2 and latest_val > max(2.0, latest_train * 8.0): + label = "Health: validation gap" + tip = "Training loss is very low while validation loss is high. Check overfitting, validation split, tokenizer match, or eval settings." + elif len(val_points) >= 3 and val_points[-1][1] > val_points[-2][1] > val_points[-3][1]: + label = "Health: overfitting?" + tip = "Validation loss has increased for three checks. Consider stopping, reducing epochs, or improving validation data." + elif latest_train is not None and (latest_train > 20.0 or not math.isfinite(latest_train)): + label = "Health: diverging" + tip = "Training loss is unstable or extremely high. Lower learning rate and check gradients/data." + elif latest_val is not None and latest_val > 10.0: + label = "Health: high val loss" + tip = "Validation loss is high. This may be early training, a difficult validation split, or a dataset/tokenizer mismatch." + elif latest_train is not None and latest_val is not None and latest_val <= latest_train * 1.8: + label = "Health: stable" + tip = "Training and validation loss are reasonably close." + else: + label = "Health: watching" + tip = "Collecting more loss points before making a stronger diagnosis." + self.training_health_metric.setText(label) + self._tip(self.training_health_metric, tip) + + @staticmethod + def _finite_metric(value: Any) -> Optional[float]: + """Return a finite metric value or ``None``. + + Args: + value: Raw metric value. + + Returns: + Finite float, or ``None`` when invalid. + """ + + if value is None: + return None + try: + numeric = float(value) + except (TypeError, ValueError): + return None + return numeric if math.isfinite(numeric) else None + + def _update_live_training_metrics( + self, + step: int, + event: dict[str, Any], + train_loss: Optional[float], + learning_rate: Optional[float], + grad_norm: Optional[float], + update_ratio: Optional[float], + tokens_per_second: Optional[float], + samples_per_second: Optional[float], + vram_allocated: Optional[float], + vram_reserved: Optional[float], + gpu_memory: Optional[float], + system_cpu: Optional[float], + system_ram: Optional[float], + data_workers: Optional[int], + ) -> None: + """Update live tracker widgets from one training progress event. + + Args: + step: Current optimizer step. + event: Progress event emitted by training. + train_loss: Latest training loss. + learning_rate: Current learning rate. + grad_norm: Current gradient norm. + update_ratio: Current parameter update ratio. + tokens_per_second: Current token throughput. + samples_per_second: Current sample throughput. + vram_allocated: Current CUDA allocated memory in GB. + vram_reserved: Current CUDA reserved memory in GB. + gpu_memory: Current GPU memory pressure percentage. + system_cpu: Current system CPU utilization percentage. + system_ram: Current system RAM utilization percentage. + data_workers: CPU data-loader worker count. + """ + + total_steps = event.get("total_steps") + if "epoch" in event and "total_epochs" in event: + self.live_epoch_metric.setText(f"Epoch: {event['epoch']}/{event['total_epochs']}") + if total_steps: + self.live_step_metric.setText(f"Step: {step:,}/{int(total_steps):,}") + data_percent = min(100.0, max(0.0, (step / max(1, int(total_steps))) * 100.0)) + self.live_data_metric.setText(f"Data: {data_percent:.1f}%") + self.live_progress.setValue(int(data_percent)) + else: + self.live_step_metric.setText(f"Step: {step:,}") + if tokens_per_second is not None: + self.live_tokens_metric.setText(f"Tokens/sec: {float(tokens_per_second):,.0f}") + if train_loss is not None: + self.live_loss_metric.setText(f"Loss: {float(train_loss):.4f}") + if learning_rate is not None: + self.live_lr_metric.setText(f"LR: {float(learning_rate):.2e}") + sample_text = str(event.get("sample_text") or "").strip() + if sample_text: + self.live_sample_text.setText(f"Training text: {self._compact_preview_text(sample_text, 220)}") + self.live_layer_status.setText(f"Layers: {self.n_layer.value()}") + self.live_head_status.setText(f"Heads: {self.n_head.value()}") + self.live_hidden_status.setText(f"Hidden size: {self.n_embd.value()}") + self.live_batch_status.setText(f"Batch size: {self.batch_size.value()}") + self.live_context_status.setText(f"Context: {self.train_context_length.value()}") + self.live_device_status.setText(f"Device: {self.device.currentText()}") + self.live_worker_status.setText(f"CPU workers: {data_workers if data_workers is not None else self.data_loader_workers.value()}") + self._set_meter(self.live_cpu_bar, "CPU", system_cpu if system_cpu is not None else self._system_cpu_value()) + self._set_meter(self.live_gpu_bar, "GPU memory", gpu_memory) + if vram_allocated is not None or vram_reserved is not None: + allocated = float(vram_allocated or 0.0) + reserved = float(vram_reserved or 0.0) + reserved_percent = None + if self.device.currentText().startswith("cuda") and torch.cuda.is_available(): + try: + _, total_vram = torch.cuda.mem_get_info() + reserved_percent = min(100.0, 100.0 * reserved * (1024 ** 3) / max(total_vram, 1)) + except Exception: + reserved_percent = None + self._set_meter(self.live_vram_bar, "VRAM reserved", reserved_percent) + self.live_vram_label.setText(f"VRAM reserved: {reserved:.2f} GB ({allocated:.2f} GB active)") + self._set_meter(self.live_ram_bar, "System RAM", system_ram if system_ram is not None else self._system_ram_value()) + latest_loss = float(train_loss) if train_loss is not None else None + self.live_flow.set_state(self.n_layer.value(), self.n_head.value(), step, latest_loss) + self.live_prediction_chart.update_distribution(step, latest_loss) + self.live_attention_chart.update_heatmap(step, grad_norm) + self.live_activation_chart.update_histogram(step, tokens_per_second) + self.live_gradient_chart.update_flow(self.n_layer.value(), grad_norm, step) + + def _system_ram_value(self) -> Optional[float]: + """Read system RAM utilization for live telemetry. + + Returns: + System RAM percentage, or None when unavailable. + """ + + if psutil is None: + return None + return float(psutil.virtual_memory().percent) + diff --git a/llm_trainer/ui/markdown_renderer.py b/interface/markdown_renderer.py similarity index 99% rename from llm_trainer/ui/markdown_renderer.py rename to interface/markdown_renderer.py index 4990087..8516881 100644 --- a/llm_trainer/ui/markdown_renderer.py +++ b/interface/markdown_renderer.py @@ -175,7 +175,7 @@ def code_block_html(label: str, highlighted_html: str, raw_code: str) -> str: return ( "
" f"
{escape_html(label or 'Code')}" - f"⧉ Copy
" + f"Copy
" f"
{highlighted_html}
" "" ) diff --git a/llm_trainer/ui/micro_llm_creator_lightning.ico b/interface/micro_llm_creator_lightning.ico similarity index 100% rename from llm_trainer/ui/micro_llm_creator_lightning.ico rename to interface/micro_llm_creator_lightning.ico diff --git a/interface/startup.py b/interface/startup.py new file mode 100644 index 0000000..a8613d8 --- /dev/null +++ b/interface/startup.py @@ -0,0 +1,461 @@ +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[1] / "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] = ["") + 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/interface/startup_splash.py similarity index 88% rename from llm_trainer/ui/startup_splash.py rename to interface/startup_splash.py index 19d09ac..a8fa826 100644 --- a/llm_trainer/ui/startup_splash.py +++ b/interface/startup_splash.py @@ -26,8 +26,8 @@ def __init__(self) -> None: logo = QLabel() logo.setFixedSize(128, 128) logo_candidates = [ + Path(__file__).resolve().parents[1] / "drunken_bot_logo_small.png", 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") @@ -82,6 +82,20 @@ def set_checks(self, checks: list[str]) -> None: 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)) @@ -103,7 +117,7 @@ def _render_checks(self) -> None: for label in self._check_order: state = self._checks.get(label, "pending") escaped = html.escape(label) - marker = {"done": "✓", "running": "●", "failed": "✗"}.get(state, "•") + 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("") diff --git a/interface/startup_validation.py b/interface/startup_validation.py new file mode 100644 index 0000000..29f2df8 --- /dev/null +++ b/interface/startup_validation.py @@ -0,0 +1,119 @@ +"""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[1] + tests_root = repo_root / "tests" + required_modules = [ + "PySide6", "torch", "PyPDF2", "numpy", "tokenizers", + "engine.dataset_build", "engine.training", "interface.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/interface/styles.qss similarity index 100% rename from llm_trainer/ui/styles.qss rename to interface/styles.qss diff --git a/llm_trainer/ui/tabs/__init__.py b/interface/tabs/__init__.py similarity index 100% rename from llm_trainer/ui/tabs/__init__.py rename to interface/tabs/__init__.py diff --git a/llm_trainer/ui/tabs/benchmark_tab.py b/interface/tabs/benchmark_tab.py similarity index 98% rename from llm_trainer/ui/tabs/benchmark_tab.py rename to interface/tabs/benchmark_tab.py index 0eab9a2..278ce0b 100644 --- a/llm_trainer/ui/tabs/benchmark_tab.py +++ b/interface/tabs/benchmark_tab.py @@ -13,7 +13,7 @@ QWidget, ) -from llm_trainer.evaluation import DEFAULT_BENCHMARK_PROMPTS +from engine.evaluation import DEFAULT_BENCHMARK_PROMPTS def build_benchmark_tab(window) -> QWidget: @@ -85,3 +85,4 @@ def build_benchmark_tab(window) -> QWidget: window.benchmark_progress = window._thin_progress() outer.addWidget(window.benchmark_progress) return page + diff --git a/llm_trainer/ui/tabs/chat_tab.py b/interface/tabs/chat_tab.py similarity index 99% rename from llm_trainer/ui/tabs/chat_tab.py rename to interface/tabs/chat_tab.py index 00fb5dc..1632d49 100644 --- a/llm_trainer/ui/tabs/chat_tab.py +++ b/interface/tabs/chat_tab.py @@ -15,7 +15,7 @@ QWidget, ) -from llm_trainer.ui.chat_widgets import ChatInputEdit +from interface.chat_widgets import ChatInputEdit def build_chat_tab(window) -> QWidget: @@ -171,3 +171,4 @@ def build_chat_tab(window) -> QWidget: layout.addWidget(window.chat_progress) window._update_chat_backend_controls() return page + diff --git a/llm_trainer/ui/tabs/dataset_plan_tab.py b/interface/tabs/dataset_plan_tab.py similarity index 99% rename from llm_trainer/ui/tabs/dataset_plan_tab.py rename to interface/tabs/dataset_plan_tab.py index 8f5b283..196f550 100644 --- a/llm_trainer/ui/tabs/dataset_plan_tab.py +++ b/interface/tabs/dataset_plan_tab.py @@ -24,7 +24,7 @@ QWidgetAction, ) -from llm_trainer.conversation_datasets import CONVERSATION_DATASET_PRESETS +from engine.conversation_datasets import CONVERSATION_DATASET_PRESETS @@ -455,3 +455,4 @@ def populate_default_data_tree(window: Any, root: Path) -> None: 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/interface/tabs/dataset_tab.py similarity index 99% rename from llm_trainer/ui/tabs/dataset_tab.py rename to interface/tabs/dataset_tab.py index 7f481f5..6287428 100644 --- a/llm_trainer/ui/tabs/dataset_tab.py +++ b/interface/tabs/dataset_tab.py @@ -21,8 +21,8 @@ QWidgetAction, ) -from llm_trainer.conversation_datasets import CONVERSATION_DATASET_PRESETS -from llm_trainer.ui.charts import DatasetBarChartWidget +from engine.conversation_datasets import CONVERSATION_DATASET_PRESETS +from interface.charts import DatasetBarChartWidget def build_dataset_tab(window) -> QWidget: @@ -297,4 +297,4 @@ def build_dataset_tab(window) -> QWidget: 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 + return page diff --git a/llm_trainer/ui/tabs/export_tab.py b/interface/tabs/export_tab.py similarity index 99% rename from llm_trainer/ui/tabs/export_tab.py rename to interface/tabs/export_tab.py index 6e02eaf..3a07eef 100644 --- a/llm_trainer/ui/tabs/export_tab.py +++ b/interface/tabs/export_tab.py @@ -116,3 +116,4 @@ def build_export_tab(window) -> QWidget: 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/interface/tabs/fine_tuning_tab.py similarity index 99% rename from llm_trainer/ui/tabs/fine_tuning_tab.py rename to interface/tabs/fine_tuning_tab.py index 755c2e2..26d89f2 100644 --- a/llm_trainer/ui/tabs/fine_tuning_tab.py +++ b/interface/tabs/fine_tuning_tab.py @@ -252,4 +252,4 @@ def _single_widget_layout(widget: QWidget) -> QVBoxLayout: layout = QVBoxLayout() layout.addWidget(widget, 1) - return layout \ No newline at end of file + return layout diff --git a/llm_trainer/ui/tabs/job_manager_tab.py b/interface/tabs/job_manager_tab.py similarity index 99% rename from llm_trainer/ui/tabs/job_manager_tab.py rename to interface/tabs/job_manager_tab.py index 6c6d204..4a88294 100644 --- a/llm_trainer/ui/tabs/job_manager_tab.py +++ b/interface/tabs/job_manager_tab.py @@ -289,3 +289,4 @@ def set_table_rows(table: QTableWidget, rows: list[list[str]]) -> None: item = QTableWidgetItem(value) item.setToolTip(value) table.setItem(row_index, column_index, item) + diff --git a/llm_trainer/ui/tabs/live_tab.py b/interface/tabs/live_tab.py similarity index 95% rename from llm_trainer/ui/tabs/live_tab.py rename to interface/tabs/live_tab.py index 5d7a1ba..d758c38 100644 --- a/llm_trainer/ui/tabs/live_tab.py +++ b/interface/tabs/live_tab.py @@ -12,8 +12,8 @@ QWidget, ) -from llm_trainer.ui.charts import LossChartWidget -from llm_trainer.ui.live_widgets import ( +from interface.charts import LossChartWidget +from interface.live_widgets import ( LiveDistributionWidget, LiveGradientFlowWidget, LiveHeatmapWidget, @@ -47,7 +47,7 @@ def build_live_training_tab(window) -> QWidget: header = QHBoxLayout() title = window._page_title("Model Training Live") - live_badge = QLabel("● 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.") @@ -84,12 +84,12 @@ def build_live_training_tab(window) -> QWidget: 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: -") + 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, diff --git a/llm_trainer/ui/tabs/training_tab.py b/interface/tabs/training_tab.py similarity index 99% rename from llm_trainer/ui/tabs/training_tab.py rename to interface/tabs/training_tab.py index 8d69bb9..92b3e77 100644 --- a/llm_trainer/ui/tabs/training_tab.py +++ b/interface/tabs/training_tab.py @@ -425,4 +425,4 @@ def build_training_tab(window) -> QWidget: 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 + return page diff --git a/interface/wiki_download.py b/interface/wiki_download.py new file mode 100644 index 0000000..13b6de2 --- /dev/null +++ b/interface/wiki_download.py @@ -0,0 +1,19 @@ +"""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/interface/wiki_download_actions.py b/interface/wiki_download_actions.py new file mode 100644 index 0000000..2174d72 --- /dev/null +++ b/interface/wiki_download_actions.py @@ -0,0 +1,312 @@ +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/interface/wiki_download_backend.py b/interface/wiki_download_backend.py new file mode 100644 index 0000000..a27d458 --- /dev/null +++ b/interface/wiki_download_backend.py @@ -0,0 +1,117 @@ +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/interface/wiki_download_layout.py b/interface/wiki_download_layout.py new file mode 100644 index 0000000..f1def46 --- /dev/null +++ b/interface/wiki_download_layout.py @@ -0,0 +1,355 @@ +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/interface/wiki_download_processing.py b/interface/wiki_download_processing.py new file mode 100644 index 0000000..44b5092 --- /dev/null +++ b/interface/wiki_download_processing.py @@ -0,0 +1,243 @@ +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/interface/wiki_download_style.py b/interface/wiki_download_style.py new file mode 100644 index 0000000..ebc886e --- /dev/null +++ b/interface/wiki_download_style.py @@ -0,0 +1,206 @@ +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/interface/wiki_download_worker.py b/interface/wiki_download_worker.py new file mode 100644 index 0000000..192ea5f --- /dev/null +++ b/interface/wiki_download_worker.py @@ -0,0 +1,136 @@ +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/ui/workers.py b/interface/workers.py similarity index 97% rename from llm_trainer/ui/workers.py rename to interface/workers.py index 972b6c8..f496082 100644 --- a/llm_trainer/ui/workers.py +++ b/interface/workers.py @@ -104,6 +104,13 @@ def _should_stop(self) -> bool: 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.""" @@ -206,4 +213,4 @@ def _drain_child_progress(self, child_progress_queue: mp.Queue) -> None: event = child_progress_queue.get_nowait() except Exception: break - self.progress_queue.put(event) \ No newline at end of file + self.progress_queue.put(event) diff --git a/llm_trainer/__init__.py b/llm_trainer/__init__.py deleted file mode 100644 index 67c1e62..0000000 --- a/llm_trainer/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -"""Backend package for preparing data and training small language models.""" - -__all__ = [ - "config", - "data", - "dataset_build", - "dataset_preview", - "dataset_mixture", - "export", - "model", - "resume_checks", - "services", - "tokenizer", - "training", - "training_orchestrator", -] diff --git a/llm_trainer/app_logging.py b/llm_trainer/app_logging.py deleted file mode 100644 index cfd0e53..0000000 --- a/llm_trainer/app_logging.py +++ /dev/null @@ -1,116 +0,0 @@ -from __future__ import annotations - -import logging -import faulthandler -from logging.handlers import RotatingFileHandler -from pathlib import Path -import sys -import threading -import traceback -from typing import Optional - - -DEFAULT_LOG_DIR = Path.home() / ".drunkenbot_ide" / "logs" -DEFAULT_LOG_PATH = DEFAULT_LOG_DIR / "drunkenbot_ide.log" -_FAULT_LOG_HANDLE: Optional[object] = None -_CONFIGURED_PATH: Optional[Path] = None - - -def setup_logging(log_path: Optional[Path] = None) -> Path: - """Configure console and rotating file logging for the desktop app. - - Args: - log_path: Optional explicit log file path. - - Returns: - Path to the active log file. - """ - - global _CONFIGURED_PATH - active_path = log_path or DEFAULT_LOG_PATH - active_path.parent.mkdir(parents=True, exist_ok=True) - root_logger = logging.getLogger() - root_logger.setLevel(logging.INFO) - - if not any(getattr(handler, "_micro_llm_handler", False) for handler in root_logger.handlers): - formatter = logging.Formatter("%(asctime)s | %(levelname)s | %(name)s | %(message)s") - file_handler = RotatingFileHandler(active_path, maxBytes=5_000_000, backupCount=5, encoding="utf-8") - file_handler.setFormatter(formatter) - file_handler._micro_llm_handler = True # type: ignore[attr-defined] - root_logger.addHandler(file_handler) - - console_handler = logging.StreamHandler(sys.stdout) - console_handler.setFormatter(formatter) - console_handler._micro_llm_handler = True # type: ignore[attr-defined] - root_logger.addHandler(console_handler) - - logging.captureWarnings(True) - sys.excepthook = _log_uncaught_exception - threading.excepthook = _log_thread_exception - _enable_fault_logging(active_path) - for logger_name in ("datasets", "huggingface_hub", "urllib3", "filelock", "pyarrow"): - logging.getLogger(logger_name).setLevel(logging.INFO) - if _CONFIGURED_PATH is None: - _CONFIGURED_PATH = active_path - logging.getLogger(__name__).info("Logging initialized: %s", active_path) - return active_path - - -def qt_message_handler(mode: object, context: object, message: str) -> None: - """Route Qt runtime messages to the app log. - - Args: - mode: Qt message type. - context: Qt message context. - message: Message text. - """ - - logger = logging.getLogger("qt") - file_name = getattr(context, "file", "") or "" - line = getattr(context, "line", 0) or 0 - location = f" ({file_name}:{line})" if file_name else "" - logger.warning("%s%s", message, location) - - -def _log_uncaught_exception(exc_type: type[BaseException], exc_value: BaseException, exc_tb: object) -> None: - """Log exceptions that reach the Python top level. - - Args: - exc_type: Exception type. - exc_value: Exception instance. - exc_tb: Traceback object. - """ - - if issubclass(exc_type, KeyboardInterrupt): - sys.__excepthook__(exc_type, exc_value, exc_tb) - return - formatted = "".join(traceback.format_exception(exc_type, exc_value, exc_tb)) - logging.getLogger(__name__).critical("Uncaught exception:\n%s", formatted) - - -def _log_thread_exception(args: threading.ExceptHookArgs) -> None: - """Log uncaught exceptions from Python threads. - - Args: - args: Thread exception hook arguments. - """ - - if args.exc_type is None or args.exc_value is None: - return - formatted = "".join(traceback.format_exception(args.exc_type, args.exc_value, args.exc_traceback)) - logging.getLogger(__name__).critical("Uncaught thread exception in %s:\n%s", args.thread.name if args.thread else "-", formatted) - - -def _enable_fault_logging(active_path: Path) -> None: - """Enable crash dumps for native faults when possible. - - Args: - active_path: Main application log path. - """ - - global _FAULT_LOG_HANDLE - if _FAULT_LOG_HANDLE is not None: - return - fault_path = active_path.with_name("drunkenbot_ide_faults.log") - _FAULT_LOG_HANDLE = fault_path.open("a", encoding="utf-8") - faulthandler.enable(file=_FAULT_LOG_HANDLE, all_threads=True) diff --git a/llm_trainer/backends/__init__.py b/llm_trainer/backends/__init__.py deleted file mode 100644 index 784de2e..0000000 --- a/llm_trainer/backends/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -from __future__ import annotations - -from .base import TrainerBackend -from .local_backend import LocalTrainerBackend -from .registry import BackendRegistry, DEFAULT_BACKEND_REGISTRY - -__all__ = ["BackendRegistry", "DEFAULT_BACKEND_REGISTRY", "LocalTrainerBackend", "TrainerBackend"] diff --git a/llm_trainer/backends/base.py b/llm_trainer/backends/base.py deleted file mode 100644 index da47b39..0000000 --- a/llm_trainer/backends/base.py +++ /dev/null @@ -1,33 +0,0 @@ -from __future__ import annotations - -from typing import Any, Callable, Optional, Protocol - -from llm_trainer.contracts import TrainingJobSpec -from llm_trainer.training import TrainingResult - - -ProgressCallback = Callable[[Any], None] -StopCallback = Callable[[], bool] - - -class TrainerBackend(Protocol): - """Protocol implemented by training backends.""" - - name: str - - def run( - self, - job: TrainingJobSpec, - progress: Optional[ProgressCallback] = None, - should_stop: Optional[StopCallback] = None, - ) -> TrainingResult: - """Run a training job. - - Args: - job: Backend-neutral training job spec. - progress: Optional progress callback. - should_stop: Optional cooperative cancellation callback. - - Returns: - Training result. - """ diff --git a/llm_trainer/backends/local_backend.py b/llm_trainer/backends/local_backend.py deleted file mode 100644 index a5c3a6b..0000000 --- a/llm_trainer/backends/local_backend.py +++ /dev/null @@ -1,58 +0,0 @@ -from __future__ import annotations - -import json -from typing import Optional - -from llm_trainer.backends.base import ProgressCallback, StopCallback -from llm_trainer.contracts import JobStatus, TrainingJobSpec -from llm_trainer.training_orchestrator import train_from_dataset -from llm_trainer.training import TrainingResult - - -class LocalTrainerBackend: - """Training backend that runs jobs in the current Python process.""" - - name = "local" - - def run( - self, - job: TrainingJobSpec, - progress: Optional[ProgressCallback] = None, - should_stop: Optional[StopCallback] = None, - ) -> TrainingResult: - """Run a local training job. - - Args: - job: Backend-neutral training job specification. - progress: Optional progress callback. - should_stop: Optional cooperative cancellation callback. - - Returns: - Training result. - """ - - job.status = JobStatus.RUNNING - try: - job.artifacts.output_dir.mkdir(parents=True, exist_ok=True) - job_manifest_path = job.artifacts.output_dir / "training_job.json" - job_manifest_path.write_text(json.dumps(job.to_jsonable(), indent=2), encoding="utf-8") - if progress: - progress( - { - "message": f"Local backend started job {job.job_id}", - "job_id": job.job_id, - "backend": self.name, - } - ) - result = train_from_dataset( - job.dataset.dataset_dir, - job.model.config, - job.training, - progress=progress, - should_stop=should_stop, - ) - job.status = JobStatus.CANCELLED if result.stopped else JobStatus.COMPLETED - return result - except Exception: - job.status = JobStatus.FAILED - raise diff --git a/llm_trainer/backends/registry.py b/llm_trainer/backends/registry.py deleted file mode 100644 index 621018b..0000000 --- a/llm_trainer/backends/registry.py +++ /dev/null @@ -1,47 +0,0 @@ -from __future__ import annotations - -from llm_trainer.backends.base import TrainerBackend -from llm_trainer.backends.local_backend import LocalTrainerBackend -from llm_trainer.contracts import BackendKind - - -class BackendRegistry: - """Registry for training backends.""" - - def __init__(self) -> None: - """Create a registry with built-in local backend.""" - - self._backends: dict[BackendKind, TrainerBackend] = { - BackendKind.LOCAL: LocalTrainerBackend(), - } - - def register(self, kind: BackendKind, backend: TrainerBackend) -> None: - """Register a backend implementation. - - Args: - kind: Backend kind. - backend: Backend implementation. - """ - - self._backends[kind] = backend - - def get(self, kind: BackendKind) -> TrainerBackend: - """Return a backend by kind. - - Args: - kind: Backend kind. - - Returns: - Backend implementation. - - Raises: - ValueError: If the backend kind is not registered. - """ - - backend = self._backends.get(kind) - if backend is None: - raise ValueError(f"No training backend registered for {kind.value}") - return backend - - -DEFAULT_BACKEND_REGISTRY = BackendRegistry() diff --git a/llm_trainer/chat_test.py b/llm_trainer/chat_test.py deleted file mode 100644 index eaaf03c..0000000 --- a/llm_trainer/chat_test.py +++ /dev/null @@ -1,74 +0,0 @@ - -TOKENIZER_PATH = r"F:\Micro_LLM_Projects\peach\models\tokenizer.json" -MODEL_PATH = r"F:\Micro_LLM_Projects\peach\models\final_model.pt" - - - -from pathlib import Path - -from llm_trainer.microgpt_chat import load_microgpt_chat_session - -# --------------------------------------------------------- -# Configuration -# --------------------------------------------------------- - -# Point this to your MODEL FOLDER, not final_model.pt -# Example: -# E:\AI_Projects\Models\DrunkenBot - - -# "auto" = CUDA if available, otherwise CPU -DEVICE = "auto" - -# --------------------------------------------------------- -# Load model -# --------------------------------------------------------- - -print("Loading model...") - -session = load_microgpt_chat_session( - MODEL_PATH, - device=DEVICE -) - -print("Model loaded successfully!") -print(session.runtime_summary) -print() - -print("Type 'exit' to quit.\n") - -# --------------------------------------------------------- -# Chat Loop -# --------------------------------------------------------- - -while True: - - user = input("You: ").strip() - - if not user: - continue - - if user.lower() in ("exit", "quit"): - break - - result = session.generate_stream( - prompt=user, - system_prompt="", - max_tokens=1024, - temperature=0.7, - top_p=0.9, - repeat_penalty=1.1, - reasoning_effort="Balanced", - thinking_enabled=True, - ) - - print("\nDrunkenBot:") - print(result["reply"]) - print() - - print( - f"[{result['token_count']} tokens | " - f"{result['tokens_per_second']:.2f} tok/s | " - f"{result['elapsed_seconds']:.2f} sec]" - ) - print("-" * 60) \ No newline at end of file diff --git a/llm_trainer/cli.py b/llm_trainer/cli.py deleted file mode 100644 index 1f1c7a6..0000000 --- a/llm_trainer/cli.py +++ /dev/null @@ -1,403 +0,0 @@ -from __future__ import annotations - -import argparse -import json -import os -from pathlib import Path - -import torch - -from .config import DatasetConfig, ModelConfig, TrainingConfig -from .coordinator import run_coordinator_api -from .coordinator.artifacts import create_job_artifact_bundle -from .contracts import BackendKind -from .contracts.jobs import RuntimeSpec, TrainingJobSpec -from .evaluation import evaluate_checkpoint, normalize_prompts -from .export import export_hf_microgpt_package -from .dataset_build import build_dataset -from .training_orchestrator import train_from_dataset -from .tokenizer import load_tokenizer -from .worker import WorkerClientConfig, run_worker_client - - -def prepare(args: argparse.Namespace) -> None: - """Prepare a dataset from command-line arguments. - - Args: - args: Parsed command-line arguments for the prepare command. - """ - - def print_progress(event: object) -> None: - """Print a progress event in CLI-friendly form. - - Args: - event: Progress dictionary or message. - """ - - if isinstance(event, dict): - message = event.get("message") - percent = event.get("percent") - prefix = f"[{percent:>3}%] " if percent is not None else "" - if message: - print(prefix + str(message)) - else: - print(event) - - config = DatasetConfig( - input_dir=Path(args.input_dir), - output_dir=Path(args.output_dir), - vocab_size=args.vocab_size, - min_frequency=args.min_frequency, - context_length=args.context_length, - validation_split=args.validation_split, - lowercase=False, - max_workers=args.max_workers, - code_training_mode=args.code_training_mode, - include_prose=not args.exclude_prose, - include_source_code=not args.exclude_source_code, - extract_code_blocks=not args.no_extract_code_blocks, - preserve_indentation=not args.no_preserve_indentation, - generate_instruction_samples=not args.no_instruction_samples, - reasoning_sample_mode=args.reasoning_sample_mode, - prepare_mode=args.prepare_mode, - tokenizer_strategy=args.tokenizer_strategy, - tokenizer_path=Path(args.tokenizer_path) if args.tokenizer_path else None, - dataset_stage=args.dataset_stage, - conversation_datasets=[item.strip() for item in args.conversation_datasets.split(",") if item.strip()], - conversation_sample_limit=args.conversation_sample_limit, - fast_scan_mode=args.fast_scan_mode, - fast_scan_sample_bytes=args.fast_scan_sample_bytes, - strict_duplicate_verification=args.strict_duplicate_verification, - ) - result = build_dataset(config, progress=print_progress) - print( - f"Documents: {result.document_count} | Characters: {result.character_count} | " - f"Tokens: {result.token_count} | Vocab: {result.vocab_size}" - ) - print(f"Cache: reused {result.cached_file_count} file(s) | processed {result.processed_file_count} file(s)") - - -def train(args: argparse.Namespace) -> None: - """Train a model from command-line arguments. - - Args: - args: Parsed command-line arguments for the train command. - """ - - data_dir = Path(args.data_dir) - tokenizer = load_tokenizer(data_dir / "tokenizer.json") - - model_config = ModelConfig( - vocab_size=tokenizer.get_vocab_size(), - context_length=args.context_length, - embedding_size=args.embedding_size, - head_count=args.head_count, - layer_count=args.layer_count, - dropout=args.dropout, - norm_type=args.norm_type, - position_encoding=args.position_encoding, - mlp_type=args.mlp_type, - rope_theta=args.rope_theta, - ) - training_config = TrainingConfig( - output_dir=Path(args.output_dir), - epochs=args.epochs, - batch_size=args.batch_size, - learning_rate=args.learning_rate, - gradient_accumulation=args.gradient_accumulation, - sample_stride=args.sample_stride, - eval_interval=args.eval_interval, - save_interval=args.save_interval, - use_amp=args.use_amp, - device=args.device, - resume=not args.no_resume, - resume_from_checkpoint=Path(args.resume_checkpoint) if args.resume_checkpoint else None, - require_compatible_resume=not args.no_resume_safety, - ) - result = train_from_dataset(data_dir, model_config, training_config) - print(f"Saved model: {result.checkpoint_path}") - print(f"Saved summary: {result.summary_path}") - - -def benchmark(args: argparse.Namespace) -> None: - """Run benchmark prompts against a trained checkpoint. - - Args: - args: Parsed command-line arguments for the benchmark command. - """ - - prompts = normalize_prompts(Path(args.prompts_file).read_text(encoding="utf-8") if args.prompts_file else args.prompts) - result = evaluate_checkpoint( - Path(args.model_dir), - prompts, - output_dir=Path(args.output_dir) if args.output_dir else None, - max_new_tokens=args.max_new_tokens, - temperature=args.temperature, - top_k=args.top_k, - device=args.device, - use_kv_cache=not args.no_kv_cache, - ) - print(f"Benchmark prompts: {result.prompt_count}") - print(f"Benchmark time: {result.total_seconds:.2f}s") - print(f"Saved benchmark: {result.output_path}") - - -def export_hf(args: argparse.Namespace) -> None: - """Export a MicroGPT checkpoint as an HF-style package. - - Args: - args: Parsed command-line arguments for the export-hf command. - """ - - output = export_hf_microgpt_package( - Path(args.model_dir), - output_dir=Path(args.output_dir) if args.output_dir else None, - ) - print(f"Saved HF-style MicroGPT package: {output}") - - -def coordinator_server(args: argparse.Namespace) -> None: - """Run the coordinator HTTP API server. - - Args: - args: Parsed command-line arguments for the coordinator command. - """ - - print(f"Coordinator API listening on http://{args.host}:{args.port}") - print(f"Artifact root: {args.artifact_root}") - run_coordinator_api(args.host, args.port, Path(args.artifact_root) if args.artifact_root else None) - - -def create_job_bundle(args: argparse.Namespace) -> None: - """Create a portable job artifact bundle. - - Args: - args: Parsed command-line arguments for the bundle command. - """ - - tokenizer = load_tokenizer(Path(args.dataset_dir) / "tokenizer.json") - model_config = ModelConfig( - vocab_size=tokenizer.get_vocab_size(), - context_length=args.context_length, - embedding_size=args.embedding_size, - head_count=args.head_count, - layer_count=args.layer_count, - dropout=args.dropout, - ) - training_config = TrainingConfig( - output_dir=Path(args.output_dir), - epochs=args.epochs, - batch_size=args.batch_size, - learning_rate=args.learning_rate, - sample_stride=args.sample_stride, - device=args.device, - ) - job = TrainingJobSpec.local(Path(args.dataset_dir), model_config, training_config) - if args.tags is None: - default_tag = "gpu" if str(args.device).lower().startswith("cuda") else "cpu" - tags = [default_tag] - else: - tags = [item.strip() for item in args.tags.split(",") if item.strip()] - job.runtime = RuntimeSpec( - backend=BackendKind.REMOTE_CLIENT, - device=args.device, - min_vram_gb=args.min_vram_gb, - tags=tags, - ) - bundle = create_job_artifact_bundle( - job, - artifact_root=Path(args.artifact_root) if args.artifact_root else None, - base_url=args.base_url, - ) - print(f"Created artifact bundle: {bundle}") - print(f"Bundle URL: {job.metadata['artifact_bundle_url']}") - print("Job JSON:") - print(json.dumps(job.to_jsonable(), indent=2)) - - -def worker_client(args: argparse.Namespace) -> None: - """Run a remote worker client. - - Args: - args: Parsed command-line arguments for the worker client command. - """ - - labels = [item.strip() for item in args.labels.split(",") if item.strip()] - config = WorkerClientConfig( - coordinator_url=args.coordinator_url, - worker_id=args.worker_id, - device=args.device, - labels=labels, - heartbeat_interval_seconds=args.heartbeat_interval, - execute_jobs=args.execute, - claim_once=args.claim_once, - workspace_dir=Path(args.workspace_dir), - ) - print(f"Worker {config.worker_id} connecting to {config.coordinator_url}") - run_worker_client(config) - - -def build_parser() -> argparse.ArgumentParser: - """Build the command-line argument parser. - - Returns: - Configured argument parser. - """ - - parser = argparse.ArgumentParser(description="Small LLM trainer backend") - subparsers = parser.add_subparsers(required=True) - - prepare_parser = subparsers.add_parser("prepare", help="Load documents and train tokenizer") - prepare_parser.add_argument("--input_dir", required=True) - prepare_parser.add_argument("--output_dir", required=True) - prepare_parser.add_argument("--vocab_size", type=int, default=None) - prepare_parser.add_argument("--min_frequency", type=int, default=2) - prepare_parser.add_argument("--context_length", type=int, default=128) - prepare_parser.add_argument("--validation_split", type=float, default=0.1) - prepare_parser.add_argument("--max_workers", type=int, default=4) - prepare_parser.add_argument("--code_training_mode", action="store_true") - prepare_parser.add_argument("--exclude_prose", action="store_true") - prepare_parser.add_argument("--exclude_source_code", action="store_true") - prepare_parser.add_argument("--no_extract_code_blocks", action="store_true") - prepare_parser.add_argument("--no_preserve_indentation", action="store_true") - prepare_parser.add_argument("--no_instruction_samples", action="store_true") - prepare_parser.add_argument( - "--reasoning_sample_mode", - choices=["none", "scaffold", "detailed"], - default="scaffold", - ) - prepare_parser.add_argument( - "--prepare_mode", - choices=["incremental", "full_rebuild", "force_reprocess"], - default="incremental", - ) - prepare_parser.add_argument( - "--tokenizer_strategy", - choices=["auto", "train_new", "reuse_dataset", "import_tokenizer"], - default="auto", - ) - prepare_parser.add_argument("--tokenizer_path", default=None) - prepare_parser.add_argument( - "--dataset_stage", - choices=["base", "instruction", "conversation"], - default="base", - help="Purpose for online datasets: base pretraining, instruction fine-tune, or conversation fine-tune.", - ) - prepare_parser.add_argument( - "--conversation_datasets", - default="", - help="Comma-separated built-in online dataset IDs. TinyStories is for base; chat/instruction sets are for fine-tuning.", - ) - prepare_parser.add_argument("--conversation_sample_limit", type=int, default=20000) - prepare_parser.add_argument( - "--fast_scan_mode", - action="store_true", - help="Use cheaper dataset fingerprints and cached preview stats for faster large-corpus scans.", - ) - prepare_parser.add_argument( - "--fast_scan_sample_bytes", - type=int, - default=64 * 1024, - help="Bytes sampled from file head/tail for fast fingerprints.", - ) - prepare_parser.add_argument( - "--strict_duplicate_verification", - action="store_true", - help="In fast scan mode, re-hash only suspected duplicate groups with full SHA-256.", - ) - prepare_parser.set_defaults(func=prepare) - - train_parser = subparsers.add_parser("train", help="Train a MicroGPT model") - train_parser.add_argument("--data_dir", required=True) - train_parser.add_argument("--output_dir", required=True) - train_parser.add_argument("--epochs", type=int, default=5) - train_parser.add_argument("--batch_size", type=int, default=16) - train_parser.add_argument("--context_length", type=int, default=128) - train_parser.add_argument("--embedding_size", type=int, default=256) - train_parser.add_argument("--head_count", type=int, default=4) - train_parser.add_argument("--layer_count", type=int, default=4) - train_parser.add_argument("--dropout", type=float, default=0.1) - train_parser.add_argument("--norm_type", choices=["layernorm", "rmsnorm"], default="layernorm") - train_parser.add_argument("--position_encoding", choices=["learned", "rope"], default="learned") - train_parser.add_argument("--mlp_type", choices=["gelu", "swiglu"], default="gelu") - train_parser.add_argument("--rope_theta", type=float, default=10000.0) - train_parser.add_argument("--learning_rate", type=float, default=3e-4) - train_parser.add_argument("--gradient_accumulation", type=int, default=1) - train_parser.add_argument("--sample_stride", type=int, default=1) - train_parser.add_argument("--eval_interval", type=int, default=100) - train_parser.add_argument("--save_interval", type=int, default=500) - train_parser.add_argument("--use_amp", action="store_true") - train_parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") - train_parser.add_argument("--no_resume", action="store_true") - train_parser.add_argument("--resume_checkpoint", default=None) - train_parser.add_argument("--no_resume_safety", action="store_true") - train_parser.set_defaults(func=train) - - benchmark_parser = subparsers.add_parser("benchmark", help="Run fixed prompts against a trained model") - benchmark_parser.add_argument("--model_dir", required=True) - benchmark_parser.add_argument("--prompts", default="") - benchmark_parser.add_argument("--prompts_file", default=None) - benchmark_parser.add_argument("--output_dir", default=None) - benchmark_parser.add_argument("--max_new_tokens", type=int, default=128) - benchmark_parser.add_argument("--temperature", type=float, default=0.7) - benchmark_parser.add_argument("--top_k", type=int, default=50) - benchmark_parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") - benchmark_parser.add_argument("--no_kv_cache", action="store_true") - benchmark_parser.set_defaults(func=benchmark) - - export_hf_parser = subparsers.add_parser("export-hf", help="Export a MicroGPT model as an HF-style package") - export_hf_parser.add_argument("--model_dir", required=True) - export_hf_parser.add_argument("--output_dir", default=None) - export_hf_parser.set_defaults(func=export_hf) - - coordinator_parser = subparsers.add_parser("coordinator-server", help="Run the distributed training coordinator API") - coordinator_parser.add_argument("--host", default="127.0.0.1") - coordinator_parser.add_argument("--port", type=int, default=8765) - coordinator_parser.add_argument("--artifact-root", default=None) - coordinator_parser.set_defaults(func=coordinator_server) - - bundle_parser = subparsers.add_parser("create-job-bundle", help="Create a remote-worker dataset artifact bundle") - bundle_parser.add_argument("--dataset-dir", required=True) - bundle_parser.add_argument("--output-dir", required=True) - bundle_parser.add_argument("--artifact-root", default=None) - bundle_parser.add_argument("--base-url", default="/artifacts") - bundle_parser.add_argument("--epochs", type=int, default=5) - bundle_parser.add_argument("--batch-size", type=int, default=16) - bundle_parser.add_argument("--context-length", type=int, default=128) - bundle_parser.add_argument("--embedding-size", type=int, default=256) - bundle_parser.add_argument("--head-count", type=int, default=4) - bundle_parser.add_argument("--layer-count", type=int, default=4) - bundle_parser.add_argument("--dropout", type=float, default=0.1) - bundle_parser.add_argument("--learning-rate", type=float, default=3e-4) - bundle_parser.add_argument("--sample-stride", type=int, default=1) - bundle_parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") - bundle_parser.add_argument("--min-vram-gb", type=float, default=None) - bundle_parser.add_argument("--tags", default=None) - bundle_parser.set_defaults(func=create_job_bundle) - - worker_parser = subparsers.add_parser("worker-client", help="Run a remote training worker client") - worker_parser.add_argument("--coordinator-url", default="http://127.0.0.1:8765") - worker_parser.add_argument("--worker-id", default=f"worker-{os.getpid()}") - worker_parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") - worker_parser.add_argument("--labels", default="gpu" if torch.cuda.is_available() else "cpu") - worker_parser.add_argument("--heartbeat-interval", type=int, default=10) - worker_parser.add_argument( - "--workspace-dir", - default=str(Path.home() / ".drunkenbot_ide" / "worker_workspace"), - ) - worker_parser.add_argument("--claim-once", action="store_true") - worker_parser.add_argument("--execute", action="store_true") - worker_parser.set_defaults(func=worker_client) - return parser - - -def main() -> None: - """Run the command-line interface.""" - - parser = build_parser() - args = parser.parse_args() - args.func(args) - - -if __name__ == "__main__": - main() diff --git a/llm_trainer/config.py b/llm_trainer/config.py deleted file mode 100644 index 6986e3e..0000000 --- a/llm_trainer/config.py +++ /dev/null @@ -1,308 +0,0 @@ -from __future__ import annotations - -from dataclasses import asdict, dataclass, field -from pathlib import Path -from typing import Any, Optional - -import torch - - -@dataclass -class DatasetConfig: - """Configuration for building a tokenizer-ready dataset. - - Attributes: - input_dir: Folder containing source PDFs, text, JSONL, or code files. - output_dir: Folder where prepared dataset artifacts are written. - vocab_size: Optional manual tokenizer vocabulary size. - min_frequency: Minimum token frequency for BPE vocabulary entries. - context_length: Token window length used by downstream training. - validation_split: Fraction of tokens reserved for validation. - lowercase: Whether to lowercase text during ingestion. - max_workers: Number of parallel file readers. - code_training_mode: Enables code/prose tagging and code preservation. - include_prose: Keeps prose/explanation samples when code mode is active. - include_source_code: Includes source-code files when code mode is active. - extract_code_blocks: Detects code-like blocks in PDFs/text. - preserve_indentation: Keeps code line breaks and indentation. - generate_instruction_samples: Wraps code with simple instruction tags. - reasoning_sample_mode: Instruction/reasoning format: none, scaffold, or detailed. - prepare_mode: Dataset update mode: incremental, full_rebuild, or force_reprocess. - tokenizer_strategy: Tokenizer policy: auto, train_new, reuse_dataset, or import_tokenizer. - tokenizer_path: Optional existing tokenizer JSON used by import_tokenizer. - dataset_stage: Intended dataset purpose: base, instruction, conversation, or code. - conversation_datasets: Built-in Hugging Face conversation dataset IDs to include. - conversation_sample_limit: Maximum rows to read from each selected conversation dataset. Zero means no limit. - conversation_dataset_path: Optional local JSON/JSONL file or folder containing conversation samples. - instruction_dataset_path: Optional local JSON/JSONL file or folder containing instruction samples. - conversation_dataset_paths: Local JSON/JSONL files or folders containing conversation samples. - instruction_dataset_paths: Local JSON/JSONL files or folders containing instruction samples. - default_data_paths: Bundled starter data files selected from the Dataset Blueprint panel. - mixture_weights: Planned dataset mixture percentages by source family. - fast_scan_mode: Uses sampled fingerprints for faster large-corpus scans. - fast_scan_sample_bytes: Head/tail bytes per file used for fast fingerprints. - strict_duplicate_verification: In fast mode, fully re-hashes only suspected duplicate groups. - tokenizer_training_max_gb: Maximum corpus size, in GiB, shown to the BPE - tokenizer trainer (which holds a frequency table sized to whatever it - is shown, in memory, for the whole training run). Sampling is applied - only above this size; the full corpus is still encoded into training - tokens regardless of this setting. 0 or a negative value disables the - cap and trains on the entire corpus -- only safe if you have enough - RAM to hold a frequency table for your full corpus size at once. - """ - - input_dir: Path - output_dir: Path - vocab_size: Optional[int] = None - min_frequency: int = 2 - context_length: int = 512 - validation_split: float = 0.1 - lowercase: bool = False - max_workers: int = 4 - code_training_mode: bool = False - include_prose: bool = True - include_source_code: bool = True - extract_code_blocks: bool = True - preserve_indentation: bool = True - generate_instruction_samples: bool = True - reasoning_sample_mode: str = "scaffold" - prepare_mode: str = "incremental" - tokenizer_strategy: str = "auto" - tokenizer_path: Optional[Path] = None - dataset_stage: str = "base" - conversation_datasets: list[str] = field(default_factory=list) - conversation_sample_limit: int = 20000 - conversation_dataset_path: Optional[Path] = None - instruction_dataset_path: Optional[Path] = None - conversation_dataset_paths: list[Path] = field(default_factory=list) - instruction_dataset_paths: list[Path] = field(default_factory=list) - default_data_paths: list[Path] = field(default_factory=list) - mixture_weights: dict[str, float] = field(default_factory=dict) - fast_scan_mode: bool = False - fast_scan_sample_bytes: int = 64 * 1024 - strict_duplicate_verification: bool = False - tokenizer_training_max_gb: float = 2.0 - - -@dataclass -class ModelConfig: - """Configuration for the GPT-style model architecture. - - Attributes: - vocab_size: Tokenizer vocabulary size. - context_length: Maximum tokens visible to the model at once. - embedding_size: Width of token embeddings and transformer channels. - head_count: Number of causal attention heads. - layer_count: Number of transformer blocks. - dropout: Dropout probability for regularization. - bias: Whether linear and normalization layers include bias terms. - norm_type: Normalization type: layernorm or rmsnorm. - position_encoding: Position encoding type: learned or rope. - mlp_type: Feed-forward type: gelu or swiglu. - rope_theta: RoPE frequency base when position_encoding is rope. - attention_type: Attention layout: mha, gqa, or mqa. - kv_head_count: Key/value head count for grouped-query attention. - attention_backend: Attention kernel backend: manual or sdpa. - attention_window: Sliding-window attention size. Zero means full context. - """ - - vocab_size: int - context_length: int = 512 - embedding_size: int = 256 - head_count: int = 4 - layer_count: int = 6 - dropout: float = 0.1 - bias: bool = True - norm_type: str = "layernorm" - position_encoding: str = "learned" - mlp_type: str = "gelu" - rope_theta: float = 10000.0 - attention_type: str = "mha" - kv_head_count: int = 0 - attention_backend: str = "sdpa" - attention_window: int = 0 - - def validate(self) -> None: - """Validate architecture constraints. - - Raises: - ValueError: If dimensions are incompatible or too small. - """ - - if self.embedding_size % self.head_count != 0: - raise ValueError("embedding_size must be divisible by head_count") - if self.context_length < 8: - raise ValueError("context_length must be at least 8") - if self.vocab_size < 16: - raise ValueError("vocab_size is too small for language modeling") - if self.norm_type not in {"layernorm", "rmsnorm"}: - raise ValueError("norm_type must be layernorm or rmsnorm") - if self.position_encoding not in {"learned", "rope"}: - raise ValueError("position_encoding must be learned or rope") - if self.mlp_type not in {"gelu", "swiglu"}: - raise ValueError("mlp_type must be gelu or swiglu") - if self.position_encoding == "rope": - head_size = self.embedding_size // self.head_count - if head_size % 2 != 0: - raise ValueError("RoPE requires an even attention head size") - if self.attention_type not in {"mha", "gqa", "mqa"}: - raise ValueError("attention_type must be mha, gqa, or mqa") - if self.attention_backend not in {"manual", "sdpa"}: - raise ValueError("attention_backend must be manual or sdpa") - kv_heads = self.resolved_kv_head_count() - if kv_heads < 1 or kv_heads > self.head_count: - raise ValueError("kv_head_count must be between 1 and head_count") - if self.head_count % kv_heads != 0: - raise ValueError("head_count must be divisible by kv_head_count") - if self.attention_window < 0: - raise ValueError("attention_window cannot be negative") - - def resolved_kv_head_count(self) -> int: - """Return the effective key/value head count. - - Returns: - Key/value head count after applying the attention type. - """ - - if self.attention_type == "mqa": - return 1 - if self.attention_type == "gqa": - return self.kv_head_count if self.kv_head_count > 0 else max(1, self.head_count // 2) - return self.head_count - - -@dataclass -class TrainingConfig: - """Configuration for model optimization and checkpointing. - - Attributes: - output_dir: Folder where checkpoints and summaries are saved. - epochs: Number of full passes over the training dataset. - batch_size: Number of token windows per training batch. - learning_rate: Base optimizer learning rate. - weight_decay: Optimizer weight decay regularization. - optimizer_name: Optimizer family: adamw, adam, lion, or adafactor. - scheduler_name: Learning-rate schedule: warmup_linear, cosine, polynomial, one_cycle, or constant. - scheduler_min_lr_ratio: Minimum learning-rate multiplier after decay. - polynomial_power: Power used by polynomial decay. - gradient_accumulation: Batches to accumulate before optimizer step. - sample_stride: Token offset step between consecutive training windows. - warmup_steps: Steps used to ramp up learning rate. - eval_interval: Steps between validation loss checks. - max_eval_batches: Maximum validation batches per interval evaluation. Zero evaluates all validation batches. - save_interval: Steps between checkpoint writes. - data_loader_workers: CPU worker processes used to prepare token batches. - max_grad_norm: Gradient clipping norm. - use_amp: Enables mixed precision on CUDA. - precision: Numeric precision policy: fp32, fp16, or bf16. - device: Training device, usually "cuda" or "cpu". - seed: Random seed for repeatability. - training_mode: Training mode: pretrain or fine_tune. - fine_tune_from_checkpoint: Optional checkpoint used as the base model for fine-tuning. - peft_method: Parameter-efficient fine-tuning method: none or lora. - lora_rank: LoRA adapter rank. - lora_alpha: LoRA scaling alpha. - lora_dropout: Dropout applied before LoRA adapters. - lora_target_modules: Comma-separated LoRA target groups. - resume: Whether to resume from checkpoints. - resume_from_checkpoint: Optional exact checkpoint path to resume from. - require_compatible_resume: Validate tokenizer/model compatibility before resuming. - early_stopping: Stop training when validation loss stops improving. - early_stopping_patience: Consecutive evaluations without improvement before stopping. - """ - - output_dir: Path - epochs: int = 5 - batch_size: int = 16 - learning_rate: float = 3e-4 - weight_decay: float = 0.1 - optimizer_name: str = "adamw" - scheduler_name: str = "warmup_linear" - scheduler_min_lr_ratio: float = 0.1 - polynomial_power: float = 1.0 - gradient_accumulation: int = 1 - sample_stride: int = 128 #1 - warmup_steps: int = 100 - eval_interval: int = 100 - max_eval_batches: int = 50 - save_interval: int = 500 - data_loader_workers: int = 0 - max_grad_norm: float = 1.0 - activation_checkpointing: bool = False - use_amp: bool = True - precision: str = "fp16" - device: str = "cuda" if torch.cuda.is_available() else "cpu" - seed: int = 1337 - training_mode: str = "pretrain" - fine_tune_from_checkpoint: Optional[Path] = None - peft_method: str = "none" - lora_rank: int = 8 - lora_alpha: float = 16.0 - lora_dropout: float = 0.05 - lora_target_modules: str = "attention" - resume: bool = True - resume_from_checkpoint: Optional[Path] = None - require_compatible_resume: bool = True - early_stopping: bool = True - early_stopping_patience: int = 3 - - def validate(self) -> None: - """Validate optimizer and schedule settings. - - Raises: - ValueError: If any optimization setting is unsupported. - """ - - if self.optimizer_name not in {"adamw", "adam", "lion", "adafactor"}: - raise ValueError("optimizer_name must be adamw, adam, lion, or adafactor") - if self.scheduler_name not in {"warmup_linear", "cosine", "polynomial", "one_cycle", "constant"}: - raise ValueError("scheduler_name must be warmup_linear, cosine, polynomial, one_cycle, or constant") - if self.precision not in {"fp32", "fp16", "bf16"}: - raise ValueError("precision must be fp32, fp16, or bf16") - if self.training_mode not in {"pretrain", "fine_tune"}: - raise ValueError("training_mode must be pretrain or fine_tune") - if self.training_mode == "fine_tune" and self.fine_tune_from_checkpoint is None: - raise ValueError("fine_tune_from_checkpoint is required for fine_tune mode") - if self.peft_method not in {"none", "lora"}: - raise ValueError("peft_method must be none or lora") - if self.peft_method == "lora": - if self.training_mode != "fine_tune": - raise ValueError("LoRA requires fine_tune training mode") - if self.lora_rank <= 0: - raise ValueError("lora_rank must be greater than 0") - if self.lora_alpha <= 0.0: - raise ValueError("lora_alpha must be greater than 0") - if self.lora_dropout < 0.0 or self.lora_dropout > 0.9: - raise ValueError("lora_dropout must be between 0 and 0.9") - if self.scheduler_min_lr_ratio < 0.0 or self.scheduler_min_lr_ratio > 1.0: - raise ValueError("scheduler_min_lr_ratio must be between 0 and 1") - if self.polynomial_power <= 0.0: - raise ValueError("polynomial_power must be greater than 0") - if self.sample_stride <= 0: - raise ValueError("sample_stride must be greater than 0") - - -def dataclass_to_jsonable(value: Any) -> dict[str, Any]: - """Convert a dataclass into JSON-friendly values. - - Args: - value: Dataclass instance to convert. - - Returns: - Dictionary safe to pass to ``json.dumps``. - """ - - def convert(item: Any) -> Any: - """Convert nested values into JSON-friendly values.""" - - if isinstance(item, Path): - return str(item) - if isinstance(item, list): - return [convert(child) for child in item] - if isinstance(item, tuple): - return [convert(child) for child in item] - if isinstance(item, dict): - return {str(key): convert(child) for key, child in item.items()} - return item - - return convert(asdict(value)) \ No newline at end of file diff --git a/llm_trainer/contracts/__init__.py b/llm_trainer/contracts/__init__.py deleted file mode 100644 index f1ae20c..0000000 --- a/llm_trainer/contracts/__init__.py +++ /dev/null @@ -1,63 +0,0 @@ -from __future__ import annotations - -from .jobs import ( - ArtifactSpec, - BackendKind, - DatasetSpec, - JobPriority, - JobStatus, - ModelSpec, - RuntimeSpec, - TrainingJobSpec, - TrainingMetrics, - TrainingResultSpec, - utc_now_iso, -) -from .protocol import ( - ClaimJobRequest, - ClaimJobResponse, - CompleteJobRequest, - CompleteJobResponse, - FailJobRequest, - FailJobResponse, - HeartbeatRequest, - HeartbeatResponse, - ProgressReportRequest, - ProgressReportResponse, - ProtocolMessageKind, - ProtocolStatus, - RegisterWorkerRequest, - RegisterWorkerResponse, - WorkerAvailability, - WorkerCapabilities, -) - -__all__ = [ - "ArtifactSpec", - "BackendKind", - "ClaimJobRequest", - "ClaimJobResponse", - "CompleteJobRequest", - "CompleteJobResponse", - "DatasetSpec", - "FailJobRequest", - "FailJobResponse", - "HeartbeatRequest", - "HeartbeatResponse", - "JobPriority", - "JobStatus", - "ModelSpec", - "ProgressReportRequest", - "ProgressReportResponse", - "ProtocolMessageKind", - "ProtocolStatus", - "RegisterWorkerRequest", - "RegisterWorkerResponse", - "RuntimeSpec", - "TrainingJobSpec", - "TrainingMetrics", - "TrainingResultSpec", - "WorkerAvailability", - "WorkerCapabilities", - "utc_now_iso", -] diff --git a/llm_trainer/contracts/jobs.py b/llm_trainer/contracts/jobs.py deleted file mode 100644 index 83f6f13..0000000 --- a/llm_trainer/contracts/jobs.py +++ /dev/null @@ -1,412 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass, field -from datetime import datetime, timezone -from enum import Enum -from pathlib import Path -from typing import Any, Optional -from uuid import uuid4 - -from llm_trainer.config import ModelConfig, TrainingConfig, dataclass_to_jsonable - - -class BackendKind(str, Enum): - """Training backend type.""" - - LOCAL = "local" - REMOTE_CLIENT = "remote_client" - HUGGINGFACE = "huggingface" - CLOUD = "cloud" - - -class JobStatus(str, Enum): - """Training job lifecycle state.""" - - QUEUED = "queued" - ASSIGNED = "assigned" - RUNNING = "running" - PAUSED = "paused" - STOPPING = "stopping" - COMPLETED = "completed" - FAILED = "failed" - CANCELLED = "cancelled" - - -class JobPriority(str, Enum): - """Relative scheduler priority.""" - - LOW = "low" - NORMAL = "normal" - HIGH = "high" - - -def utc_now_iso() -> str: - """Return the current UTC timestamp. - - Returns: - ISO formatted UTC timestamp. - """ - - return datetime.now(timezone.utc).isoformat(timespec="seconds") - - -@dataclass -class DatasetSpec: - """Dataset artifacts required for a training job. - - Attributes: - dataset_dir: Prepared dataset directory. - tokenizer_path: Tokenizer path, usually inside ``dataset_dir``. - train_tokens_path: Training token file. - val_tokens_path: Validation token file. - summary_path: Dataset summary path. - dataset_id: Optional dataset identifier from lineage. - dataset_version: Optional dataset version identifier. - """ - - dataset_dir: Path - tokenizer_path: Optional[Path] = None - train_tokens_path: Optional[Path] = None - val_tokens_path: Optional[Path] = None - summary_path: Optional[Path] = None - dataset_id: Optional[str] = None - dataset_version: Optional[str] = None - - @classmethod - def from_dataset_dir(cls, dataset_dir: Path) -> "DatasetSpec": - """Create a dataset spec from a prepared dataset directory. - - Args: - dataset_dir: Prepared dataset directory. - - Returns: - Dataset specification. - """ - - dataset_dir = Path(dataset_dir) - train_tokens_path = dataset_dir / "train_tokens.npy" - val_tokens_path = dataset_dir / "val_tokens.npy" - if not train_tokens_path.exists() or not val_tokens_path.exists(): - train_tokens_path = dataset_dir / "train_tokens.json" - val_tokens_path = dataset_dir / "val_tokens.json" - return cls( - dataset_dir=dataset_dir, - tokenizer_path=dataset_dir / "tokenizer.json", - train_tokens_path=train_tokens_path, - val_tokens_path=val_tokens_path, - summary_path=dataset_dir / "dataset_summary.json", - ) - - -@dataclass -class ModelSpec: - """Model architecture payload for a training job. - - Attributes: - config: Model configuration. - base_checkpoint: Optional base checkpoint used for fine-tuning. - """ - - config: ModelConfig - base_checkpoint: Optional[Path] = None - - -@dataclass -class RuntimeSpec: - """Runtime and scheduling requirements for a training job. - - Attributes: - backend: Backend kind selected for this job. - device: Requested training device. - min_vram_gb: Optional minimum VRAM requirement. - preferred_worker_id: Optional worker ID to target. - priority: Scheduler priority. - tags: Free-form job tags. - """ - - backend: BackendKind = BackendKind.LOCAL - device: str = "auto" - min_vram_gb: Optional[float] = None - preferred_worker_id: Optional[str] = None - priority: JobPriority = JobPriority.NORMAL - tags: list[str] = field(default_factory=list) - - -@dataclass -class ArtifactSpec: - """Artifact destinations for a training job. - - Attributes: - output_dir: Model output directory. - checkpoints_dir: Checkpoint directory. - final_checkpoint: Final checkpoint path. - summary_path: Training summary path. - telemetry_db: Optional telemetry database path. - """ - - output_dir: Path - checkpoints_dir: Optional[Path] = None - final_checkpoint: Optional[Path] = None - summary_path: Optional[Path] = None - telemetry_db: Optional[Path] = None - - @classmethod - def from_output_dir(cls, output_dir: Path) -> "ArtifactSpec": - """Create artifact paths from a model output directory. - - Args: - output_dir: Model output directory. - - Returns: - Artifact specification. - """ - - output_dir = Path(output_dir) - return cls( - output_dir=output_dir, - checkpoints_dir=output_dir / "checkpoints", - final_checkpoint=output_dir / "final_model.pt", - summary_path=output_dir / "training_summary.json", - ) - - -@dataclass -class TrainingMetrics: - """Serializable training metrics emitted by trainers. - - Attributes: - step: Current optimizer step. - total_steps: Planned optimizer steps. - epoch: Current epoch. - total_epochs: Planned epochs. - train_loss: Latest training loss. - val_loss: Latest validation loss. - learning_rate: Current learning rate. - tokens_per_second: Token throughput. - samples_per_second: Sample throughput. - gpu_memory_percent: GPU memory usage percentage. - system_ram_percent: System RAM usage percentage. - message: Optional status message. - """ - - step: Optional[int] = None - total_steps: Optional[int] = None - epoch: Optional[int] = None - total_epochs: Optional[int] = None - train_loss: Optional[float] = None - val_loss: Optional[float] = None - learning_rate: Optional[float] = None - tokens_per_second: Optional[float] = None - samples_per_second: Optional[float] = None - gpu_memory_percent: Optional[float] = None - system_ram_percent: Optional[float] = None - message: Optional[str] = None - - -@dataclass -class TrainingResultSpec: - """Serializable result returned by a training backend. - - Attributes: - job_id: Job identifier. - status: Final job status. - checkpoint_path: Final or stopped checkpoint path. - summary_path: Training summary JSON path. - final_train_loss: Final training loss. - final_val_loss: Final validation loss. - stopped: Whether the job stopped by request. - error: Optional error text. - artifact_bundle_url: Optional coordinator URL for downloaded worker outputs. - """ - - job_id: str - status: JobStatus - checkpoint_path: Optional[Path] = None - summary_path: Optional[Path] = None - final_train_loss: Optional[float] = None - final_val_loss: Optional[float] = None - stopped: bool = False - error: Optional[str] = None - artifact_bundle_url: Optional[str] = None - - -@dataclass -class TrainingJobSpec: - """Complete backend-neutral training job contract. - - Attributes: - job_id: Stable job identifier. - created_at: UTC creation timestamp. - dataset: Dataset artifact specification. - model: Model architecture specification. - training: Training configuration. - runtime: Runtime/scheduler specification. - artifacts: Output artifact specification. - status: Current job status. - metadata: Free-form metadata for UI, manager, or cloud adapters. - """ - - dataset: DatasetSpec - model: ModelSpec - training: TrainingConfig - artifacts: ArtifactSpec - runtime: RuntimeSpec = field(default_factory=RuntimeSpec) - job_id: str = field(default_factory=lambda: f"job_{uuid4().hex}") - created_at: str = field(default_factory=utc_now_iso) - status: JobStatus = JobStatus.QUEUED - metadata: dict[str, Any] = field(default_factory=dict) - - @classmethod - def local( - cls, - dataset_dir: Path, - model_config: ModelConfig, - training_config: TrainingConfig, - metadata: Optional[dict[str, Any]] = None, - ) -> "TrainingJobSpec": - """Create a local training job spec. - - Args: - dataset_dir: Prepared dataset directory. - model_config: Model architecture configuration. - training_config: Training configuration. - metadata: Optional metadata. - - Returns: - Training job specification. - """ - - base_checkpoint = training_config.fine_tune_from_checkpoint - return cls( - dataset=DatasetSpec.from_dataset_dir(dataset_dir), - model=ModelSpec(model_config, base_checkpoint=base_checkpoint), - training=training_config, - artifacts=ArtifactSpec.from_output_dir(training_config.output_dir), - runtime=RuntimeSpec(backend=BackendKind.LOCAL, device=training_config.device), - metadata=metadata or {}, - ) - - def to_jsonable(self) -> dict[str, Any]: - """Convert the job spec to JSON-friendly values. - - Returns: - Serializable dictionary. - """ - - return { - "job_id": self.job_id, - "created_at": self.created_at, - "status": self.status.value, - "dataset": _paths_to_strings(self.dataset.__dict__), - "model": { - "config": dataclass_to_jsonable(self.model.config), - "base_checkpoint": str(self.model.base_checkpoint) if self.model.base_checkpoint else None, - }, - "training": dataclass_to_jsonable(self.training), - "runtime": _enum_values(_paths_to_strings(self.runtime.__dict__)), - "artifacts": _paths_to_strings(self.artifacts.__dict__), - "metadata": self.metadata, - } - - @classmethod - def from_jsonable(cls, data: dict[str, Any]) -> "TrainingJobSpec": - """Create a training job spec from JSON-friendly values. - - Args: - data: Serialized training job data. - - Returns: - Training job specification. - """ - - model_data = dict(data["model"]["config"]) - training_data = dict(data["training"]) - dataset_data = _strings_to_paths(data["dataset"], _DATASET_PATH_FIELDS) - artifact_data = _strings_to_paths(data["artifacts"], _ARTIFACT_PATH_FIELDS) - runtime_data = dict(data["runtime"]) - model_payload = dict(data["model"]) - for key in _TRAINING_PATH_FIELDS: - if training_data.get(key): - training_data[key] = Path(training_data[key]) - base_checkpoint = model_payload.get("base_checkpoint") - return cls( - dataset=DatasetSpec(**dataset_data), - model=ModelSpec( - config=ModelConfig(**model_data), - base_checkpoint=Path(base_checkpoint) if base_checkpoint else None, - ), - training=TrainingConfig(**training_data), - artifacts=ArtifactSpec(**artifact_data), - runtime=RuntimeSpec( - backend=BackendKind(runtime_data.get("backend", BackendKind.LOCAL.value)), - device=runtime_data.get("device", "auto"), - min_vram_gb=runtime_data.get("min_vram_gb"), - preferred_worker_id=runtime_data.get("preferred_worker_id"), - priority=JobPriority(runtime_data.get("priority", JobPriority.NORMAL.value)), - tags=list(runtime_data.get("tags") or []), - ), - job_id=data["job_id"], - created_at=data["created_at"], - status=JobStatus(data.get("status", JobStatus.QUEUED.value)), - metadata=dict(data.get("metadata") or {}), - ) - - -def _paths_to_strings(data: dict[str, Any]) -> dict[str, Any]: - """Convert path values in a dictionary to strings. - - Args: - data: Dictionary to convert. - - Returns: - Converted dictionary. - """ - - output: dict[str, Any] = {} - for key, value in data.items(): - if isinstance(value, Path): - output[key] = str(value) - else: - output[key] = value - return output - - -def _enum_values(data: dict[str, Any]) -> dict[str, Any]: - """Convert enum values in a dictionary to their raw values. - - Args: - data: Dictionary to convert. - - Returns: - Converted dictionary. - """ - - output: dict[str, Any] = {} - for key, value in data.items(): - output[key] = value.value if isinstance(value, Enum) else value - return output - - -_DATASET_PATH_FIELDS = {"dataset_dir", "tokenizer_path", "train_tokens_path", "val_tokens_path", "summary_path"} -_ARTIFACT_PATH_FIELDS = {"output_dir", "checkpoints_dir", "final_checkpoint", "summary_path", "telemetry_db"} -_TRAINING_PATH_FIELDS = {"output_dir", "fine_tune_from_checkpoint", "resume_from_checkpoint"} - - -def _strings_to_paths(data: dict[str, Any], path_fields: set[str]) -> dict[str, Any]: - """Convert selected string fields in a dictionary to paths. - - Args: - data: Dictionary to convert. - path_fields: Keys that should become paths. - - Returns: - Converted dictionary. - """ - - output: dict[str, Any] = {} - for key, value in data.items(): - if key in path_fields and value: - output[key] = Path(value) - else: - output[key] = value - return output diff --git a/llm_trainer/contracts/protocol.py b/llm_trainer/contracts/protocol.py deleted file mode 100644 index 1d578eb..0000000 --- a/llm_trainer/contracts/protocol.py +++ /dev/null @@ -1,862 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass, field -from enum import Enum -from typing import Any, Optional -from uuid import uuid4 - -from llm_trainer.contracts.jobs import BackendKind, TrainingJobSpec, TrainingMetrics, TrainingResultSpec, utc_now_iso - - -class ProtocolMessageKind(str, Enum): - """Coordinator protocol message kind.""" - - REGISTER_WORKER_REQUEST = "register_worker_request" - REGISTER_WORKER_RESPONSE = "register_worker_response" - HEARTBEAT_REQUEST = "heartbeat_request" - HEARTBEAT_RESPONSE = "heartbeat_response" - CLAIM_JOB_REQUEST = "claim_job_request" - CLAIM_JOB_RESPONSE = "claim_job_response" - PROGRESS_REPORT_REQUEST = "progress_report_request" - PROGRESS_REPORT_RESPONSE = "progress_report_response" - COMPLETE_JOB_REQUEST = "complete_job_request" - COMPLETE_JOB_RESPONSE = "complete_job_response" - FAIL_JOB_REQUEST = "fail_job_request" - FAIL_JOB_RESPONSE = "fail_job_response" - - -class ProtocolStatus(str, Enum): - """Coordinator protocol response status.""" - - OK = "ok" - REJECTED = "rejected" - ERROR = "error" - - -class WorkerAvailability(str, Enum): - """Remote worker availability status.""" - - AVAILABLE = "available" - BUSY = "busy" - OFFLINE = "offline" - - -@dataclass -class ProtocolEnvelope: - """Base metadata for coordinator protocol messages. - - Attributes: - message_id: Unique protocol message identifier. - kind: Protocol message kind. - sent_at: UTC timestamp when the message was created. - protocol_version: Protocol version string. - """ - - kind: ProtocolMessageKind - message_id: str = field(default_factory=lambda: f"msg_{uuid4().hex}") - sent_at: str = field(default_factory=utc_now_iso) - protocol_version: str = "0.1" - - def envelope_json(self) -> dict[str, Any]: - """Return JSON-friendly envelope fields. - - Returns: - Serializable envelope dictionary. - """ - - return { - "message_id": self.message_id, - "kind": self.kind.value, - "sent_at": self.sent_at, - "protocol_version": self.protocol_version, - } - - -@dataclass -class WorkerCapabilities: - """Hardware and runtime capabilities advertised by a worker. - - Attributes: - hostname: Optional host name. - platform: Operating system or runtime platform label. - cpu_count: Logical CPU count. - system_ram_gb: System RAM in GB. - gpu_names: GPU names visible to the worker. - total_vram_gb: Total visible GPU VRAM in GB. - supports_cuda: Whether CUDA is available. - supports_bf16: Whether BF16 is available. - supports_fp16: Whether FP16 is available. - extra: Free-form implementation details. - """ - - hostname: Optional[str] = None - platform: Optional[str] = None - cpu_count: Optional[int] = None - system_ram_gb: Optional[float] = None - gpu_names: list[str] = field(default_factory=list) - total_vram_gb: Optional[float] = None - supports_cuda: bool = False - supports_bf16: bool = False - supports_fp16: bool = False - extra: dict[str, Any] = field(default_factory=dict) - - def to_jsonable(self) -> dict[str, Any]: - """Convert capabilities to JSON-friendly values. - - Returns: - Serializable dictionary. - """ - - return { - "hostname": self.hostname, - "platform": self.platform, - "cpu_count": self.cpu_count, - "system_ram_gb": self.system_ram_gb, - "gpu_names": self.gpu_names, - "total_vram_gb": self.total_vram_gb, - "supports_cuda": self.supports_cuda, - "supports_bf16": self.supports_bf16, - "supports_fp16": self.supports_fp16, - "extra": self.extra, - } - - @classmethod - def from_jsonable(cls, data: dict[str, Any]) -> "WorkerCapabilities": - """Create capabilities from JSON-friendly values. - - Args: - data: Serialized capabilities. - - Returns: - Worker capabilities. - """ - - return cls( - hostname=data.get("hostname"), - platform=data.get("platform"), - cpu_count=data.get("cpu_count"), - system_ram_gb=data.get("system_ram_gb"), - gpu_names=list(data.get("gpu_names") or []), - total_vram_gb=data.get("total_vram_gb"), - supports_cuda=bool(data.get("supports_cuda")), - supports_bf16=bool(data.get("supports_bf16")), - supports_fp16=bool(data.get("supports_fp16")), - extra=dict(data.get("extra") or {}), - ) - - -@dataclass -class RegisterWorkerRequest(ProtocolEnvelope): - """Request sent by a worker to join the coordinator.""" - - worker_id: str = "" - backend: BackendKind = BackendKind.REMOTE_CLIENT - device: str = "auto" - capabilities: WorkerCapabilities = field(default_factory=WorkerCapabilities) - labels: list[str] = field(default_factory=list) - - def __init__( - self, - worker_id: str, - backend: BackendKind = BackendKind.REMOTE_CLIENT, - device: str = "auto", - capabilities: Optional[WorkerCapabilities] = None, - labels: Optional[list[str]] = None, - ) -> None: - """Create a register worker request. - - Args: - worker_id: Worker identifier. - backend: Backend kind the worker can run. - device: Preferred runtime device. - capabilities: Worker hardware/runtime capabilities. - labels: Free-form worker labels for scheduling. - """ - - super().__init__(ProtocolMessageKind.REGISTER_WORKER_REQUEST) - self.worker_id = worker_id - self.backend = backend - self.device = device - self.capabilities = capabilities or WorkerCapabilities() - self.labels = labels or [] - - def to_jsonable(self) -> dict[str, Any]: - """Convert the request to JSON-friendly values. - - Returns: - Serializable dictionary. - """ - - return { - **self.envelope_json(), - "worker_id": self.worker_id, - "backend": self.backend.value, - "device": self.device, - "capabilities": self.capabilities.to_jsonable(), - "labels": self.labels, - } - - @classmethod - def from_jsonable(cls, data: dict[str, Any]) -> "RegisterWorkerRequest": - """Create a request from JSON-friendly values. - - Args: - data: Serialized request. - - Returns: - Register worker request. - """ - - request = cls( - worker_id=data["worker_id"], - backend=BackendKind(data.get("backend", BackendKind.REMOTE_CLIENT.value)), - device=data.get("device", "auto"), - capabilities=WorkerCapabilities.from_jsonable(data.get("capabilities") or {}), - labels=list(data.get("labels") or []), - ) - _restore_envelope(request, data) - return request - - -@dataclass -class RegisterWorkerResponse(ProtocolEnvelope): - """Response returned after worker registration.""" - - status: ProtocolStatus = ProtocolStatus.OK - worker_id: str = "" - accepted: bool = True - heartbeat_interval_seconds: int = 10 - message: str = "" - - def __init__( - self, - worker_id: str, - accepted: bool = True, - status: ProtocolStatus = ProtocolStatus.OK, - heartbeat_interval_seconds: int = 10, - message: str = "", - ) -> None: - """Create a register worker response. - - Args: - worker_id: Worker identifier. - accepted: Whether registration was accepted. - status: Protocol response status. - heartbeat_interval_seconds: Requested heartbeat interval. - message: Human-readable response message. - """ - - super().__init__(ProtocolMessageKind.REGISTER_WORKER_RESPONSE) - self.status = status - self.worker_id = worker_id - self.accepted = accepted - self.heartbeat_interval_seconds = heartbeat_interval_seconds - self.message = message - - def to_jsonable(self) -> dict[str, Any]: - """Convert the response to JSON-friendly values. - - Returns: - Serializable dictionary. - """ - - return { - **self.envelope_json(), - "status": self.status.value, - "worker_id": self.worker_id, - "accepted": self.accepted, - "heartbeat_interval_seconds": self.heartbeat_interval_seconds, - "message": self.message, - } - - -@dataclass -class HeartbeatRequest(ProtocolEnvelope): - """Heartbeat request sent by a worker.""" - - worker_id: str = "" - availability: WorkerAvailability = WorkerAvailability.AVAILABLE - backend: BackendKind = BackendKind.REMOTE_CLIENT - active_job_id: Optional[str] = None - device: str = "auto" - metrics: dict[str, Any] = field(default_factory=dict) - - def __init__( - self, - worker_id: str, - availability: WorkerAvailability = WorkerAvailability.AVAILABLE, - backend: BackendKind = BackendKind.REMOTE_CLIENT, - active_job_id: Optional[str] = None, - device: str = "auto", - metrics: Optional[dict[str, Any]] = None, - ) -> None: - """Create a heartbeat request. - - Args: - worker_id: Worker identifier. - availability: Current worker availability. - backend: Backend kind the worker can run. - active_job_id: Active job identifier when busy. - device: Runtime device. - metrics: Worker metrics. - """ - - super().__init__(ProtocolMessageKind.HEARTBEAT_REQUEST) - self.worker_id = worker_id - self.availability = availability - self.backend = backend - self.active_job_id = active_job_id - self.device = device - self.metrics = metrics or {} - - def to_jsonable(self) -> dict[str, Any]: - """Convert the request to JSON-friendly values. - - Returns: - Serializable dictionary. - """ - - return { - **self.envelope_json(), - "worker_id": self.worker_id, - "availability": self.availability.value, - "backend": self.backend.value, - "active_job_id": self.active_job_id, - "device": self.device, - "metrics": self.metrics, - } - - @classmethod - def from_jsonable(cls, data: dict[str, Any]) -> "HeartbeatRequest": - """Create a request from JSON-friendly values. - - Args: - data: Serialized request. - - Returns: - Heartbeat request. - """ - - request = cls( - worker_id=data["worker_id"], - availability=WorkerAvailability(data.get("availability", WorkerAvailability.AVAILABLE.value)), - backend=BackendKind(data.get("backend", BackendKind.REMOTE_CLIENT.value)), - active_job_id=data.get("active_job_id"), - device=data.get("device", "auto"), - metrics=dict(data.get("metrics") or {}), - ) - _restore_envelope(request, data) - return request - - -@dataclass -class HeartbeatResponse(ProtocolEnvelope): - """Response returned after a worker heartbeat.""" - - status: ProtocolStatus = ProtocolStatus.OK - should_stop_job: bool = False - should_pause_job: bool = False - message: str = "" - - def __init__( - self, - status: ProtocolStatus = ProtocolStatus.OK, - should_stop_job: bool = False, - should_pause_job: bool = False, - message: str = "", - ) -> None: - """Create a heartbeat response. - - Args: - status: Protocol response status. - should_stop_job: Whether active job should stop. - should_pause_job: Whether active job should pause. - message: Human-readable response message. - """ - - super().__init__(ProtocolMessageKind.HEARTBEAT_RESPONSE) - self.status = status - self.should_stop_job = should_stop_job - self.should_pause_job = should_pause_job - self.message = message - - def to_jsonable(self) -> dict[str, Any]: - """Convert the response to JSON-friendly values. - - Returns: - Serializable dictionary. - """ - - return { - **self.envelope_json(), - "status": self.status.value, - "should_stop_job": self.should_stop_job, - "should_pause_job": self.should_pause_job, - "message": self.message, - } - - -@dataclass -class ClaimJobRequest(ProtocolEnvelope): - """Request sent by a worker asking for a job.""" - - worker_id: str = "" - backend: BackendKind = BackendKind.REMOTE_CLIENT - capabilities: WorkerCapabilities = field(default_factory=WorkerCapabilities) - - def __init__( - self, - worker_id: str, - backend: BackendKind = BackendKind.REMOTE_CLIENT, - capabilities: Optional[WorkerCapabilities] = None, - ) -> None: - """Create a claim job request. - - Args: - worker_id: Worker identifier. - backend: Backend kind requested by the worker. - capabilities: Latest worker capabilities. - """ - - super().__init__(ProtocolMessageKind.CLAIM_JOB_REQUEST) - self.worker_id = worker_id - self.backend = backend - self.capabilities = capabilities or WorkerCapabilities() - - def to_jsonable(self) -> dict[str, Any]: - """Convert the request to JSON-friendly values. - - Returns: - Serializable dictionary. - """ - - return { - **self.envelope_json(), - "worker_id": self.worker_id, - "backend": self.backend.value, - "capabilities": self.capabilities.to_jsonable(), - } - - @classmethod - def from_jsonable(cls, data: dict[str, Any]) -> "ClaimJobRequest": - """Create a request from JSON-friendly values. - - Args: - data: Serialized request. - - Returns: - Claim job request. - """ - - request = cls( - worker_id=data["worker_id"], - backend=BackendKind(data.get("backend", BackendKind.REMOTE_CLIENT.value)), - capabilities=WorkerCapabilities.from_jsonable(data.get("capabilities") or {}), - ) - _restore_envelope(request, data) - return request - - -@dataclass -class ClaimJobResponse(ProtocolEnvelope): - """Response with an assigned job or empty assignment.""" - - status: ProtocolStatus = ProtocolStatus.OK - job: Optional[TrainingJobSpec] = None - message: str = "" - - def __init__( - self, - job: Optional[TrainingJobSpec] = None, - status: ProtocolStatus = ProtocolStatus.OK, - message: str = "", - ) -> None: - """Create a claim job response. - - Args: - job: Assigned job contract when available. - status: Protocol response status. - message: Human-readable response message. - """ - - super().__init__(ProtocolMessageKind.CLAIM_JOB_RESPONSE) - self.status = status - self.job = job - self.message = message - - def to_jsonable(self) -> dict[str, Any]: - """Convert the response to JSON-friendly values. - - Returns: - Serializable dictionary. - """ - - return { - **self.envelope_json(), - "status": self.status.value, - "job": self.job.to_jsonable() if self.job else None, - "message": self.message, - } - - @classmethod - def from_jsonable(cls, data: dict[str, Any]) -> "ClaimJobResponse": - """Create a response from JSON-friendly values. - - Args: - data: Serialized response. - - Returns: - Claim job response. - """ - - job_data = data.get("job") - response = cls( - job=TrainingJobSpec.from_jsonable(job_data) if job_data else None, - status=ProtocolStatus(data.get("status", ProtocolStatus.OK.value)), - message=data.get("message", ""), - ) - _restore_envelope(response, data) - return response - - -@dataclass -class ProgressReportRequest(ProtocolEnvelope): - """Progress update sent by a worker for a running job.""" - - worker_id: str = "" - job_id: str = "" - metrics: TrainingMetrics = field(default_factory=TrainingMetrics) - - def __init__(self, worker_id: str, job_id: str, metrics: Optional[TrainingMetrics] = None) -> None: - """Create a progress report request. - - Args: - worker_id: Worker identifier. - job_id: Job identifier. - metrics: Training metrics. - """ - - super().__init__(ProtocolMessageKind.PROGRESS_REPORT_REQUEST) - self.worker_id = worker_id - self.job_id = job_id - self.metrics = metrics or TrainingMetrics() - - def to_jsonable(self) -> dict[str, Any]: - """Convert the request to JSON-friendly values. - - Returns: - Serializable dictionary. - """ - - return { - **self.envelope_json(), - "worker_id": self.worker_id, - "job_id": self.job_id, - "metrics": self.metrics.__dict__, - } - - @classmethod - def from_jsonable(cls, data: dict[str, Any]) -> "ProgressReportRequest": - """Create a request from JSON-friendly values. - - Args: - data: Serialized request. - - Returns: - Progress report request. - """ - - request = cls( - worker_id=data["worker_id"], - job_id=data["job_id"], - metrics=TrainingMetrics(**dict(data.get("metrics") or {})), - ) - _restore_envelope(request, data) - return request - - -@dataclass -class ProgressReportResponse(ProtocolEnvelope): - """Response returned after a progress report.""" - - status: ProtocolStatus = ProtocolStatus.OK - should_stop_job: bool = False - should_pause_job: bool = False - message: str = "" - - def __init__( - self, - status: ProtocolStatus = ProtocolStatus.OK, - should_stop_job: bool = False, - should_pause_job: bool = False, - message: str = "", - ) -> None: - """Create a progress report response. - - Args: - status: Protocol response status. - should_stop_job: Whether the worker should stop the active job. - should_pause_job: Whether the worker should pause the active job. - message: Human-readable response message. - """ - - super().__init__(ProtocolMessageKind.PROGRESS_REPORT_RESPONSE) - self.status = status - self.should_stop_job = should_stop_job - self.should_pause_job = should_pause_job - self.message = message - - def to_jsonable(self) -> dict[str, Any]: - """Convert the response to JSON-friendly values. - - Returns: - Serializable dictionary. - """ - - return { - **self.envelope_json(), - "status": self.status.value, - "should_stop_job": self.should_stop_job, - "should_pause_job": self.should_pause_job, - "message": self.message, - } - - -@dataclass -class CompleteJobRequest(ProtocolEnvelope): - """Completion report sent by a worker.""" - - worker_id: str = "" - result: Optional[TrainingResultSpec] = None - - def __init__(self, worker_id: str, result: TrainingResultSpec) -> None: - """Create a complete job request. - - Args: - worker_id: Worker identifier. - result: Training result specification. - """ - - super().__init__(ProtocolMessageKind.COMPLETE_JOB_REQUEST) - self.worker_id = worker_id - self.result = result - - def to_jsonable(self) -> dict[str, Any]: - """Convert the request to JSON-friendly values. - - Returns: - Serializable dictionary. - """ - - return { - **self.envelope_json(), - "worker_id": self.worker_id, - "result": _result_to_jsonable(self.result) if self.result else None, - } - - @classmethod - def from_jsonable(cls, data: dict[str, Any]) -> "CompleteJobRequest": - """Create a request from JSON-friendly values. - - Args: - data: Serialized request. - - Returns: - Complete job request. - """ - - request = cls( - worker_id=data["worker_id"], - result=_result_from_jsonable(dict(data["result"])), - ) - _restore_envelope(request, data) - return request - - -@dataclass -class CompleteJobResponse(ProtocolEnvelope): - """Response returned after job completion.""" - - status: ProtocolStatus = ProtocolStatus.OK - message: str = "" - - def __init__(self, status: ProtocolStatus = ProtocolStatus.OK, message: str = "") -> None: - """Create a complete job response. - - Args: - status: Protocol response status. - message: Human-readable response message. - """ - - super().__init__(ProtocolMessageKind.COMPLETE_JOB_RESPONSE) - self.status = status - self.message = message - - def to_jsonable(self) -> dict[str, Any]: - """Convert the response to JSON-friendly values. - - Returns: - Serializable dictionary. - """ - - return { - **self.envelope_json(), - "status": self.status.value, - "message": self.message, - } - - -@dataclass -class FailJobRequest(ProtocolEnvelope): - """Failure report sent by a worker.""" - - worker_id: str = "" - job_id: str = "" - error: str = "" - retryable: bool = False - - def __init__(self, worker_id: str, job_id: str, error: str, retryable: bool = False) -> None: - """Create a fail job request. - - Args: - worker_id: Worker identifier. - job_id: Job identifier. - error: Failure text. - retryable: Whether the coordinator may retry the job. - """ - - super().__init__(ProtocolMessageKind.FAIL_JOB_REQUEST) - self.worker_id = worker_id - self.job_id = job_id - self.error = error - self.retryable = retryable - - def to_jsonable(self) -> dict[str, Any]: - """Convert the request to JSON-friendly values. - - Returns: - Serializable dictionary. - """ - - return { - **self.envelope_json(), - "worker_id": self.worker_id, - "job_id": self.job_id, - "error": self.error, - "retryable": self.retryable, - } - - @classmethod - def from_jsonable(cls, data: dict[str, Any]) -> "FailJobRequest": - """Create a request from JSON-friendly values. - - Args: - data: Serialized request. - - Returns: - Fail job request. - """ - - request = cls( - worker_id=data["worker_id"], - job_id=data["job_id"], - error=data.get("error", ""), - retryable=bool(data.get("retryable")), - ) - _restore_envelope(request, data) - return request - - -@dataclass -class FailJobResponse(ProtocolEnvelope): - """Response returned after job failure report.""" - - status: ProtocolStatus = ProtocolStatus.OK - message: str = "" - - def __init__(self, status: ProtocolStatus = ProtocolStatus.OK, message: str = "") -> None: - """Create a fail job response. - - Args: - status: Protocol response status. - message: Human-readable response message. - """ - - super().__init__(ProtocolMessageKind.FAIL_JOB_RESPONSE) - self.status = status - self.message = message - - def to_jsonable(self) -> dict[str, Any]: - """Convert the response to JSON-friendly values. - - Returns: - Serializable dictionary. - """ - - return { - **self.envelope_json(), - "status": self.status.value, - "message": self.message, - } - - -def _restore_envelope(message: ProtocolEnvelope, data: dict[str, Any]) -> None: - """Restore envelope fields on a protocol message. - - Args: - message: Protocol message. - data: Serialized message data. - """ - - message.message_id = data.get("message_id", message.message_id) - message.sent_at = data.get("sent_at", message.sent_at) - message.protocol_version = data.get("protocol_version", message.protocol_version) - - -def _result_to_jsonable(result: TrainingResultSpec) -> dict[str, Any]: - """Convert a result spec to JSON-friendly values. - - Args: - result: Training result specification. - - Returns: - Serializable dictionary. - """ - - output: dict[str, Any] = {} - for key, value in result.__dict__.items(): - if hasattr(value, "value"): - output[key] = value.value - elif value is None: - output[key] = None - else: - output[key] = str(value) if key.endswith("_path") else value - return output - - -def _result_from_jsonable(data: dict[str, Any]) -> TrainingResultSpec: - """Create a result spec from JSON-friendly values. - - Args: - data: Serialized result. - - Returns: - Training result specification. - """ - - from pathlib import Path - from llm_trainer.contracts.jobs import JobStatus - - checkpoint_path = data.get("checkpoint_path") - summary_path = data.get("summary_path") - return TrainingResultSpec( - job_id=data["job_id"], - status=JobStatus(data["status"]), - checkpoint_path=Path(checkpoint_path) if checkpoint_path else None, - summary_path=Path(summary_path) if summary_path else None, - final_train_loss=data.get("final_train_loss"), - final_val_loss=data.get("final_val_loss"), - stopped=bool(data.get("stopped")), - error=data.get("error"), - artifact_bundle_url=data.get("artifact_bundle_url"), - ) diff --git a/llm_trainer/conversation_datasets.py b/llm_trainer/conversation_datasets.py deleted file mode 100644 index 1200577..0000000 --- a/llm_trainer/conversation_datasets.py +++ /dev/null @@ -1,715 +0,0 @@ -from __future__ import annotations - -import argparse -from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait -from dataclasses import dataclass -import json -import logging -import os -from pathlib import Path -from queue import Empty, Queue -import subprocess -import sys -from threading import Thread -from typing import Any, Callable, Optional - -from llm_trainer.data import Document, clean_text - - -LOGGER = logging.getLogger(__name__) - - -@dataclass(frozen=True) -class ConversationDatasetPreset: - """Built-in Hugging Face dataset recipe for conversation/instruction data. - - Attributes: - dataset_id: Stable UI/config identifier. - label: User-facing label with size hint. - hf_path: Hugging Face dataset path. - config_name: Optional Hugging Face dataset configuration name. - split: Dataset split to load. - stage: Recommended training stage: base, instruction, or conversation. - description: Short user-facing purpose hint. - """ - - dataset_id: str - label: str - hf_path: str - config_name: Optional[str] - split: str - stage: str - description: str - - -CONVERSATION_DATASET_PRESETS: dict[str, ConversationDatasetPreset] = { - "tinystories": ConversationDatasetPreset( - "tinystories", - "TinyStories (~2M short stories)", - "roneneldan/TinyStories", - None, - "train", - "base", - "Language fluency, simple narrative structure, and basic world knowledge.", - ), - "wikitext_103": ConversationDatasetPreset( - "wikitext_103", - "WikiText-103 (~100M tokens)", - "Salesforce/wikitext", - "wikitext-103-raw-v1", - "train", - "base", - "Clean Wikipedia-style long-form text for grammar, facts, and language modeling.", - ), - "wikipedia_en": ConversationDatasetPreset( - "wikipedia_en", - "Wikipedia EN 2023 (large encyclopedia)", - "wikimedia/wikipedia", - "20250101.en", - "train", - "base", - "Broad encyclopedia prose. Use a row limit unless you intentionally want a large download.", - ), - "openwebtext": ConversationDatasetPreset( - "openwebtext", "OpenWebText (~8M web documents)", "Skylion007/openwebtext", - None, "train", "base", "Broad web text for general language pretraining.", - ), - "bookcorpusopen": ConversationDatasetPreset( - "bookcorpusopen", "BookCorpusOpen (books)", "kmfoda/bookcorpus", - None, "train", "base", "Long-form literary text and narrative language.", - ), - "scientific_papers": ConversationDatasetPreset( - "scientific_papers", "Scientific Papers (ArXiv)", "scientific_papers", - "arxiv", "train", "base", "Scientific writing and technical vocabulary.", - ), - "pubmed_qa": ConversationDatasetPreset( - "pubmed_qa", "PubMed QA", "pubmed_qa", "pqa_labeled", - "train", "instruction", "Biomedical question answering.", - ), - "open_orca": ConversationDatasetPreset( - "open_orca", "OpenOrca (~1M instructions)", "Open-Orca/OpenOrca", - None, "train", "instruction", "Diverse instruction and reasoning answers.", - ), - "wizardlm_evol_instruct": ConversationDatasetPreset( - "wizardlm_evol_instruct", "WizardLM Evol-Instruct", "WizardLM/WizardLM_evol_instruct_V2_196k", - None, "train", "instruction", "Evolved multi-step instruction following.", - ), - "no_robots": ConversationDatasetPreset( - "no_robots", "No Robots (10K conversations)", "HuggingFaceH4/no_robots", - None, "train", "conversation", "High-quality multi-turn assistant conversations.", - ), - "fineweb_edu": ConversationDatasetPreset( - "fineweb_edu", - "FineWeb-Edu sample (large educational web)", - "HuggingFaceFW/fineweb-edu", - "sample-10BT", - "train", - "base", - "High-quality educational web text for base language pretraining. Use a row limit.", - ), - "ultrachat_200k": ConversationDatasetPreset( - "ultrachat_200k", - "UltraChat 200K (~200K conversations)", - "HuggingFaceH4/ultrachat_200k", - None, - "train_sft", - "conversation", - "Multi-turn assistant conversation and helpful response style.", - ), - "dailydialog": ConversationDatasetPreset( - "dailydialog", - "DailyDialog (~13K dialogues)", - "pixelsandpointers/better_daily_dialog", - None, - "train", - "conversation", - "Natural everyday dialogue and short conversational turns.", - ), - "alpaca_52k": ConversationDatasetPreset( - "alpaca_52k", - "Alpaca 52K (~52K instructions)", - "tatsu-lab/alpaca", - None, - "train", - "instruction", - "Instruction following with concise task-answer pairs.", - ), - "dolly_15k": ConversationDatasetPreset( - "dolly_15k", - "Dolly 15K (~15K instructions)", - "databricks/databricks-dolly-15k", - None, - "train", - "instruction", - "Human-written instruction following, brainstorming, QA, and classification.", - ), - "oasst1": ConversationDatasetPreset( - "oasst1", - "OpenAssistant OASST1 (~88K messages)", - "OpenAssistant/oasst1", - None, - "train", - "conversation", - "Assistant-style conversational messages and preference data text.", - ), - "slimorca": ConversationDatasetPreset( - "slimorca", - "SlimOrca (~517K examples)", - "Open-Orca/SlimOrca", - None, - "train", - "instruction", - "Instruction and reasoning-style assistant answers.", - ), - "codealpaca_20k": ConversationDatasetPreset( - "codealpaca_20k", - "CodeAlpaca 20K (~20K code instructions)", - "sahil2801/CodeAlpaca-20k", - None, - "train", - "code", - "Small code instruction dataset for text-to-code, code explanation, and programming tasks.", - ), - "magicoder_oss_75k": ConversationDatasetPreset( - "magicoder_oss_75k", - "Magicoder OSS-Instruct 75K (~75K code tasks)", - "ise-uiuc/Magicoder-OSS-Instruct-75K", - None, - "train", - "code", - "Code generation instruction data built from open-source code references.", - ), - "evol_codealpaca": ConversationDatasetPreset( - "evol_codealpaca", - "Evol CodeAlpaca (~evolved code instructions)", - "theblackcat102/evol-codealpaca-v1", - None, - "train", - "code", - "Evolved programming instructions for stronger code fine-tuning variety.", - ), -} - -BASE_DATASET_IDS = [dataset_id for dataset_id, preset in CONVERSATION_DATASET_PRESETS.items() if preset.stage == "base"] -INSTRUCTION_DATASET_IDS = [dataset_id for dataset_id, preset in CONVERSATION_DATASET_PRESETS.items() if preset.stage == "instruction"] -CONVERSATION_DATASET_IDS = [dataset_id for dataset_id, preset in CONVERSATION_DATASET_PRESETS.items() if preset.stage == "conversation"] -CODE_DATASET_IDS = [dataset_id for dataset_id, preset in CONVERSATION_DATASET_PRESETS.items() if preset.stage == "code"] - - -def dataset_ids_for_stage(stage: str) -> list[str]: - """Return online dataset IDs available for a training stage. - - Args: - stage: Dataset/training stage. - - Returns: - Dataset IDs for the selected stage. Base pretraining intentionally - exposes every built-in source so users can build mixed base corpora. - """ - - if stage == "base": - return list(CONVERSATION_DATASET_PRESETS) - if stage == "instruction": - return INSTRUCTION_DATASET_IDS - if stage == "conversation": - return CONVERSATION_DATASET_IDS - if stage == "code": - return CODE_DATASET_IDS - return [] - - -def dataset_stage_label(stage: str) -> str: - """Return a user-facing stage label. - - Args: - stage: Dataset/training stage. - - Returns: - Human-readable stage name. - """ - - return { - "base": "Base pretraining", - "instruction": "Instruction fine-tune", - "conversation": "Conversation fine-tune", - "code": "Code fine-tune", - }.get(stage, "Custom") - - -def load_conversation_documents( - dataset_ids: list[str], - sample_limit: int, - cache_dir: Path, - lowercase: bool = False, - progress: Optional[Callable[[Any], None]] = None, - should_stop: Optional[Callable[[], bool]] = None, -) -> list[Document]: - """Load selected Hugging Face conversation datasets as training documents. - - Args: - dataset_ids: Preset IDs to load. - sample_limit: Maximum rows per dataset. Zero means no limit. - cache_dir: Hugging Face dataset cache directory. - lowercase: Whether to lowercase extracted text. - progress: Optional progress callback. - should_stop: Optional cancellation callback. - - Returns: - Conversation/instruction documents. - """ - - if not dataset_ids: - return [] - documents: list[Document] = [] - cache_dir.mkdir(parents=True, exist_ok=True) - _emit(progress, f"Hugging Face dataset cache: {cache_dir}") - LOGGER.info("Hugging Face dataset cache: %s", cache_dir) - def load_one(preset_index: int, dataset_id: str) -> tuple[str, list[Document]]: - """Load one preset and return its documents.""" - - if should_stop and should_stop(): - raise RuntimeError("Dataset preparation stopped by user.") - preset = CONVERSATION_DATASET_PRESETS.get(dataset_id) - if preset is None and dataset_id.startswith("hf_custom:"): - custom_path = dataset_id.removeprefix("hf_custom:") - preset = ConversationDatasetPreset( - dataset_id=dataset_id, - label=f"Custom Hugging Face dataset ({custom_path})", - hf_path=custom_path, - config_name=None, - split="train", - stage="base", - description="User-provided Hugging Face dataset.", - ) - if preset is None: - _emit(progress, f"Skipping unknown conversation dataset: {dataset_id}") - LOGGER.warning("Skipping unknown conversation dataset: %s", dataset_id) - return dataset_id, [] - _emit( - progress, - f"Downloading/loading {preset.label} into {cache_dir}...", - 8 + min(25, preset_index * 3), - ) - LOGGER.info("Downloading/loading %s into %s", preset.label, cache_dir) - rows = _load_preset_rows_in_subprocess(preset, sample_limit, cache_dir, lowercase, progress, should_stop) - total = len(rows) - loaded = 0 - preset_documents: list[Document] = [] - for row_index, row in enumerate(rows): - if should_stop and should_stop(): - raise RuntimeError("Dataset preparation stopped by user.") - kind = str(row.get("kind") or "conversation") - messages = row.get("messages") - if isinstance(messages, list): - text = _render_message_list(messages) - else: - text = str(row.get("text") or "") - if not text: - continue - preset_documents.append( - Document( - path=Path("__hf_datasets__") / preset.dataset_id / f"{row_index}.txt", - text=text, - kind=kind, - language=preset.dataset_id, - ) - ) - loaded += 1 - if loaded % 1000 == 0: - _emit(progress, f"{preset.label}: loaded {loaded:,}/{total:,} sample(s).") - _emit(progress, f"{preset.label}: added {loaded:,} sample(s).") - LOGGER.info("%s added %s sample(s)", preset.label, f"{loaded:,}") - return dataset_id, preset_documents - - max_workers = min(4, max(1, len(dataset_ids))) - if len(dataset_ids) > 1: - _emit(progress, f"Loading {len(dataset_ids)} online dataset(s) in parallel with {max_workers} worker(s).") - LOGGER.info("Loading %s online dataset(s) in parallel with %s worker(s)", len(dataset_ids), max_workers) - with ThreadPoolExecutor(max_workers=max_workers) as executor: - pending = { - executor.submit(load_one, preset_index, dataset_id) - for preset_index, dataset_id in enumerate(dataset_ids, start=1) - } - while pending: - if should_stop and should_stop(): - for future in pending: - future.cancel() - raise RuntimeError("Dataset preparation stopped by user.") - done, pending = wait(pending, timeout=0.25, return_when=FIRST_COMPLETED) - for future in done: - if should_stop and should_stop(): - for pending_future in pending: - pending_future.cancel() - raise RuntimeError("Dataset preparation stopped by user.") - _, preset_documents = future.result() - documents.extend(preset_documents) - for future in pending: - _, preset_documents = future.result() - documents.extend(preset_documents) - return documents - - -def _load_preset_rows_in_subprocess( - preset: ConversationDatasetPreset, - sample_limit: int, - cache_dir: Path, - lowercase: bool, - progress: Optional[Callable[[Any], None]], - should_stop: Optional[Callable[[], bool]], -) -> list[dict[str, str]]: - """Extract a Hugging Face preset in a child process. - - Args: - preset: Dataset preset to extract. - sample_limit: Maximum rows to extract. - cache_dir: Hugging Face cache directory. - lowercase: Whether to lowercase text. - progress: Optional progress callback. - should_stop: Optional cancellation callback. - - Returns: - Extracted text rows. - """ - - extract_dir = cache_dir / "_micro_llm_extracted" - extract_dir.mkdir(parents=True, exist_ok=True) - output_path = extract_dir / f"{preset.dataset_id}_{max(sample_limit, 0)}_{int(lowercase)}.jsonl" - if output_path.exists(): - output_path.unlink() - command = [ - sys.executable, - "-m", - "llm_trainer.conversation_datasets", - "extract", - "--dataset-id", - preset.dataset_id, - "--sample-limit", - str(sample_limit), - "--cache-dir", - str(cache_dir), - "--output-jsonl", - str(output_path), - ] - if preset.config_name: - command.extend(["--config-name", preset.config_name]) - if lowercase: - command.append("--lowercase") - LOGGER.info("Starting Hugging Face extraction subprocess: %s", " ".join(command)) - process = subprocess.Popen( - command, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - encoding="utf-8", - errors="replace", - env=_hf_subprocess_environment(cache_dir), - ) - assert process.stdout is not None - - output_queue: Queue[str] = Queue() - - def read_output() -> None: - """Read child output without blocking cancellation polling.""" - - assert process.stdout is not None - for line in process.stdout: - output_queue.put(line) - - reader = Thread(target=read_output, daemon=True) - reader.start() - while process.poll() is None: - while True: - try: - line = output_queue.get_nowait() - except Empty: - break - text = line.strip() - if text: - LOGGER.info("[hf:%s] %s", preset.dataset_id, text) - _emit(progress, text) - if should_stop and should_stop(): - _terminate_process(process, preset.dataset_id) - raise RuntimeError("Dataset preparation stopped by user.") - try: - line = output_queue.get(timeout=0.2) - except Empty: - continue - text = line.strip() - if text: - LOGGER.info("[hf:%s] %s", preset.dataset_id, text) - _emit(progress, text) - return_code = process.wait() - reader.join(timeout=1) - while True: - try: - line = output_queue.get_nowait() - except Empty: - break - text = line.strip() - if text: - LOGGER.info("[hf:%s] %s", preset.dataset_id, text) - _emit(progress, text) - LOGGER.info("Hugging Face extraction subprocess finished for %s with code %s", preset.dataset_id, return_code) - if return_code != 0: - raise RuntimeError( - f"Hugging Face dataset loader exited with code {return_code} while loading {preset.label}. " - "Check drunkenbot_ide.log and drunkenbot_ide_faults.log." - ) - if not output_path.exists(): - raise RuntimeError(f"Hugging Face extraction did not create output: {output_path}") - rows: list[dict[str, str]] = [] - with output_path.open("r", encoding="utf-8") as file: - for line in file: - if line.strip(): - rows.append(json.loads(line)) - return rows - - -def _terminate_process(process: subprocess.Popen[Any], dataset_id: str) -> None: - """Terminate a child process and escalate to kill if it stays alive. - - Args: - process: Running child process. - dataset_id: Dataset ID used for logging. - """ - - LOGGER.info("Terminating Hugging Face extraction subprocess for %s", dataset_id) - process.terminate() - try: - process.wait(timeout=3) - except subprocess.TimeoutExpired: - LOGGER.warning("Killing unresponsive Hugging Face extraction subprocess for %s", dataset_id) - process.kill() - process.wait(timeout=3) - - -def _hf_subprocess_environment(cache_dir: Path) -> dict[str, str]: - """Build a Hugging Face environment that stays inside the project cache. - - Args: - cache_dir: Project-local Hugging Face cache directory. - - Returns: - Environment variables for the extraction subprocess. - """ - - env = os.environ.copy() - env["HF_HOME"] = str(cache_dir / "hf_home") - env["HF_HUB_CACHE"] = str(cache_dir / "hub") - env["HF_DATASETS_CACHE"] = str(cache_dir / "datasets") - env["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1" - return env - - -def _emit(progress: Optional[Callable[[Any], None]], message: str, percent: Optional[int] = None) -> None: - """Emit a progress event if a callback is available.""" - - if progress: - progress({"message": message, "percent": percent}) - - -def _conversation_text_from_row(row: dict[str, Any]) -> tuple[str, str]: - """Extract tagged conversation/instruction text from a dataset row. - - Args: - row: Hugging Face row. - - Returns: - Text and document kind. - """ - - messages = row.get("messages") or row.get("conversation") or row.get("conversations") - if isinstance(messages, list): - rendered = _render_message_list(messages) - if rendered: - return rendered, "conversation" - dialogue = row.get("dialogue") - if isinstance(dialogue, list): - turns = [f"{'User' if index % 2 == 0 else 'Assistant'}: {value}" for index, value in enumerate(dialogue)] - return "\n".join(turns), "conversation" - instruction = str(row.get("instruction") or row.get("prompt") or row.get("question") or "").strip() - input_text = str(row.get("input") or row.get("context") or row.get("problem") or "").strip() - output = str( - row.get("output") - or row.get("response") - or row.get("answer") - or row.get("completion") - or row.get("solution") - or row.get("code") - or "" - ).strip() - if instruction and output: - user = instruction if not input_text else f"{instruction}\n\n{input_text}" - return f"User: {user}\nAssistant: {output}", "instruction" - for key in ("text", "story", "content"): - value = row.get(key) - if value: - return str(value), "prose" - return "", "prose" - - -def _chat_messages_from_text(text: str, kind: str) -> list[dict[str, str]]: - """Convert tagged extracted text into the canonical chat message shape.""" - messages: list[dict[str, str]] = [] - for line in text.splitlines(): - if ": " not in line: - continue - role, content = line.split(": ", 1) - role_key = role.strip().lower() - if role_key in {"system", "user", "assistant"} and content.strip(): - messages.append({"role": role_key, "content": content.strip()}) - if kind == "instruction" and not any(item["role"] == "system" for item in messages): - messages.insert(0, {"role": "system", "content": "You are a helpful assistant."}) - return messages - - -def _render_message_list(messages: list[Any]) -> str: - """Render common message-list schemas into role-prefixed turns.""" - - turns: list[str] = [] - for index, message in enumerate(messages): - if isinstance(message, dict): - role = str(message.get("role") or message.get("from") or message.get("speaker") or "").strip() - content = str(message.get("content") or message.get("value") or message.get("text") or "").strip() - else: - role = "user" if index % 2 == 0 else "assistant" - content = str(message).strip() - if not content: - continue - label = "Assistant" if role.lower() in {"assistant", "gpt", "bot"} else "User" - if role.lower() in {"system"}: - label = "System" - turns.append(f"{label}: {content}") - return "\n".join(turns) - - -def _extract_preset_to_jsonl( - dataset_id: str, - sample_limit: int, - cache_dir: Path, - output_jsonl: Path, - lowercase: bool, -) -> None: - """Extract one Hugging Face preset to JSONL for the parent app. - - Args: - dataset_id: Preset ID. - sample_limit: Maximum rows to extract. - cache_dir: Hugging Face cache directory. - output_jsonl: JSONL output path. - lowercase: Whether to lowercase extracted text. - """ - - preset = CONVERSATION_DATASET_PRESETS[dataset_id] - os.environ.update(_hf_subprocess_environment(cache_dir)) - print(f"Importing datasets package for {preset.label}.", flush=True) - try: - from datasets import load_dataset - except ImportError as exc: - raise RuntimeError("Install the datasets package to use Hugging Face conversation datasets.") from exc - - cache_dir.mkdir(parents=True, exist_ok=True) - output_jsonl.parent.mkdir(parents=True, exist_ok=True) - print(f"Downloading/loading {preset.label} into {cache_dir}.", flush=True) - if preset.config_name: - dataset = load_dataset(preset.hf_path, preset.config_name, split=preset.split, cache_dir=str(cache_dir)) - else: - dataset = load_dataset(preset.hf_path, split=preset.split, cache_dir=str(cache_dir)) - row_count = len(dataset) if hasattr(dataset, "__len__") else 0 - print(f"Loaded {preset.hf_path} split {preset.split} with {row_count or 'unknown'} row(s).", flush=True) - limit = row_count if sample_limit <= 0 or row_count <= 0 else min(sample_limit, row_count) - if limit and hasattr(dataset, "select"): - dataset = dataset.select(range(limit)) - loaded = 0 - with output_jsonl.open("w", encoding="utf-8") as file: - if dataset_id == "dailydialog": - loaded = _write_daily_dialog_rows(dataset, file, lowercase) - else: - for row in dataset: - text, kind = _conversation_text_from_row(dict(row)) - text = clean_text(text, lowercase=lowercase) - if not text: - continue - if preset.stage == "code": - kind = "code" - messages = _chat_messages_from_text(text, kind) - record: dict[str, Any] = {"kind": kind} - if messages: - record["messages"] = messages - else: - record["text"] = text - file.write(json.dumps(record, ensure_ascii=False) + "\n") - loaded += 1 - if loaded % 1000 == 0: - print(f"{preset.label}: extracted {loaded:,}/{limit:,} sample(s).", flush=True) - print(f"{preset.label}: wrote {loaded:,} sample(s) to {output_jsonl}.", flush=True) - - -def _write_daily_dialog_rows(dataset: Any, file: Any, lowercase: bool) -> int: - """Group DailyDialog utterance rows into dialogue samples. - - Args: - dataset: Hugging Face dataset rows. - file: Open JSONL file handle. - lowercase: Whether to lowercase extracted text. - - Returns: - Number of written dialogue samples. - """ - - dialogues: dict[str, list[str]] = {} - order: list[str] = [] - for row in dataset: - value = dict(row) - dialog_id = str(value.get("dialog_id", len(order))) - utterance = str(value.get("utterance") or "").strip() - if not utterance: - continue - if dialog_id not in dialogues: - dialogues[dialog_id] = [] - order.append(dialog_id) - dialogues[dialog_id].append(utterance) - written = 0 - for dialog_id in order: - turns = dialogues[dialog_id] - if not turns: - continue - text = "\n".join( - f"{'User' if index % 2 == 0 else 'Assistant'}: {utterance}" - for index, utterance in enumerate(turns) - ) - text = clean_text(text, lowercase=lowercase) - if not text: - continue - file.write(json.dumps({"text": text, "kind": "conversation"}, ensure_ascii=False) + "\n") - written += 1 - if written % 1000 == 0: - print(f"DailyDialog: extracted {written:,} dialogue sample(s).", flush=True) - return written - - -def main() -> None: - """Run conversation dataset helper commands.""" - - parser = argparse.ArgumentParser(description="Micro LLM conversation dataset helper") - subparsers = parser.add_subparsers(dest="command", required=True) - extract_parser = subparsers.add_parser("extract") - extract_parser.add_argument("--dataset-id", required=True, choices=sorted(CONVERSATION_DATASET_PRESETS)) - extract_parser.add_argument("--sample-limit", type=int, default=20000) - extract_parser.add_argument("--cache-dir", required=True) - extract_parser.add_argument("--output-jsonl", required=True) - extract_parser.add_argument("--config-name", default=None) - extract_parser.add_argument("--lowercase", action="store_true") - args = parser.parse_args() - if args.command == "extract": - _extract_preset_to_jsonl( - dataset_id=args.dataset_id, - sample_limit=args.sample_limit, - cache_dir=Path(args.cache_dir), - output_jsonl=Path(args.output_jsonl), - lowercase=args.lowercase, - ) - - -if __name__ == "__main__": - main() diff --git a/llm_trainer/convert_jsonl.py b/llm_trainer/convert_jsonl.py deleted file mode 100644 index 859aebd..0000000 --- a/llm_trainer/convert_jsonl.py +++ /dev/null @@ -1,205 +0,0 @@ -import json -import re -from pathlib import Path - -# ========================================================== -# Configuration -# ========================================================== - -INPUT_FOLDER = "F:\\Micro_LLM_Projects\\Nero\\training_data\\dnet_scape\\" -OUTPUT_FOLDER = "F:\\Micro_LLM_Projects\\Nero\\training_data\\dnet_scape\\Converted" -MIN_DOCUMENT_SIZE_KB = 10 -MIN_DOCUMENT_SIZE_BYTES = MIN_DOCUMENT_SIZE_KB * 1024 - -SAVE_AS_MARKDOWN = True - -Path(OUTPUT_FOLDER).mkdir(parents=True, exist_ok=True) - -# ========================================================== -# Things to remove -# ========================================================== - -REMOVE_EXACT = { - "[TABLE]", - "[/TABLE]", - "Project Data", - "Overview", - "Description", - "Background", - "Assumptions", - "Objectives", - "RACI(J)", - "Development Management", - "Testing Plan", - "Additional Development Documentation", - "Additional Close Information", - "Project Lessons Learned", - "Scope Sign Offs", - "Deliverables", - "Feasibility", - "Development", - "Closure", - "What", - "Why", - "Who", - "How", -} - -REMOVE_PREFIXES = ( - "Viewed By", - "Additional ", - "Project Ticket", - "Idea Ticket", - "Timetracking Entry", - "Lifetime Assessment", - "Initial Effort", - "Date Started", - "Date Finished", - "Last Revised", - "Key Stakeholder", - "Project Sponsor", - "Project Manager", - "Jira Project", - "Versions", - "Feature Requests", - "Migration Plan", - "Rollback plan", - "Support Team", - "Announcements", -) - -INVALID_FILENAME = r'[<>:"/\\|?*]' - -# ========================================================== -# Helpers -# ========================================================== - -def clean_filename(name: str) -> str: - name = re.sub(INVALID_FILENAME, "_", name) - name = re.sub(r"\s+", " ", name) - return name.strip()[:180] - - -def clean_text(text: str): - - if not text: - return "" - - # ----------------------------------------- - # Remove URLs - # ----------------------------------------- - text = re.sub(r"http\S+", "", text) - - # ----------------------------------------- - # Remove Windows paths - # ----------------------------------------- - text = re.sub(r"[A-Za-z]:\\[^\s]+", "", text) - - cleaned = [] - - for line in text.splitlines(): - - line = line.strip() - - if not line: - continue - - # Remove table markers - if line in ("[TABLE]", "[/TABLE]"): - continue - - # Remove all table rows - if "|" in line: - continue - - # Remove "1 flat", "2 flat" - if re.fullmatch(r"\d+\s+flat", line): - continue - - # Remove exact headings - if line in REMOVE_EXACT: - continue - - # Remove common prefixes - if any(line.startswith(prefix) for prefix in REMOVE_PREFIXES): - continue - - # Remove XML / HTML tags - if re.match(r"<.*?>", line): - continue - - # Remove Jira IDs - line = re.sub(r"\b[A-Z]+-\d+\b", "", line) - - # Collapse whitespace - line = re.sub(r"\s+", " ", line).strip() - - if len(line) < 2: - continue - - cleaned.append(line) - - # Remove excessive blank lines - text = "\n".join(cleaned) - text = re.sub(r"\n{3,}", "\n\n", text) - - return text.strip() - - -# ========================================================== -# Convert -# ========================================================== - -count = 0 - -for jsonl_file in Path(INPUT_FOLDER).glob("*.jsonl"): - - print(f"Processing {jsonl_file.name}") - - with open(jsonl_file, "r", encoding="utf-8") as f: - - for line in f: - - line = line.strip() - - if not line: - continue - - try: - obj = json.loads(line) - except Exception: - continue - - title = clean_text(obj.get("title", "Untitled")) - body = clean_text(obj.get("text", "")) - - if len(body) < 50: - continue - - if len(body.encode("utf-8")) < MIN_DOCUMENT_SIZE_BYTES: - continue - - filename = clean_filename(title) - - ext = ".md" if SAVE_AS_MARKDOWN else ".txt" - - outfile = Path(OUTPUT_FOLDER) / f"{filename}{ext}" - - if SAVE_AS_MARKDOWN: - - content = f"# {title}\n\n{body}\n" - - else: - - content = f"{title}\n\n{body}\n" - - with open(outfile, "w", encoding="utf-8") as out: - out.write(content) - - count += 1 - -print() -print("=" * 60) -print(f"Converted {count} documents.") -print(f"Saved to: {OUTPUT_FOLDER}") -print("=" * 60) \ No newline at end of file diff --git a/llm_trainer/convert_to_gguf.py b/llm_trainer/convert_to_gguf.py deleted file mode 100644 index c835257..0000000 --- a/llm_trainer/convert_to_gguf.py +++ /dev/null @@ -1,298 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- - -from __future__ import annotations - -import argparse -import logging -import os -import sys -from pathlib import Path - -import torch - -if 'NO_LOCAL_GGUF' not in os.environ: - sys.path.insert(1, str(Path(__file__).parent / 'gguf-py')) -import gguf - -from conversion import ( - ModelBase, - ModelType, - get_model_architecture, - get_model_class, - logger, - print_registered_models, - _mistral_common_installed, - _mistral_import_error_msg, -) - - -def split_str_to_n_bytes(split_str: str) -> int: - if split_str.endswith("K"): - n = int(split_str[:-1]) * 1000 - elif split_str.endswith("M"): - n = int(split_str[:-1]) * 1000 * 1000 - elif split_str.endswith("G"): - n = int(split_str[:-1]) * 1000 * 1000 * 1000 - elif split_str.isnumeric(): - n = int(split_str) - else: - raise ValueError(f"Invalid split size: {split_str}, must be a number, optionally followed by K, M, or G") - - if n < 0: - raise ValueError(f"Invalid split size: {split_str}, must be positive") - - return n - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="Convert a huggingface model to a GGML compatible file") - parser.add_argument( - "--vocab-only", action="store_true", - help="extract only the vocab", - ) - parser.add_argument( - "--outfile", type=Path, - help="path to write to; default: based on input. {ftype} will be replaced by the outtype.", - ) - parser.add_argument( - "--outtype", type=str, choices=["f32", "f16", "bf16", "q8_0", "tq1_0", "tq2_0", "auto"], default="auto", - help="output format - use f32 for float32, f16 for float16, bf16 for bfloat16, q8_0 for Q8_0, tq1_0 or tq2_0 for ternary, and auto for the highest-fidelity 16-bit float type", - ) - parser.add_argument( - "--bigendian", action="store_true", - help="model is executed on big endian machine", - ) - parser.add_argument( - "model", type=str, - help="directory containing model file or huggingface repository ID (if --remote)", - nargs="?", - ) - parser.add_argument( - "--use-temp-file", action="store_true", - help="use the tempfile library while processing (helpful when running out of memory, process killed)", - ) - parser.add_argument( - "--no-lazy", action="store_true", - help="use more RAM by computing all outputs before writing (use in case lazy evaluation is broken)", - ) - parser.add_argument( - "--model-name", type=str, default=None, - help="name of the model", - ) - parser.add_argument( - "--verbose", action="store_true", - help="increase output verbosity", - ) - parser.add_argument( - "--split-max-tensors", type=int, default=0, - help="max tensors in each split", - ) - parser.add_argument( - "--split-max-size", type=str, default="0", - help="max size per split N(M|G)", - ) - parser.add_argument( - "--dry-run", action="store_true", - help="only print out a split plan and exit, without writing any new files", - ) - parser.add_argument( - "--no-tensor-first-split", action="store_true", - help="do not add tensors to the first split (disabled by default)" - ) - parser.add_argument( - "--metadata", type=Path, - help="Specify the path for an authorship metadata override file" - ) - parser.add_argument( - "--print-supported-models", action="store_true", - help="Print the supported models" - ) - parser.add_argument( - "--remote", action="store_true", - help="(Experimental) Read safetensors file remotely without downloading to disk. Config and tokenizer files will still be downloaded. To use this feature, you need to specify Hugging Face model repo name instead of a local directory. For example: 'HuggingFaceTB/SmolLM2-1.7B-Instruct'. Note: To access gated repo, set HF_TOKEN environment variable to your Hugging Face token.", - ) - parser.add_argument( - "--mmproj", action="store_true", - help="Export multimodal projector (mmproj) for vision models. This will only work on some vision models. An 'mmproj-' prefix will be added to the output file name.", - ) - parser.add_argument( - "--mtp", action="store_true", - help="Export only the multi-token prediction (MTP) head as a separate GGUF, suitable for use as a speculative draft. An 'mtp-' prefix will be added to the output file name.", - ) - parser.add_argument( - "--no-mtp", action="store_true", - help="Exclude the multi-token prediction (MTP) head from the converted GGUF. Pair with --mtp on a second run to publish trunk and MTP as two files. Note: the split form duplicates embeddings, but even though the bundled default is more space-efficient overall, this allows differing quantization which may be more performant.", - ) - parser.add_argument( - "--mistral-format", action="store_true", - help="Whether the model is stored following the Mistral format.", - ) - parser.add_argument( - "--disable-mistral-community-chat-template", action="store_true", - help=( - "Whether to disable usage of Mistral community chat templates. If set, use the Mistral official `mistral-common` library for tokenization and detokenization of Mistral models. " - "Using `mistral-common` ensure correctness and zero-day support of tokenization for models converted from the Mistral format but requires to manually setup the tokenization server." - ) - ) - - parser.add_argument( - "--sentence-transformers-dense-modules", action="store_true", - help=("Whether to include sentence-transformers dense modules. " - "It can be used for sentence-transformers models, like google/embeddinggemma-300m. " - "Default these modules are not included.") - ) - - parser.add_argument( - "--fuse-gate-up-exps", action="store_true", - help="Fuse gate_exps and up_exps tensors into a single gate_up_exps tensor for MoE models.", - ) - parser.add_argument( - "--fp8-as-q8", action="store_true", - help="Store tensors dequantized from FP8 as Q8_0 instead of BF16/F16.", - ) - - parser.add_argument( - "--target-model-dir", type=str, default=None, - help=( - "path to the target model directory; required when converting a standalone draft model " - "(e.g. EAGLE3 / DFlash) that needs target-model metadata such as tokenizer, hidden size, and " - "layer count to populate its GGUF." - ), - ) - - args = parser.parse_args() - if not args.print_supported_models and args.model is None: - parser.error("the following arguments are required: model") - return args - - -def main() -> None: - args = parse_args() - - if args.print_supported_models: - logger.error("Supported models:") - print_registered_models() - sys.exit(0) - - if args.verbose: - logging.basicConfig(level=logging.DEBUG) - else: - logging.basicConfig(level=logging.INFO) - - if args.remote: - hf_repo_id = args.model - from huggingface_hub import snapshot_download - allowed_patterns = ["LICENSE", "*.json", "*.md", "*.txt", "tokenizer.model"] - if args.sentence_transformers_dense_modules: - # include sentence-transformers dense modules safetensors files - allowed_patterns.append("*.safetensors") - local_dir = snapshot_download( - repo_id=hf_repo_id, - allow_patterns=allowed_patterns) - dir_model = Path(local_dir) - logger.info(f"Downloaded config and tokenizer to {local_dir}") - else: - hf_repo_id = None - dir_model = Path(args.model) - - if not dir_model.is_dir(): - logger.error(f'Error: {dir_model} is not a directory') - sys.exit(1) - - ftype_map: dict[str, gguf.LlamaFileType] = { - "f32": gguf.LlamaFileType.ALL_F32, - "f16": gguf.LlamaFileType.MOSTLY_F16, - "bf16": gguf.LlamaFileType.MOSTLY_BF16, - "q8_0": gguf.LlamaFileType.MOSTLY_Q8_0, - "tq1_0": gguf.LlamaFileType.MOSTLY_TQ1_0, - "tq2_0": gguf.LlamaFileType.MOSTLY_TQ2_0, - "auto": gguf.LlamaFileType.GUESSED, - } - - is_split = args.split_max_tensors > 0 or args.split_max_size != "0" - if args.use_temp_file and is_split: - logger.error("Error: Cannot use temp file when splitting") - sys.exit(1) - - if args.outfile is not None: - fname_out = args.outfile - elif hf_repo_id: - # if remote, use the model ID as the output file name - fname_out = Path("./" + hf_repo_id.replace("/", "-") + "-{ftype}.gguf") - else: - fname_out = dir_model - - logger.info(f"Loading model: {dir_model.name}") - - is_mistral_format = args.mistral_format - if is_mistral_format and not _mistral_common_installed: - raise ImportError(_mistral_import_error_msg) - disable_mistral_community_chat_template = args.disable_mistral_community_chat_template - - with torch.inference_mode(): - output_type = ftype_map[args.outtype] - model_type = ModelType.MMPROJ if args.mmproj else ModelType.TEXT - hparams = ModelBase.load_hparams(dir_model, is_mistral_format) - if not is_mistral_format: - model_architecture = get_model_architecture(hparams, model_type) - logger.info(f"Model architecture: {model_architecture}") - try: - model_class = get_model_class(model_architecture, mmproj=(model_type == ModelType.MMPROJ)) - except NotImplementedError: - logger.error(f"Model {model_architecture} is not supported") - sys.exit(1) - elif args.mmproj: - assert hparams.get("vision_encoder") is not None, "This model does not support multimodal" - from conversion.pixtral import PixtralModel - model_class = PixtralModel - elif hparams.get("moe") is not None: - from conversion.mistral import MistralMoeModel - model_class = MistralMoeModel - else: - from conversion.mistral import MistralModel - model_class = MistralModel - - if args.mtp and args.no_mtp: - logger.error("--mtp and --no-mtp are mutually exclusive") - sys.exit(1) - - if args.mtp or args.no_mtp: - from conversion.qwen import _Qwen35MtpMixin - from conversion.step3 import Step35Model - if not (issubclass(model_class, _Qwen35MtpMixin) or issubclass(model_class, Step35Model)): - logger.error("--mtp / --no-mtp are only supported for Qwen3.5/3.6 and Step3.5 text variants today") - sys.exit(1) - if args.no_mtp: - model_class.no_mtp = True - if args.mtp: - model_class.mtp_only = True - - model_instance = model_class(dir_model, output_type, fname_out, - is_big_endian=args.bigendian, use_temp_file=args.use_temp_file, - eager=args.no_lazy, - metadata_override=args.metadata, model_name=args.model_name, - split_max_tensors=args.split_max_tensors, - split_max_size=split_str_to_n_bytes(args.split_max_size), dry_run=args.dry_run, - small_first_shard=args.no_tensor_first_split, - remote_hf_model_id=hf_repo_id, disable_mistral_community_chat_template=disable_mistral_community_chat_template, - sentence_transformers_dense_modules=args.sentence_transformers_dense_modules, - target_model_dir=Path(args.target_model_dir) if args.target_model_dir else None, - fuse_gate_up_exps=args.fuse_gate_up_exps, - fp8_as_q8=args.fp8_as_q8, - ) - - if args.vocab_only: - logger.info("Exporting model vocab...") - model_instance.write_vocab() - logger.info(f"Model vocab successfully exported to {model_instance.fname_out}") - else: - logger.info("Exporting model...") - model_instance.write() - out_path = f"{model_instance.fname_out.parent}{os.sep}" if is_split else model_instance.fname_out - logger.info(f"Model successfully exported to {out_path}") - - -if __name__ == '__main__': - main() \ No newline at end of file diff --git a/llm_trainer/coordinator/__init__.py b/llm_trainer/coordinator/__init__.py deleted file mode 100644 index 6d95bd0..0000000 --- a/llm_trainer/coordinator/__init__.py +++ /dev/null @@ -1,18 +0,0 @@ -from __future__ import annotations - -from .api_server import CoordinatorApiServer, run_coordinator_api -from .artifacts import create_job_artifact_bundle, create_result_artifact_bundle, default_artifact_root -from .job_manager import JobManager, ManagedJob, WorkerDescriptor, WorkerHeartbeat, WorkerStatus - -__all__ = [ - "CoordinatorApiServer", - "create_job_artifact_bundle", - "create_result_artifact_bundle", - "default_artifact_root", - "JobManager", - "ManagedJob", - "WorkerDescriptor", - "WorkerHeartbeat", - "WorkerStatus", - "run_coordinator_api", -] diff --git a/llm_trainer/coordinator/api_server.py b/llm_trainer/coordinator/api_server.py deleted file mode 100644 index 021f786..0000000 --- a/llm_trainer/coordinator/api_server.py +++ /dev/null @@ -1,397 +0,0 @@ -from __future__ import annotations - -import json -import mimetypes -import threading -from http import HTTPStatus -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from pathlib import Path -from typing import Any, Callable, Optional -from urllib.parse import unquote - -from llm_trainer.contracts import ( - ClaimJobRequest, - CompleteJobRequest, - FailJobRequest, - HeartbeatRequest, - ProgressReportRequest, - RegisterWorkerRequest, -) -from llm_trainer.coordinator.job_manager import JobManager -from llm_trainer.coordinator.artifacts import default_artifact_root - - -class CoordinatorApiServer: - """HTTP API wrapper around the job manager.""" - - def __init__( - self, - manager: Optional[JobManager] = None, - host: str = "127.0.0.1", - port: int = 8765, - artifact_root: Optional[Path] = None, - ) -> None: - """Create a coordinator API server. - - Args: - manager: Job manager instance. - host: Host address to bind. - port: TCP port to bind. - artifact_root: Root folder served by the artifact endpoint. - """ - - self.manager = manager or JobManager() - self.host = host - self.port = port - self.artifact_root = Path(artifact_root) if artifact_root else default_artifact_root() - self.artifact_root.mkdir(parents=True, exist_ok=True) - self.httpd: Optional[ThreadingHTTPServer] = None - self._manager_lock = threading.RLock() - - def serve_forever(self) -> None: - """Start serving coordinator API requests.""" - - handler = self._handler_class() - self.httpd = ThreadingHTTPServer((self.host, self.port), handler) - self.httpd.serve_forever() - - def shutdown(self) -> None: - """Stop the coordinator API server.""" - - if self.httpd is not None: - self.httpd.shutdown() - - def _handler_class(self) -> type[BaseHTTPRequestHandler]: - """Create a request handler bound to this server. - - Returns: - HTTP request handler class. - """ - - api = self - - class CoordinatorRequestHandler(BaseHTTPRequestHandler): - """HTTP request handler for coordinator protocol routes.""" - - server_version = "MicroLLMCoordinator/0.1" - - def do_GET(self) -> None: - """Handle GET requests.""" - - if self.path.startswith("/artifacts/"): - self._send_artifact(self.path.removeprefix("/artifacts/")) - return - routes: dict[str, Callable[[], dict[str, Any]]] = { - "/health": api._health, - "/workers": api._workers, - "/jobs": api._jobs, - } - handler = routes.get(self.path) - if handler is None: - self._send_json({"error": "not found"}, HTTPStatus.NOT_FOUND) - return - with api._manager_lock: - self._send_json(handler()) - - def do_POST(self) -> None: - """Handle POST requests.""" - - routes: dict[str, Callable[[dict[str, Any]], dict[str, Any]]] = { - "/register": api._register, - "/heartbeat": api._heartbeat, - "/claim-job": api._claim_job, - "/progress": api._progress, - "/complete": api._complete, - "/fail": api._fail, - "/pause-all": api._pause_all, - "/resume-all": api._resume_all, - "/stop-all": api._stop_all, - "/stale-workers": api._stale_workers, - } - handler = routes.get(self.path) - if handler is None: - self._send_json({"error": "not found"}, HTTPStatus.NOT_FOUND) - return - try: - payload = self._read_json() - with api._manager_lock: - response = handler(payload) - except ValueError as exc: - self._send_json({"error": str(exc)}, HTTPStatus.BAD_REQUEST) - return - except Exception as exc: - self._send_json({"error": str(exc)}, HTTPStatus.INTERNAL_SERVER_ERROR) - return - self._send_json(response) - - def do_PUT(self) -> None: - """Handle artifact upload requests.""" - - if not self.path.startswith("/artifacts/"): - self._send_json({"error": "not found"}, HTTPStatus.NOT_FOUND) - return - self._receive_artifact(self.path.removeprefix("/artifacts/")) - - def log_message(self, format: str, *args: Any) -> None: - """Silence default stderr logging.""" - - def _read_json(self) -> dict[str, Any]: - """Read JSON payload from the request. - - Returns: - Request payload. - - Raises: - ValueError: If the payload is not valid JSON. - """ - - content_length = int(self.headers.get("Content-Length", "0") or "0") - if content_length <= 0: - return {} - raw = self.rfile.read(content_length) - try: - payload = json.loads(raw.decode("utf-8")) - except json.JSONDecodeError as exc: - raise ValueError(f"Invalid JSON: {exc}") from exc - if not isinstance(payload, dict): - raise ValueError("JSON payload must be an object") - return payload - - def _send_json(self, payload: dict[str, Any], status: HTTPStatus = HTTPStatus.OK) -> None: - """Send a JSON response. - - Args: - payload: Response payload. - status: HTTP status. - """ - - body = json.dumps(payload, indent=2).encode("utf-8") - self.send_response(int(status)) - self.send_header("Content-Type", "application/json; charset=utf-8") - self.send_header("Content-Length", str(len(body))) - self.end_headers() - self.wfile.write(body) - - def _send_artifact(self, relative_url_path: str) -> None: - """Send an artifact file. - - Args: - relative_url_path: URL path relative to the artifact root. - """ - - relative_path = Path(unquote(relative_url_path)) - try: - resolved = (api.artifact_root / relative_path).resolve() - root = api.artifact_root.resolve() - if root not in resolved.parents and resolved != root: - raise ValueError("Artifact path escapes artifact root") - except ValueError: - self._send_json({"error": "invalid artifact path"}, HTTPStatus.BAD_REQUEST) - return - if not resolved.is_file(): - self._send_json({"error": "artifact not found"}, HTTPStatus.NOT_FOUND) - return - content_type = mimetypes.guess_type(str(resolved))[0] or "application/octet-stream" - self.send_response(int(HTTPStatus.OK)) - self.send_header("Content-Type", content_type) - self.send_header("Content-Length", str(resolved.stat().st_size)) - self.end_headers() - with resolved.open("rb") as artifact: - while chunk := artifact.read(1024 * 1024): - self.wfile.write(chunk) - - def _receive_artifact(self, relative_url_path: str) -> None: - """Receive an uploaded artifact file. - - Args: - relative_url_path: URL path relative to the artifact root. - """ - - relative_path = Path(unquote(relative_url_path)) - try: - resolved = (api.artifact_root / relative_path).resolve() - root = api.artifact_root.resolve() - if root not in resolved.parents and resolved != root: - raise ValueError("Artifact path escapes artifact root") - except ValueError: - self._send_json({"error": "invalid artifact path"}, HTTPStatus.BAD_REQUEST) - return - content_length = int(self.headers.get("Content-Length", "0") or "0") - if content_length <= 0: - self._send_json({"error": "empty artifact upload"}, HTTPStatus.BAD_REQUEST) - return - resolved.parent.mkdir(parents=True, exist_ok=True) - remaining = content_length - with resolved.open("wb") as artifact: - while remaining > 0: - chunk = self.rfile.read(min(1024 * 1024, remaining)) - if not chunk: - break - artifact.write(chunk) - remaining -= len(chunk) - if remaining: - self._send_json({"error": "incomplete artifact upload"}, HTTPStatus.BAD_REQUEST) - return - self._send_json({"status": "ok", "artifact_url": f"/artifacts/{relative_path.as_posix()}"}) - - return CoordinatorRequestHandler - - def _health(self) -> dict[str, Any]: - """Return server health. - - Returns: - Health payload. - """ - - return {"status": "ok", "workers": len(self.manager.list_workers()), "jobs": len(self.manager.list_jobs())} - - def _workers(self) -> dict[str, Any]: - """Return workers. - - Returns: - Worker payload. - """ - - return {"workers": [worker.to_jsonable() for worker in self.manager.list_workers()]} - - def _jobs(self) -> dict[str, Any]: - """Return jobs. - - Returns: - Job payload. - """ - - return {"jobs": [job.to_jsonable() for job in self.manager.list_jobs()]} - - def _register(self, payload: dict[str, Any]) -> dict[str, Any]: - """Register a worker. - - Args: - payload: Register worker request payload. - - Returns: - Register worker response payload. - """ - - return self.manager.register_remote_worker(RegisterWorkerRequest.from_jsonable(payload)).to_jsonable() - - def _heartbeat(self, payload: dict[str, Any]) -> dict[str, Any]: - """Handle worker heartbeat. - - Args: - payload: Heartbeat request payload. - - Returns: - Heartbeat response payload. - """ - - return self.manager.handle_heartbeat(HeartbeatRequest.from_jsonable(payload)).to_jsonable() - - def _claim_job(self, payload: dict[str, Any]) -> dict[str, Any]: - """Handle job claim. - - Args: - payload: Claim job request payload. - - Returns: - Claim job response payload. - """ - - return self.manager.handle_claim_job(ClaimJobRequest.from_jsonable(payload)).to_jsonable() - - def _progress(self, payload: dict[str, Any]) -> dict[str, Any]: - """Handle progress report. - - Args: - payload: Progress report request payload. - - Returns: - Progress response payload. - """ - - return self.manager.handle_progress_report(ProgressReportRequest.from_jsonable(payload)).to_jsonable() - - def _complete(self, payload: dict[str, Any]) -> dict[str, Any]: - """Handle job completion. - - Args: - payload: Complete job request payload. - - Returns: - Completion response payload. - """ - - return self.manager.handle_complete_job(CompleteJobRequest.from_jsonable(payload)).to_jsonable() - - def _fail(self, payload: dict[str, Any]) -> dict[str, Any]: - """Handle job failure. - - Args: - payload: Fail job request payload. - - Returns: - Failure response payload. - """ - - return self.manager.handle_fail_job(FailJobRequest.from_jsonable(payload)).to_jsonable() - - def _pause_all(self, payload: dict[str, Any]) -> dict[str, Any]: - """Pause all jobs. - - Args: - payload: Ignored payload. - - Returns: - Pause summary. - """ - - return {"paused": self.manager.pause_all_jobs()} - - def _resume_all(self, payload: dict[str, Any]) -> dict[str, Any]: - """Resume all jobs. - - Args: - payload: Ignored payload. - - Returns: - Resume summary. - """ - - return {"resumed": self.manager.resume_all_jobs()} - - def _stop_all(self, payload: dict[str, Any]) -> dict[str, Any]: - """Stop all jobs. - - Args: - payload: Ignored payload. - - Returns: - Stop summary. - """ - - return {"stopping": self.manager.stop_all_jobs()} - - def _stale_workers(self, payload: dict[str, Any]) -> dict[str, Any]: - """Mark stale workers offline. - - Args: - payload: Payload with optional timeout_seconds. - - Returns: - Stale worker summary. - """ - - timeout_seconds = int(payload.get("timeout_seconds", 30)) - return {"offline_workers": self.manager.mark_stale_workers_offline(timeout_seconds=timeout_seconds)} - - -def run_coordinator_api(host: str = "127.0.0.1", port: int = 8765, artifact_root: Optional[Path] = None) -> None: - """Run the coordinator API server. - - Args: - host: Host address to bind. - port: TCP port to bind. - artifact_root: Root folder served by the artifact endpoint. - """ - - CoordinatorApiServer(host=host, port=port, artifact_root=artifact_root).serve_forever() diff --git a/llm_trainer/coordinator/artifacts.py b/llm_trainer/coordinator/artifacts.py deleted file mode 100644 index b103bd6..0000000 --- a/llm_trainer/coordinator/artifacts.py +++ /dev/null @@ -1,97 +0,0 @@ -from __future__ import annotations - -import json -import zipfile -from pathlib import Path -from typing import Optional - -from llm_trainer.contracts import TrainingJobSpec - - -def default_artifact_root() -> Path: - """Return the default artifact serving root. - - Returns: - Artifact root path. - """ - - return Path.home() / ".drunkenbot_ide" / "artifacts" - - -def create_job_artifact_bundle( - job: TrainingJobSpec, - artifact_root: Optional[Path] = None, - base_url: str = "/artifacts", -) -> Path: - """Create a portable artifact bundle for a training job. - - The bundle contains the prepared dataset under ``dataset/`` and a - ``job.json`` manifest. Workers extract it into their local workspace and - rewrite dataset/output paths before training. - - Args: - job: Training job specification. - artifact_root: Root folder where bundles are written. - base_url: URL prefix used by the coordinator artifact route. - - Returns: - Bundle path. - """ - - root = Path(artifact_root) if artifact_root else default_artifact_root() - root.mkdir(parents=True, exist_ok=True) - bundle_path = root / f"{job.job_id}.zip" - with zipfile.ZipFile(bundle_path, "w", compression=zipfile.ZIP_DEFLATED) as archive: - _write_directory(archive, job.dataset.dataset_dir, "dataset") - resume_path = job.training.resume_from_checkpoint - if resume_path and Path(resume_path).is_file(): - archive.write(Path(resume_path), "checkpoints/resume_checkpoint.pt") - job.metadata["resume_checkpoint_artifact"] = "checkpoints/resume_checkpoint.pt" - base_path = job.training.fine_tune_from_checkpoint or job.model.base_checkpoint - if base_path and Path(base_path).is_file(): - archive.write(Path(base_path), "checkpoints/base_checkpoint.pt") - job.metadata["base_checkpoint_artifact"] = "checkpoints/base_checkpoint.pt" - archive.writestr("job.json", json.dumps(job.to_jsonable(), indent=2)) - job.metadata["artifact_bundle_url"] = f"{base_url.rstrip('/')}/{bundle_path.name}" - job.metadata["artifact_bundle_name"] = bundle_path.name - return bundle_path - - -def create_result_artifact_bundle(job_id: str, output_dir: Path, bundle_path: Path) -> Path: - """Create a portable output artifact bundle for a completed job. - - Args: - job_id: Training job identifier. - output_dir: Worker-local output directory to bundle. - bundle_path: Destination zip path. - - Returns: - Bundle path. - """ - - bundle_path = Path(bundle_path) - bundle_path.parent.mkdir(parents=True, exist_ok=True) - with zipfile.ZipFile(bundle_path, "w", compression=zipfile.ZIP_DEFLATED) as archive: - _write_directory(archive, output_dir, "model") - archive.writestr("result.json", json.dumps({"job_id": job_id}, indent=2)) - return bundle_path - - -def _write_directory(archive: zipfile.ZipFile, source: Path, archive_root: str) -> None: - """Write a directory into a zip archive. - - Args: - archive: Zip archive. - source: Source directory. - archive_root: Archive root folder. - - Raises: - FileNotFoundError: If the source directory does not exist. - """ - - source = Path(source) - if not source.exists(): - raise FileNotFoundError(f"Artifact source not found: {source}") - for path in source.rglob("*"): - if path.is_file(): - archive.write(path, Path(archive_root) / path.relative_to(source)) diff --git a/llm_trainer/coordinator/job_manager.py b/llm_trainer/coordinator/job_manager.py deleted file mode 100644 index f1b7b49..0000000 --- a/llm_trainer/coordinator/job_manager.py +++ /dev/null @@ -1,976 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass, field -from datetime import datetime, timedelta -from enum import Enum -from pathlib import Path -from typing import Any, Optional - -from llm_trainer.backends.base import ProgressCallback, StopCallback, TrainerBackend -from llm_trainer.backends.registry import DEFAULT_BACKEND_REGISTRY, BackendRegistry -from llm_trainer.contracts import ( - BackendKind, - ClaimJobRequest, - ClaimJobResponse, - CompleteJobRequest, - CompleteJobResponse, - FailJobRequest, - FailJobResponse, - HeartbeatRequest, - HeartbeatResponse, - JobStatus, - ProgressReportRequest, - ProgressReportResponse, - ProtocolStatus, - RegisterWorkerRequest, - RegisterWorkerResponse, - TrainingMetrics, - TrainingJobSpec, - TrainingResultSpec, - WorkerAvailability, - utc_now_iso, -) -from llm_trainer.coordinator.state_store import JobStateStore -from llm_trainer.training import TrainingResult - - -class WorkerStatus(str, Enum): - """Worker availability state.""" - - AVAILABLE = "available" - BUSY = "busy" - OFFLINE = "offline" - - -@dataclass -class WorkerDescriptor: - """Training worker registered with the job manager. - - Attributes: - worker_id: Stable worker identifier. - backend: Backend kind this worker can execute. - status: Current availability state. - device: Device advertised by the worker. - hostname: Optional worker host name. - capabilities: Free-form hardware/runtime capabilities. - """ - - worker_id: str - backend: BackendKind = BackendKind.LOCAL - status: WorkerStatus = WorkerStatus.AVAILABLE - device: str = "auto" - hostname: Optional[str] = None - capabilities: dict[str, Any] = field(default_factory=dict) - last_heartbeat_at: Optional[str] = None - - def to_jsonable(self) -> dict[str, Any]: - """Convert the worker descriptor to JSON-friendly values. - - Returns: - Serializable dictionary. - """ - - return { - "worker_id": self.worker_id, - "backend": self.backend.value, - "status": self.status.value, - "device": self.device, - "hostname": self.hostname, - "capabilities": self.capabilities, - "last_heartbeat_at": self.last_heartbeat_at, - } - - @classmethod - def from_jsonable(cls, data: dict[str, Any]) -> "WorkerDescriptor": - """Create a worker descriptor from JSON-friendly values. - - Args: - data: Serialized worker data. - - Returns: - Worker descriptor. - """ - - return cls( - worker_id=data["worker_id"], - backend=BackendKind(data.get("backend", BackendKind.LOCAL.value)), - status=WorkerStatus(data.get("status", WorkerStatus.AVAILABLE.value)), - device=data.get("device", "auto"), - hostname=data.get("hostname"), - capabilities=dict(data.get("capabilities") or {}), - last_heartbeat_at=data.get("last_heartbeat_at"), - ) - - -@dataclass -class WorkerHeartbeat: - """Heartbeat reported by a training worker. - - Attributes: - worker_id: Worker identifier. - status: Worker availability state. - backend: Backend kind the worker can execute. - active_job_id: Job currently running on the worker. - device: Worker device. - metrics: Runtime metrics reported by the worker. - timestamp: UTC heartbeat timestamp. - """ - - worker_id: str - status: WorkerStatus - backend: BackendKind = BackendKind.LOCAL - active_job_id: Optional[str] = None - device: str = "auto" - metrics: dict[str, Any] = field(default_factory=dict) - timestamp: str = field(default_factory=utc_now_iso) - - def to_jsonable(self) -> dict[str, Any]: - """Convert the heartbeat to JSON-friendly values. - - Returns: - Serializable dictionary. - """ - - return { - "worker_id": self.worker_id, - "status": self.status.value, - "backend": self.backend.value, - "active_job_id": self.active_job_id, - "device": self.device, - "metrics": self.metrics, - "timestamp": self.timestamp, - } - - -@dataclass -class ManagedJob: - """Job state tracked by the manager. - - Attributes: - spec: Training job contract. - assigned_worker_id: Worker currently assigned to this job. - result: Serializable result metadata when finished. - error: Error text when failed. - """ - - spec: TrainingJobSpec - assigned_worker_id: Optional[str] = None - result: Optional[TrainingResultSpec] = None - error: Optional[str] = None - latest_metrics: Optional[TrainingMetrics] = None - updated_at: str = field(default_factory=utc_now_iso) - - def to_jsonable(self) -> dict[str, Any]: - """Convert the managed job to JSON-friendly values. - - Returns: - Serializable dictionary. - """ - - return { - "spec": self.spec.to_jsonable(), - "assigned_worker_id": self.assigned_worker_id, - "result": _jsonable_result(self.result) if self.result else None, - "error": self.error, - "latest_metrics": self.latest_metrics.__dict__ if self.latest_metrics else None, - "updated_at": self.updated_at, - } - - -class JobManager: - """Coordinates training jobs across available backends. - - The first implementation runs local jobs in-process. The API is intentionally - shaped like a distributed coordinator so remote clients can be added without - rewriting the desktop app's training calls. - """ - - def __init__( - self, - registry: Optional[BackendRegistry] = None, - state_store: Optional[JobStateStore] = None, - ) -> None: - """Create a job manager. - - Args: - registry: Backend registry used to resolve job backends. - state_store: Optional persistent state store. - """ - - self.registry = registry or DEFAULT_BACKEND_REGISTRY - self.state_store = state_store or JobStateStore() - self._jobs: dict[str, ManagedJob] = {} - self._workers: dict[str, WorkerDescriptor] = {} - self._stop_requested: set[str] = set() - self._restore_state() - self.register_worker( - WorkerDescriptor( - worker_id="local", - backend=BackendKind.LOCAL, - status=WorkerStatus.AVAILABLE, - device="auto", - hostname="localhost", - ) - ) - - def register_worker(self, worker: WorkerDescriptor) -> None: - """Register or update a worker. - - Args: - worker: Worker descriptor. - """ - - self._workers[worker.worker_id] = worker - self._persist_worker(worker) - - def register_remote_worker(self, request: RegisterWorkerRequest) -> RegisterWorkerResponse: - """Register a remote worker from a protocol request. - - Args: - request: Register worker request. - - Returns: - Register worker response. - """ - - if not request.worker_id.strip(): - return RegisterWorkerResponse( - worker_id=request.worker_id, - accepted=False, - status=ProtocolStatus.REJECTED, - message="worker_id is required", - ) - capabilities = request.capabilities.to_jsonable() - if request.labels: - capabilities["labels"] = request.labels - self.register_worker( - WorkerDescriptor( - worker_id=request.worker_id, - backend=request.backend, - status=WorkerStatus.AVAILABLE, - device=request.device, - hostname=request.capabilities.hostname, - capabilities=capabilities, - last_heartbeat_at=utc_now_iso(), - ) - ) - return RegisterWorkerResponse( - worker_id=request.worker_id, - accepted=True, - status=ProtocolStatus.OK, - heartbeat_interval_seconds=10, - message="worker registered", - ) - - def record_heartbeat(self, heartbeat: WorkerHeartbeat) -> None: - """Record a worker heartbeat and update worker availability. - - Args: - heartbeat: Worker heartbeat. - """ - - worker = self._workers.get( - heartbeat.worker_id, - WorkerDescriptor( - worker_id=heartbeat.worker_id, - backend=heartbeat.backend, - device=heartbeat.device, - ), - ) - worker.backend = heartbeat.backend - worker.status = heartbeat.status - worker.device = heartbeat.device - worker.last_heartbeat_at = heartbeat.timestamp - self._workers[worker.worker_id] = worker - self._persist_worker(worker) - self.state_store.record_heartbeat(worker.worker_id, heartbeat.to_jsonable()) - - def handle_heartbeat(self, request: HeartbeatRequest) -> HeartbeatResponse: - """Handle a protocol heartbeat request from a worker. - - Args: - request: Heartbeat request. - - Returns: - Heartbeat response. - """ - - if not request.worker_id.strip(): - return HeartbeatResponse( - status=ProtocolStatus.REJECTED, - should_stop_job=True, - message="worker_id is required", - ) - heartbeat = WorkerHeartbeat( - worker_id=request.worker_id, - status=_availability_to_worker_status(request.availability), - backend=request.backend, - active_job_id=request.active_job_id, - device=request.device, - metrics=request.metrics, - timestamp=request.sent_at, - ) - self.record_heartbeat(heartbeat) - should_stop = bool(request.active_job_id and self._should_stop_remote_job(request.active_job_id)) - should_pause = bool(request.active_job_id and self._should_pause_remote_job(request.active_job_id)) - return HeartbeatResponse( - status=ProtocolStatus.OK, - should_stop_job=should_stop, - should_pause_job=should_pause, - message=_control_message(should_stop, should_pause, "heartbeat accepted"), - ) - - def handle_claim_job(self, request: ClaimJobRequest) -> ClaimJobResponse: - """Assign a queued job to a worker when compatible work exists. - - Args: - request: Claim job request. - - Returns: - Claim job response containing an assigned job when available. - """ - - if not request.worker_id.strip(): - return ClaimJobResponse(status=ProtocolStatus.REJECTED, message="worker_id is required") - worker = self._workers.get( - request.worker_id, - WorkerDescriptor( - worker_id=request.worker_id, - backend=request.backend, - status=WorkerStatus.AVAILABLE, - hostname=request.capabilities.hostname, - capabilities=request.capabilities.to_jsonable(), - ), - ) - worker.backend = request.backend - worker.status = WorkerStatus.AVAILABLE - worker.hostname = request.capabilities.hostname or worker.hostname - worker.capabilities.update(request.capabilities.to_jsonable()) - worker.last_heartbeat_at = utc_now_iso() - self._workers[worker.worker_id] = worker - self._persist_worker(worker) - - for managed in self._jobs.values(): - job = managed.spec - if not self._worker_can_claim_job(worker, job): - continue - managed.assigned_worker_id = worker.worker_id - worker.status = WorkerStatus.BUSY - job.status = JobStatus.ASSIGNED - self._persist_worker(worker) - self._persist_job(job.job_id) - return ClaimJobResponse(job=job, status=ProtocolStatus.OK, message="job assigned") - return ClaimJobResponse(status=ProtocolStatus.OK, message="no compatible queued job") - - def handle_progress_report(self, request: ProgressReportRequest) -> ProgressReportResponse: - """Handle progress metrics from a worker. - - Args: - request: Progress report request. - - Returns: - Progress report response. - """ - - managed = self._jobs.get(request.job_id) - if managed is None: - return ProgressReportResponse( - status=ProtocolStatus.REJECTED, - should_stop_job=True, - message="unknown job", - ) - if managed.assigned_worker_id and managed.assigned_worker_id != request.worker_id: - return ProgressReportResponse( - status=ProtocolStatus.REJECTED, - should_stop_job=True, - message="job is assigned to a different worker", - ) - managed.assigned_worker_id = request.worker_id - managed.latest_metrics = request.metrics - managed.updated_at = request.sent_at - if managed.spec.status == JobStatus.ASSIGNED: - managed.spec.status = JobStatus.RUNNING - worker = self._workers.get(request.worker_id) - if worker: - worker.status = WorkerStatus.BUSY - worker.last_heartbeat_at = request.sent_at - self._persist_worker(worker) - self._persist_job(request.job_id) - should_stop = self._should_stop_remote_job(request.job_id) - should_pause = self._should_pause_remote_job(request.job_id) - return ProgressReportResponse( - status=ProtocolStatus.OK, - should_stop_job=should_stop, - should_pause_job=should_pause, - message=_control_message(should_stop, should_pause, "progress accepted"), - ) - - def handle_complete_job(self, request: CompleteJobRequest) -> CompleteJobResponse: - """Handle successful remote job completion. - - Args: - request: Complete job request. - - Returns: - Complete job response. - """ - - if request.result is None: - return CompleteJobResponse(status=ProtocolStatus.REJECTED, message="result is required") - managed = self._jobs.get(request.result.job_id) - if managed is None: - return CompleteJobResponse(status=ProtocolStatus.REJECTED, message="unknown job") - managed.assigned_worker_id = request.worker_id - managed.result = request.result - managed.spec.status = request.result.status - managed.updated_at = request.sent_at - managed.error = request.result.error - worker = self._workers.get(request.worker_id) - if worker: - worker.status = WorkerStatus.AVAILABLE - worker.last_heartbeat_at = request.sent_at - self._persist_worker(worker) - self._persist_job(request.result.job_id) - self._stop_requested.discard(request.result.job_id) - return CompleteJobResponse(status=ProtocolStatus.OK, message="job completion accepted") - - def handle_fail_job(self, request: FailJobRequest) -> FailJobResponse: - """Handle remote job failure. - - Args: - request: Fail job request. - - Returns: - Fail job response. - """ - - managed = self._jobs.get(request.job_id) - if managed is None: - return FailJobResponse(status=ProtocolStatus.REJECTED, message="unknown job") - managed.assigned_worker_id = None if request.retryable else request.worker_id - managed.spec.status = JobStatus.QUEUED if request.retryable else JobStatus.FAILED - managed.error = request.error - managed.updated_at = request.sent_at - worker = self._workers.get(request.worker_id) - if worker: - worker.status = WorkerStatus.AVAILABLE - worker.last_heartbeat_at = request.sent_at - self._persist_worker(worker) - self._persist_job(request.job_id) - self._stop_requested.discard(request.job_id) - return FailJobResponse( - status=ProtocolStatus.OK, - message="job requeued after failure" if request.retryable else "job failure accepted", - ) - - def mark_stale_workers_offline(self, timeout_seconds: int = 30) -> list[str]: - """Mark workers offline when their last heartbeat is too old. - - Args: - timeout_seconds: Age in seconds after which a worker is stale. - - Returns: - Worker IDs marked offline. - """ - - marked: list[str] = [] - cutoff = datetime.now().astimezone() - timedelta(seconds=timeout_seconds) - for worker in self._workers.values(): - if worker.worker_id == "local" or worker.status == WorkerStatus.OFFLINE: - continue - heartbeat_at = _parse_timestamp(worker.last_heartbeat_at) - if heartbeat_at and heartbeat_at < cutoff: - worker.status = WorkerStatus.OFFLINE - self._persist_worker(worker) - marked.append(worker.worker_id) - return marked - - def list_workers(self) -> list[WorkerDescriptor]: - """Return known workers. - - Returns: - Registered workers. - """ - - return list(self._workers.values()) - - def submit(self, job: TrainingJobSpec) -> str: - """Submit a job to the manager queue. - - Args: - job: Training job contract. - - Returns: - Job identifier. - """ - - job.status = JobStatus.QUEUED - self._jobs[job.job_id] = ManagedJob(spec=job) - self._persist_job(job.job_id) - return job.job_id - - def get_job(self, job_id: str) -> ManagedJob: - """Return a managed job by ID. - - Args: - job_id: Job identifier. - - Returns: - Managed job. - - Raises: - KeyError: If the job is unknown. - """ - - return self._jobs[job_id] - - def list_jobs(self) -> list[ManagedJob]: - """Return tracked jobs. - - Returns: - Managed jobs. - """ - - return list(self._jobs.values()) - - def cancel(self, job_id: str) -> None: - """Request cooperative cancellation for a job. - - Args: - job_id: Job identifier. - """ - - managed = self.get_job(job_id) - managed.spec.status = JobStatus.STOPPING - self._stop_requested.add(job_id) - self._persist_job(job_id) - - def stop_all_jobs(self) -> int: - """Request cooperative stop for all active jobs. - - Returns: - Number of jobs marked for stopping. - """ - - count = 0 - for managed in self._jobs.values(): - if managed.spec.status in {JobStatus.QUEUED, JobStatus.ASSIGNED, JobStatus.RUNNING, JobStatus.PAUSED}: - managed.spec.status = JobStatus.STOPPING - self._stop_requested.add(managed.spec.job_id) - self._persist_job(managed.spec.job_id) - count += 1 - return count - - def pause_all_jobs(self) -> int: - """Pause queued or active remote jobs. - - Returns: - Number of jobs marked paused. - """ - - count = 0 - for managed in self._jobs.values(): - if managed.spec.status in {JobStatus.QUEUED, JobStatus.ASSIGNED, JobStatus.RUNNING}: - managed.spec.status = JobStatus.PAUSED - self._persist_job(managed.spec.job_id) - count += 1 - return count - - def resume_all_jobs(self) -> int: - """Resume paused jobs by returning them to the queue. - - Returns: - Number of jobs resumed. - """ - - count = 0 - for managed in self._jobs.values(): - if managed.spec.status == JobStatus.PAUSED: - managed.spec.status = JobStatus.QUEUED - managed.assigned_worker_id = None - self._persist_job(managed.spec.job_id) - count += 1 - return count - - def run_next( - self, - progress: Optional[ProgressCallback] = None, - should_stop: Optional[StopCallback] = None, - ) -> TrainingResult: - """Run the next queued job. - - Args: - progress: Optional progress callback. - should_stop: Optional external stop callback. - - Returns: - Training result. - - Raises: - ValueError: If no queued job is available. - """ - - for managed in self._jobs.values(): - if managed.spec.status == JobStatus.QUEUED: - return self.run_job(managed.spec.job_id, progress=progress, should_stop=should_stop) - raise ValueError("No queued training jobs are available") - - def run_job( - self, - job_id: str, - progress: Optional[ProgressCallback] = None, - should_stop: Optional[StopCallback] = None, - ) -> TrainingResult: - """Run one job through an available backend. - - Args: - job_id: Job identifier. - progress: Optional progress callback. - should_stop: Optional external stop callback. - - Returns: - Training result. - """ - - managed = self.get_job(job_id) - job = managed.spec - worker = self._select_worker(job.runtime.backend) - managed.assigned_worker_id = worker.worker_id - worker.status = WorkerStatus.BUSY - job.status = JobStatus.ASSIGNED - self._persist_worker(worker) - self._persist_job(job_id) - backend = self.registry.get(job.runtime.backend) - try: - result = backend.run( - job, - progress=progress, - should_stop=lambda: self._should_stop(job_id, should_stop), - ) - managed.result = self._result_spec(job, result) - job.status = managed.result.status - self._persist_job(job_id) - return result - except Exception as exc: - job.status = JobStatus.FAILED - managed.error = str(exc) - self._persist_job(job_id) - raise - finally: - worker.status = WorkerStatus.AVAILABLE - worker.last_heartbeat_at = utc_now_iso() - self._persist_worker(worker) - self._stop_requested.discard(job_id) - - def _select_worker(self, backend: BackendKind) -> WorkerDescriptor: - """Select an available worker for a backend. - - Args: - backend: Backend kind required by the job. - - Returns: - Available worker. - - Raises: - ValueError: If no worker is available. - """ - - for worker in self._workers.values(): - if worker.backend == backend and worker.status == WorkerStatus.AVAILABLE: - return worker - raise ValueError(f"No available worker for backend {backend.value}") - - def _worker_can_claim_job(self, worker: WorkerDescriptor, job: TrainingJobSpec) -> bool: - """Return whether a remote worker can claim a job. - - Args: - worker: Worker descriptor. - job: Training job contract. - - Returns: - Whether the worker can claim the job. - """ - - if job.status != JobStatus.QUEUED: - return False - if job.runtime.backend != worker.backend: - return False - if job.runtime.preferred_worker_id and job.runtime.preferred_worker_id != worker.worker_id: - return False - if job.runtime.min_vram_gb is not None: - worker_vram = _worker_total_vram_gb(worker) - if worker_vram is None or worker_vram < job.runtime.min_vram_gb: - return False - if job.runtime.tags: - worker_labels = set(worker.capabilities.get("labels") or []) - if not set(job.runtime.tags).issubset(worker_labels): - return False - return True - - def _should_stop(self, job_id: str, external_stop: Optional[StopCallback]) -> bool: - """Return whether a job should stop. - - Args: - job_id: Job identifier. - external_stop: Optional external stop callback. - - Returns: - Whether the job should stop. - """ - - return job_id in self._stop_requested or bool(external_stop and external_stop()) - - def _should_stop_remote_job(self, job_id: str) -> bool: - """Return whether a remote worker should stop a job. - - Args: - job_id: Job identifier. - - Returns: - Whether the worker should stop the job. - """ - - if job_id in self._stop_requested: - return True - managed = self._jobs.get(job_id) - return bool(managed and managed.spec.status in {JobStatus.STOPPING, JobStatus.CANCELLED, JobStatus.FAILED}) - - def _should_pause_remote_job(self, job_id: str) -> bool: - """Return whether a remote worker should pause a job. - - Args: - job_id: Job identifier. - - Returns: - Whether the worker should pause the job. - """ - - managed = self._jobs.get(job_id) - return bool(managed and managed.spec.status == JobStatus.PAUSED) - - def _result_spec(self, job: TrainingJobSpec, result: TrainingResult) -> TrainingResultSpec: - """Create a serializable result contract from a training result. - - Args: - job: Training job contract. - result: Concrete training result. - - Returns: - Serializable result specification. - """ - - return TrainingResultSpec( - job_id=job.job_id, - status=JobStatus.CANCELLED if result.stopped else JobStatus.COMPLETED, - 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, - ) - - def _restore_state(self) -> None: - """Restore persisted jobs and workers from the state store.""" - - for worker_data in self.state_store.load_workers(): - worker = WorkerDescriptor.from_jsonable(worker_data) - if worker.status == WorkerStatus.BUSY: - worker.status = WorkerStatus.OFFLINE - self._workers[worker.worker_id] = worker - for job_data in self.state_store.load_jobs(): - managed = _managed_job_from_jsonable(job_data) - if managed.spec.status in {JobStatus.ASSIGNED, JobStatus.RUNNING, JobStatus.STOPPING}: - managed.spec.status = JobStatus.QUEUED - managed.assigned_worker_id = None - managed.error = "Recovered after app restart before job completion." - self._jobs[managed.spec.job_id] = managed - self._persist_job(managed.spec.job_id) - - def _persist_job(self, job_id: str) -> None: - """Persist a managed job. - - Args: - job_id: Job identifier. - """ - - managed = self._jobs.get(job_id) - if managed: - self.state_store.save_job(job_id, managed.spec.status.value, managed.to_jsonable()) - - def _persist_worker(self, worker: WorkerDescriptor) -> None: - """Persist a worker descriptor. - - Args: - worker: Worker descriptor. - """ - - self.state_store.save_worker(worker.worker_id, worker.status.value, worker.to_jsonable()) - - -def run_local_job( - job: TrainingJobSpec, - backend: Optional[TrainerBackend] = None, - progress: Optional[ProgressCallback] = None, - should_stop: Optional[StopCallback] = None, -) -> TrainingResult: - """Run a job through a temporary local manager. - - Args: - job: Training job contract. - backend: Optional backend override for tests or embedded use. - progress: Optional progress callback. - should_stop: Optional cooperative cancellation callback. - - Returns: - Training result. - """ - - registry = BackendRegistry() - if backend is not None: - registry.register(job.runtime.backend, backend) - manager = JobManager(registry=registry) - manager.submit(job) - return manager.run_job(job.job_id, progress=progress, should_stop=should_stop) - - -def _jsonable_result(result: TrainingResultSpec) -> dict[str, Any]: - """Convert a result spec to JSON-friendly values. - - Args: - result: Training result specification. - - Returns: - Serializable dictionary. - """ - - output: dict[str, Any] = {} - for key, value in result.__dict__.items(): - if isinstance(value, Path): - output[key] = str(value) - elif isinstance(value, Enum): - output[key] = value.value - else: - output[key] = value - return output - - -def _managed_job_from_jsonable(data: dict[str, Any]) -> ManagedJob: - """Create a managed job from JSON-friendly values. - - Args: - data: Serialized managed job payload. - - Returns: - Managed job. - """ - - result_data = data.get("result") - return ManagedJob( - spec=TrainingJobSpec.from_jsonable(data["spec"]), - assigned_worker_id=data.get("assigned_worker_id"), - result=_result_from_jsonable(result_data) if result_data else None, - error=data.get("error"), - latest_metrics=TrainingMetrics(**dict(data.get("latest_metrics") or {})) if data.get("latest_metrics") else None, - updated_at=data.get("updated_at", utc_now_iso()), - ) - - -def _result_from_jsonable(data: dict[str, Any]) -> TrainingResultSpec: - """Create a training result spec from JSON-friendly values. - - Args: - data: Serialized result data. - - Returns: - Training result specification. - """ - - checkpoint_path = data.get("checkpoint_path") - summary_path = data.get("summary_path") - return TrainingResultSpec( - job_id=data["job_id"], - status=JobStatus(data["status"]), - checkpoint_path=Path(checkpoint_path) if checkpoint_path else None, - summary_path=Path(summary_path) if summary_path else None, - final_train_loss=data.get("final_train_loss"), - final_val_loss=data.get("final_val_loss"), - stopped=bool(data.get("stopped")), - error=data.get("error"), - artifact_bundle_url=data.get("artifact_bundle_url"), - ) - - -def _availability_to_worker_status(availability: WorkerAvailability) -> WorkerStatus: - """Convert protocol availability to manager worker status. - - Args: - availability: Protocol worker availability. - - Returns: - Manager worker status. - """ - - if availability == WorkerAvailability.BUSY: - return WorkerStatus.BUSY - if availability == WorkerAvailability.OFFLINE: - return WorkerStatus.OFFLINE - return WorkerStatus.AVAILABLE - - -def _parse_timestamp(value: Optional[str]) -> Optional[datetime]: - """Parse an ISO timestamp. - - Args: - value: ISO timestamp. - - Returns: - Parsed timestamp when valid. - """ - - if not value: - return None - try: - parsed = datetime.fromisoformat(value) - except ValueError: - return None - if parsed.tzinfo is None: - return parsed.astimezone() - return parsed - - -def _worker_total_vram_gb(worker: WorkerDescriptor) -> Optional[float]: - """Return worker total VRAM in GB when known. - - Args: - worker: Worker descriptor. - - Returns: - Total VRAM in GB. - """ - - value = worker.capabilities.get("total_vram_gb") - if value is None: - return None - try: - return float(value) - except (TypeError, ValueError): - return None - - -def _control_message(should_stop: bool, should_pause: bool, default: str) -> str: - """Return a human-readable control message. - - Args: - should_stop: Whether stop was requested. - should_pause: Whether pause was requested. - default: Default message. - - Returns: - Control message. - """ - - if should_stop: - return "stop requested" - if should_pause: - return "pause requested" - return default diff --git a/llm_trainer/coordinator/state_store.py b/llm_trainer/coordinator/state_store.py deleted file mode 100644 index 8fab5e8..0000000 --- a/llm_trainer/coordinator/state_store.py +++ /dev/null @@ -1,188 +0,0 @@ -from __future__ import annotations - -import json -import sqlite3 -from contextlib import closing -from pathlib import Path -from typing import Any, Optional - -from llm_trainer.contracts import utc_now_iso - - -def default_state_db_path() -> Path: - """Return the default coordinator state database path. - - Returns: - Default SQLite database path. - """ - - return Path.home() / ".drunkenbot_ide" / "coordinator_state.sqlite3" - - -class JobStateStore: - """SQLite-backed state store for coordinator jobs and workers.""" - - def __init__(self, db_path: Optional[Path] = None) -> None: - """Create a job state store. - - Args: - db_path: SQLite file path. Defaults to the user's app data folder. - """ - - self.db_path = Path(db_path) if db_path else default_state_db_path() - self.db_path.parent.mkdir(parents=True, exist_ok=True) - self._initialize() - - def save_job(self, job_id: str, status: str, payload: dict[str, Any]) -> None: - """Save or replace a managed job record. - - Args: - job_id: Job identifier. - status: Job status. - payload: Serializable managed job payload. - """ - - with closing(self._connect()) as connection: - connection.execute( - """ - INSERT INTO jobs(job_id, status, payload_json, updated_at) - VALUES(?, ?, ?, ?) - ON CONFLICT(job_id) DO UPDATE SET - status = excluded.status, - payload_json = excluded.payload_json, - updated_at = excluded.updated_at - """, - (job_id, status, json.dumps(payload, indent=2), utc_now_iso()), - ) - connection.commit() - - def load_jobs(self) -> list[dict[str, Any]]: - """Load all persisted job payloads. - - Returns: - Serialized managed job payloads. - """ - - with closing(self._connect()) as connection: - rows = connection.execute("SELECT payload_json FROM jobs ORDER BY updated_at").fetchall() - return [json.loads(row["payload_json"]) for row in rows] - - def save_worker(self, worker_id: str, status: str, payload: dict[str, Any]) -> None: - """Save or replace a worker record. - - Args: - worker_id: Worker identifier. - status: Worker status. - payload: Serializable worker payload. - """ - - with closing(self._connect()) as connection: - connection.execute( - """ - INSERT INTO workers(worker_id, status, payload_json, updated_at) - VALUES(?, ?, ?, ?) - ON CONFLICT(worker_id) DO UPDATE SET - status = excluded.status, - payload_json = excluded.payload_json, - updated_at = excluded.updated_at - """, - (worker_id, status, json.dumps(payload, indent=2), utc_now_iso()), - ) - connection.commit() - - def load_workers(self) -> list[dict[str, Any]]: - """Load all persisted worker payloads. - - Returns: - Serialized worker payloads. - """ - - with closing(self._connect()) as connection: - rows = connection.execute("SELECT payload_json FROM workers ORDER BY worker_id").fetchall() - return [json.loads(row["payload_json"]) for row in rows] - - def record_heartbeat(self, worker_id: str, payload: dict[str, Any]) -> None: - """Record a worker heartbeat. - - Args: - worker_id: Worker identifier. - payload: Serializable heartbeat payload. - """ - - with closing(self._connect()) as connection: - connection.execute( - "INSERT INTO worker_heartbeats(worker_id, payload_json, received_at) VALUES(?, ?, ?)", - (worker_id, json.dumps(payload, indent=2), utc_now_iso()), - ) - connection.commit() - - def latest_heartbeats(self) -> dict[str, dict[str, Any]]: - """Return the latest heartbeat for each worker. - - Returns: - Mapping of worker ID to heartbeat payload. - """ - - with closing(self._connect()) as connection: - rows = connection.execute( - """ - SELECT h.worker_id, h.payload_json - FROM worker_heartbeats h - JOIN ( - SELECT worker_id, MAX(received_at) AS received_at - FROM worker_heartbeats - GROUP BY worker_id - ) latest - ON h.worker_id = latest.worker_id AND h.received_at = latest.received_at - """ - ).fetchall() - return {row["worker_id"]: json.loads(row["payload_json"]) for row in rows} - - def _initialize(self) -> None: - """Create database tables when they do not exist.""" - - with closing(self._connect()) as connection: - connection.execute( - """ - CREATE TABLE IF NOT EXISTS jobs( - job_id TEXT PRIMARY KEY, - status TEXT NOT NULL, - payload_json TEXT NOT NULL, - updated_at TEXT NOT NULL - ) - """ - ) - connection.execute( - """ - CREATE TABLE IF NOT EXISTS workers( - worker_id TEXT PRIMARY KEY, - status TEXT NOT NULL, - payload_json TEXT NOT NULL, - updated_at TEXT NOT NULL - ) - """ - ) - connection.execute( - """ - CREATE TABLE IF NOT EXISTS worker_heartbeats( - id INTEGER PRIMARY KEY AUTOINCREMENT, - worker_id TEXT NOT NULL, - payload_json TEXT NOT NULL, - received_at TEXT NOT NULL - ) - """ - ) - connection.execute("CREATE INDEX IF NOT EXISTS idx_jobs_status ON jobs(status)") - connection.execute("CREATE INDEX IF NOT EXISTS idx_heartbeats_worker ON worker_heartbeats(worker_id)") - connection.commit() - - def _connect(self) -> sqlite3.Connection: - """Open a SQLite connection. - - Returns: - SQLite connection. - """ - - connection = sqlite3.connect(self.db_path) - connection.row_factory = sqlite3.Row - return connection diff --git a/llm_trainer/data.py b/llm_trainer/data.py deleted file mode 100644 index f97d1ea..0000000 --- a/llm_trainer/data.py +++ /dev/null @@ -1,838 +0,0 @@ -from __future__ import annotations - -import json -import re -import hashlib -from concurrent.futures import ThreadPoolExecutor, as_completed -from dataclasses import dataclass -from pathlib import Path -from typing import Any, Callable, Optional - -import PyPDF2 - - -class OperationCancelled(RuntimeError): - """Raised when a long-running operation is cancelled by the user.""" - - -SUPPORTED_TEXT_SUFFIXES = {".txt", ".md", ".text"} -SUPPORTED_CODE_SUFFIXES = { - ".py": "python", - ".js": "javascript", - ".ts": "typescript", - ".tsx": "typescript", - ".jsx": "javascript", - ".java": "java", - ".c": "c", - ".h": "c", - ".cpp": "cpp", - ".cc": "cpp", - ".hpp": "cpp", - ".cs": "csharp", - ".go": "go", - ".rs": "rust", - ".php": "php", - ".rb": "ruby", - ".swift": "swift", - ".kt": "kotlin", - ".kts": "kotlin", - ".scala": "scala", - ".r": "r", - ".sql": "sql", - ".sh": "bash", - ".ps1": "powershell", - ".html": "html", - ".css": "css", - ".xml": "xml", - ".json": "json", - ".yaml": "yaml", - ".yml": "yaml", - ".toml": "toml", - ".ini": "ini", -} - - -@dataclass -class Document: - """Loaded training sample. - - Attributes: - path: Original source path. - text: Loaded or extracted sample text. - kind: Sample type, usually ``prose`` or ``code``. - language: Optional programming language label for code samples. - """ - - path: Path - text: str - kind: str = "prose" - language: Optional[str] = None - - -def document_to_dict(document: Document) -> dict[str, Any]: - """Convert a document to a JSON-friendly dictionary. - - Args: - document: Document to serialize. - - Returns: - JSON-friendly document dictionary. - """ - - return { - "path": str(document.path), - "text": document.text, - "kind": document.kind, - "language": document.language, - } - - -def document_from_dict(value: dict[str, Any]) -> Document: - """Load a document from a dictionary. - - Args: - value: Serialized document. - - Returns: - Document instance. - """ - - return Document( - path=Path(value["path"]), - text=str(value.get("text", "")), - kind=str(value.get("kind", "prose")), - language=value.get("language"), - ) - - -def file_sha256(path: Path) -> str: - """Calculate a file SHA-256 digest. - - Args: - path: File path. - - Returns: - Hex digest. - """ - - digest = hashlib.sha256() - with path.open("rb") as file: - for chunk in iter(lambda: file.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - -def file_fingerprint(path: Path, fast: bool = False, sample_bytes: int = 64 * 1024) -> str: - """Calculate a file fingerprint. - - Args: - path: File path. - fast: When true, hash only sampled bytes and size metadata. - sample_bytes: Bytes read from file head/tail in fast mode. - - Returns: - Fingerprint hex digest. - """ - - if not fast: - return file_sha256(path) - sample_bytes = max(0, int(sample_bytes)) - - stat = path.stat() - size = stat.st_size - digest = hashlib.blake2b(digest_size=20) - digest.update(str(size).encode("utf-8")) - if size <= 0: - return f"fast:{digest.hexdigest()}" - with path.open("rb") as file: - head = file.read(sample_bytes) - digest.update(head) - if size > sample_bytes: - file.seek(max(0, size - sample_bytes)) - digest.update(file.read(sample_bytes)) - return f"fast:{digest.hexdigest()}" - - -def supported_source_paths(input_dir: Path, code_training_mode: bool = False, include_source_code: bool = True) -> list[Path]: - """Return supported source paths. - - Args: - input_dir: Folder to scan. - code_training_mode: Whether source-code files are supported. - include_source_code: Whether to include source-code files. - - Returns: - Sorted supported paths. - - Raises: - FileNotFoundError: If the folder does not exist. - """ - - input_dir = Path(input_dir) - if not input_dir.exists(): - raise FileNotFoundError(f"Input folder does not exist: {input_dir}") - paths = [path for path in sorted(input_dir.rglob("*")) if path.is_file()] - return [ - path - for path in paths - if path.suffix.lower() in SUPPORTED_TEXT_SUFFIXES | {".pdf", ".jsonl"} - or (code_training_mode and include_source_code and path.suffix.lower() in SUPPORTED_CODE_SUFFIXES) - ] - - -def clean_text(text: str, lowercase: bool = False) -> str: - """Normalize prose text. - - Args: - text: Raw text extracted from a document. - lowercase: Whether to convert text to lowercase. - - Returns: - Whitespace-normalized prose text. - """ - - text = text.replace("\x00", " ") - text = re.sub(r"\s+", " ", text) - text = text.strip() - return text.lower() if lowercase else text - - -def clean_code(text: str, lowercase: bool = False) -> str: - """Normalize code while preserving structure. - - Args: - text: Raw code text. - lowercase: Whether to lowercase code. Usually false for code. - - Returns: - Code text with line breaks and indentation retained. - """ - - text = text.replace("\x00", "") - text = text.replace("\r\n", "\n").replace("\r", "\n") - text = re.sub(r"\n{4,}", "\n\n\n", text) - text = text.strip() - return text.lower() if lowercase else text - - -def read_pdf(path: Path) -> str: - """Extract text from a PDF file. - - Args: - path: PDF file path. - - Returns: - Extracted text joined across pages. - """ - - chunks: list[str] = [] - with path.open("rb") as file: - reader = PyPDF2.PdfReader(file) - for page in reader.pages: - chunks.append(page.extract_text() or "") - return "\n".join(chunks) - - -def read_jsonl(path: Path) -> str: - """Read text-like values from a JSONL file. - - Args: - path: JSONL file path. - - Returns: - Combined text from string rows or common text fields. - """ - - chunks: list[str] = [] - with path.open("r", encoding="utf-8") as file: - for line in file: - if not line.strip(): - continue - value = json.loads(line) - if isinstance(value, str): - chunks.append(value) - elif isinstance(value, dict): - for key in ("text", "content", "prompt", "completion"): - if key in value and value[key]: - chunks.append(str(value[key])) - return "\n".join(chunks) - - -def _iter_json_records(path: Path) -> list[Any]: - """Read JSON or JSONL records from a file. - - Args: - path: JSON or JSONL source file. - - Returns: - List of decoded records. - """ - - if path.suffix.lower() == ".jsonl": - records: list[Any] = [] - with path.open("r", encoding="utf-8") as file: - for line in file: - if line.strip(): - records.append(json.loads(line)) - return records - - value = json.loads(path.read_text(encoding="utf-8")) - if isinstance(value, list): - return value - if isinstance(value, dict): - for key in ("data", "examples", "items", "records", "rows"): - nested = value.get(key) - if isinstance(nested, list): - return nested - return [value] - return [value] - - -def _role_name(value: Any) -> str: - """Return a readable chat role name. - - Args: - value: Raw role/from value. - - Returns: - Normalized role label. - """ - - role = str(value or "").strip().lower() - if role in {"human", "user", "prompt", "question"}: - return "User" - if role in {"gpt", "assistant", "bot", "model", "answer"}: - return "Assistant" - if role in {"system", "developer"}: - return role.title() - return role.title() if role else "Message" - - -def _format_message_list(messages: Any) -> str: - """Format OpenAI/ShareGPT-style message rows. - - Args: - messages: Message list from a structured dataset record. - - Returns: - Human-readable transcript text. - """ - - if not isinstance(messages, list): - return "" - lines: list[str] = [] - for item in messages: - if isinstance(item, str): - content = item.strip() - if content: - lines.append(content) - continue - if not isinstance(item, dict): - continue - role = _role_name(item.get("role", item.get("from", item.get("speaker", item.get("author"))))) - content = item.get("content", item.get("value", item.get("text", item.get("message", "")))) - if isinstance(content, list): - content = " ".join(str(part) for part in content if part) - content_text = str(content or "").strip() - if content_text: - lines.append(f"{role}: {content_text}") - return "\n".join(lines) - - -def _extract_structured_text(record: Any, kind: str) -> str: - """Extract training text from a structured JSON record. - - Args: - record: JSON value from a dataset file. - kind: Target sample kind, usually conversation or instruction. - - Returns: - Extracted sample text, or an empty string. - """ - - if isinstance(record, str): - return record.strip() - if isinstance(record, list): - return _format_message_list(record) - if not isinstance(record, dict): - return "" - - for message_key in ("messages", "conversations", "dialogue", "utterances", "turns"): - transcript = _format_message_list(record.get(message_key)) - if transcript: - return transcript - - instruction = str(record.get("instruction", "") or "").strip() - user_input = str(record.get("input", "") or "").strip() - output = str( - record.get("output", record.get("response", record.get("answer", record.get("completion", "")))) or "" - ).strip() - if instruction or user_input: - lines = [] - if instruction: - lines.append(f"Instruction: {instruction}") - if user_input: - lines.append(f"Input: {user_input}") - if output: - lines.append(f"Response: {output}") - return "\n".join(lines) - - prompt = str(record.get("prompt", record.get("question", "")) or "").strip() - completion = str(record.get("completion", record.get("answer", record.get("response", ""))) or "").strip() - if prompt or completion: - if kind == "conversation": - return "\n".join(part for part in (f"User: {prompt}" if prompt else "", f"Assistant: {completion}" if completion else "") if part) - return "\n".join(part for part in (f"Prompt: {prompt}" if prompt else "", f"Completion: {completion}" if completion else "") if part) - - for key in ("text", "content", "body"): - value = record.get(key) - if value: - return str(value).strip() - return "" - - -def load_structured_json_documents(path: Path, kind: str, lowercase: bool = False) -> list[Document]: - """Load conversation or instruction samples from JSON/JSONL files. - - Args: - path: JSON/JSONL file or folder containing JSON/JSONL files. - kind: Sample kind to assign, usually ``conversation`` or ``instruction``. - lowercase: Whether to lowercase extracted text. - - Returns: - Loaded structured dataset documents. - - Raises: - FileNotFoundError: If the configured path does not exist. - ValueError: If the file type is unsupported. - """ - - path = Path(path) - if not path.exists(): - raise FileNotFoundError(f"Structured dataset path does not exist: {path}") - files = [path] if path.is_file() else sorted(item for item in path.rglob("*") if item.suffix.lower() in {".json", ".jsonl"}) - if not files: - raise ValueError(f"No .json or .jsonl files found in {path}") - - documents: list[Document] = [] - for file_path in files: - if file_path.suffix.lower() not in {".json", ".jsonl"}: - raise ValueError(f"Unsupported structured dataset file: {file_path}") - for index, record in enumerate(_iter_json_records(file_path), start=1): - text = _extract_structured_text(record, kind) - text = clean_code(text, lowercase=lowercase) - if not text: - continue - documents.append( - Document( - path=Path(f"{file_path}#{index}"), - text=text, - kind=kind, - language="local_json", - ) - ) - return documents - - -def read_supported_document( - path: Path, - lowercase: bool = False, - code_training_mode: bool = False, - preserve_indentation: bool = True, -) -> Optional[Document]: - """Read one supported document or source-code file. - - Args: - path: Source file path. - lowercase: Whether to lowercase loaded content. - code_training_mode: Whether code-specific handling is enabled. - preserve_indentation: Whether code line structure should be kept. - - Returns: - Loaded document, or ``None`` when the file has no useful text. - """ - - suffix = path.suffix.lower() - # Bundled code-training corpora may use .txt or .jsonl containers while - # still being intended for code-aware preparation. Classify those files - # by their directory as well as by source-code extension. - in_code_training_folder = any( - part.lower() == "code_training" for part in path.parts - ) - if code_training_mode and suffix in SUPPORTED_CODE_SUFFIXES: - text = path.read_text(encoding="utf-8", errors="ignore") - text = clean_code(text, lowercase=lowercase) if preserve_indentation else clean_text(text, lowercase=lowercase) - if not text: - return None - return Document(path=path, text=text, kind="code", language=SUPPORTED_CODE_SUFFIXES[suffix]) - if suffix in SUPPORTED_TEXT_SUFFIXES: - text = path.read_text(encoding="utf-8", errors="ignore") - elif suffix == ".pdf": - text = read_pdf(path) - elif suffix == ".jsonl": - text = read_jsonl(path) - else: - return None - - if in_code_training_folder and code_training_mode: - text = clean_code(text, lowercase=lowercase) - if not text: - return None - language = next( - ( - language - for extension, language in SUPPORTED_CODE_SUFFIXES.items() - if path.stem.lower().startswith(extension.lstrip(".")) - ), - None, - ) - return Document(path=path, text=text, kind="code", language=language) - - text = clean_text(text, lowercase=lowercase) - if not text: - return None - return Document(path=path, text=text) - - -def is_code_like_line(line: str) -> bool: - """Estimate whether a line appears to be source code. - - Args: - line: Candidate text line. - - Returns: - True when the line contains common code markers or dense syntax. - """ - - stripped = line.strip() - if not stripped: - return False - code_markers = ( - "def ", "class ", "function ", "import ", "from ", "return ", "for ", - "while ", "if ", "else:", "elif ", "try:", "except ", "public ", - "private ", "protected ", "#include", "using ", "namespace ", "var ", - "let ", "const ", "SELECT ", "INSERT ", "UPDATE ", "DELETE ", - ) - if stripped.startswith(code_markers): - return True - symbol_count = sum(stripped.count(symbol) for symbol in "{}[]();=<>:+-*/") - return symbol_count >= 3 or line.startswith((" ", "\t")) - - -def guess_language(text: str, fallback: Optional[str] = None) -> Optional[str]: - """Guess a programming language from code text. - - Args: - text: Code sample text. - fallback: Language to return when no heuristic matches. - - Returns: - Guessed language name, fallback, or ``None``. - """ - - lowered = text.lower() - if "def " in lowered or "import " in lowered or "self." in lowered: - return "python" - if "function " in lowered or "const " in lowered or "let " in lowered or "=>" in lowered: - return "javascript" - if "public class" in lowered or "system.out" in lowered: - return "java" - if "#include" in lowered or "std::" in lowered: - return "cpp" - if "select " in lowered and " from " in lowered: - return "sql" - return fallback - - -def extract_code_blocks_from_text(document: Document, preserve_indentation: bool = True) -> list[Document]: - """Extract code-like blocks from prose/PDF text. - - Args: - document: Source document whose text may contain code snippets. - preserve_indentation: Whether extracted code should keep indentation. - - Returns: - Code sample documents extracted from the source document. - """ - - lines = document.text.replace("\r\n", "\n").replace("\r", "\n").split("\n") - blocks: list[Document] = [] - current: list[str] = [] - - def flush() -> None: - """Flush the current candidate block into ``blocks`` if code-like.""" - - nonlocal current - if len(current) >= 3: - block = "\n".join(current) - if sum(1 for line in current if is_code_like_line(line)) >= 2: - cleaned = clean_code(block) if preserve_indentation else clean_text(block) - blocks.append( - Document( - path=document.path, - text=cleaned, - kind="code", - language=guess_language(cleaned), - ) - ) - current = [] - - for line in lines: - if is_code_like_line(line): - current.append(line) - else: - flush() - flush() - return blocks - - -def expand_code_documents( - documents: list[Document], - include_prose: bool = True, - extract_code_blocks: bool = True, - preserve_indentation: bool = True, - should_stop: Optional[Callable[[], bool]] = None, -) -> list[Document]: - """Expand documents for code-aware training. - - Args: - documents: Loaded source documents. - include_prose: Whether to keep prose documents. - extract_code_blocks: Whether to extract code-like prose blocks. - preserve_indentation: Whether to preserve code indentation. - should_stop: Optional cancellation callback. - - Returns: - Expanded document list. - """ - - expanded: list[Document] = [] - for document in documents: - if should_stop and should_stop(): - raise OperationCancelled("Dataset preparation stopped by user.") - if document.kind == "code": - expanded.append(document) - continue - if include_prose: - expanded.append(document) - if extract_code_blocks: - expanded.extend(extract_code_blocks_from_text(document, preserve_indentation=preserve_indentation)) - return expanded - - -def format_document_for_training( - document: Document, - generate_instruction_samples: bool = True, - reasoning_sample_mode: str = "scaffold", -) -> str: - """Format a document with tags for the training corpus. - - Args: - document: Document to serialize. - generate_instruction_samples: Whether code samples should include a - simple instruction wrapper. - reasoning_sample_mode: Instruction/reasoning style: none, scaffold, or detailed. - - Returns: - Tagged training text for the document. - """ - - source = document.path.name - if document.kind == "code": - language = document.language or "unknown" - if generate_instruction_samples: - return format_code_instruction_sample(document, language, source, reasoning_sample_mode) - return f"\n{document.text}\n" - if document.kind == "conversation": - return f"\n{document.text}\n" - if document.kind == "instruction": - return f"\n{document.text}\n" - return f"\n{document.text}\n" - - -def format_code_instruction_sample(document: Document, language: str, source: str, reasoning_sample_mode: str) -> str: - """Format a code sample as an instruction/reasoning training example. - - Args: - document: Code document. - language: Programming language label. - source: Source file name. - reasoning_sample_mode: Instruction/reasoning style. - - Returns: - Tagged training text. - """ - - task = infer_code_task(document, language) - if reasoning_sample_mode == "none": - return ( - f"\n" - f"{task}\n" - f"\n```{language}\n{document.text}\n```\n\n" - f"" - ) - if reasoning_sample_mode == "detailed": - reasoning = ( - "1. Identify the goal implied by the file name, function names, and surrounding code.\n" - "2. Inspect inputs, outputs, control flow, data structures, and error handling.\n" - "3. Preserve language syntax, indentation, imports, and naming style.\n" - "4. Produce the code first, then explain the important design choices and edge cases." - ) - explanation = ( - "This sample teaches the model to connect a programming task with implementation details, " - "syntax, structure, and a concise explanation." - ) - else: - reasoning = ( - "Understand the requested programming task, choose the relevant language patterns, " - "preserve correct syntax, and provide the implementation." - ) - explanation = "The answer contains the implementation that satisfies the task." - return ( - f"\n" - f"{task}\n" - f"\n{reasoning}\n\n" - f"\n```{language}\n{document.text}\n```\n\n" - f"{explanation}\n" - f"" - ) - - -def infer_code_task(document: Document, language: str) -> str: - """Infer a simple task instruction for a code sample. - - Args: - document: Code document. - language: Programming language label. - - Returns: - Task instruction text. - """ - - stem = document.path.stem.replace("_", " ").replace("-", " ").strip() - if stem and stem.lower() not in {"index", "main", "app"}: - return f"Write or explain the {language} code for {stem}." - return f"Write or explain this {language} code with correct syntax and structure." - - -def load_documents( - input_dir: Path, - lowercase: bool = False, - max_workers: int = 4, - code_training_mode: bool = False, - include_prose: bool = True, - include_source_code: bool = True, - extract_code_blocks: bool = True, - preserve_indentation: bool = True, - progress: Optional[Callable[[Any], None]] = None, - should_stop: Optional[Callable[[], bool]] = None, -) -> list[Document]: - """Load supported files from a folder. - - Args: - input_dir: Folder to scan recursively. - lowercase: Whether to lowercase loaded content. - max_workers: Maximum parallel file readers. - code_training_mode: Enables code-aware loading and expansion. - include_prose: Keeps prose documents in code-aware mode. - include_source_code: Includes source-code files in code-aware mode. - extract_code_blocks: Extracts code-like blocks from prose documents. - preserve_indentation: Keeps code formatting where possible. - progress: Optional callback receiving progress event dictionaries. - should_stop: Optional callback returning true when loading should stop. - - Returns: - Sorted list of loaded document samples. - - Raises: - FileNotFoundError: If ``input_dir`` does not exist. - """ - - input_dir = Path(input_dir) - if not input_dir.exists(): - raise FileNotFoundError(f"Input folder does not exist: {input_dir}") - - documents: list[Document] = [] - supported_paths = supported_source_paths(input_dir, code_training_mode=code_training_mode, include_source_code=include_source_code) - if progress: - progress({"message": f"Found {len(supported_paths)} supported files in {input_dir}.", "percent": 8}) - - if not supported_paths: - return documents - - worker_count = max(1, min(max_workers, len(supported_paths))) - if progress: - progress({"message": f"Reading files with {worker_count} worker(s).", "percent": 10}) - - with ThreadPoolExecutor(max_workers=worker_count) as executor: - future_map = { - executor.submit(read_supported_document, path, lowercase, code_training_mode, preserve_indentation): path - for path in supported_paths - } - for index, future in enumerate(as_completed(future_map), start=1): - if should_stop and should_stop(): - for pending in future_map: - pending.cancel() - raise OperationCancelled("Dataset preparation stopped by user.") - path = future_map[future] - percent = 10 + int(32 * index / max(len(supported_paths), 1)) - try: - document = future.result() - except Exception as exc: - if progress: - progress({"message": f"Failed {path.name}: {exc}", "percent": percent}) - continue - - if document is None: - if progress: - progress({"message": f"Skipped {path.name}: no readable text found.", "percent": percent}) - continue - - documents.append(document) - if progress: - progress({"message": f"Loaded {path.name}: {len(document.text):,} characters.", "percent": percent}) - - if code_training_mode: - documents = expand_code_documents( - documents, - include_prose=include_prose, - extract_code_blocks=extract_code_blocks, - preserve_indentation=preserve_indentation, - should_stop=should_stop, - ) - - return sorted(documents, key=lambda document: (str(document.path), document.kind, document.language or "")) - - -def write_training_corpus( - documents: list[Document], - output_path: Path, - code_training_mode: bool = False, - generate_instruction_samples: bool = True, - reasoning_sample_mode: str = "scaffold", -) -> None: - """Write loaded samples into a tokenizer training corpus. - - Args: - documents: Loaded document samples. - output_path: Destination corpus text file. - code_training_mode: Whether to use code/prose tags. - generate_instruction_samples: Whether to wrap code samples with - instruction text. - reasoning_sample_mode: Instruction/reasoning style for code samples. - """ - - output_path.parent.mkdir(parents=True, exist_ok=True) - with output_path.open("w", encoding="utf-8") as file: - for doc in documents: - if code_training_mode: - file.write( - format_document_for_training( - doc, - generate_instruction_samples=generate_instruction_samples, - reasoning_sample_mode=reasoning_sample_mode, - ) - ) - else: - file.write(doc.text) - file.write("\n\n") diff --git a/llm_trainer/dataset_build.py b/llm_trainer/dataset_build.py deleted file mode 100644 index ff00de7..0000000 --- a/llm_trainer/dataset_build.py +++ /dev/null @@ -1,1475 +0,0 @@ -from __future__ import annotations - -import json -import hashlib -import logging -import multiprocessing as mp -import os -import re -import shutil -import statistics -from collections import Counter -from concurrent.futures import ProcessPoolExecutor, as_completed -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any, Callable, Optional - -import numpy as np - -from .config import DatasetConfig, dataclass_to_jsonable -from .conversation_datasets import CONVERSATION_DATASET_PRESETS, \ - dataset_ids_for_stage, load_conversation_documents -from .data import ( - Document, - SUPPORTED_CODE_SUFFIXES, - SUPPORTED_TEXT_SUFFIXES, - document_from_dict, - file_fingerprint, - file_sha256, - format_document_for_training, - load_structured_json_documents, - supported_source_paths, -) -from .dataset_mixture import ( - MAX_REPETITIVE_UNIT_RATIO, - MIN_REPETITION_CHECK_CHARS, - MIN_REPETITION_CHECK_UNITS, - _content_units_for_diversity, -) -from .document_extraction import ( - bad_extraction_reasons as _bad_extraction_reasons, - extract_documents_worker, -) -from .lineage import read_json, record_dataset_version, write_json -from .manifest_store import ManifestStore -from .tokenizer import ( - encode_file_to_npy, - load_tokenizer, - save_tokenizer_package, - token_dtype_for_vocab, - train_tokenizer, - validate_training_tokenizer, -) -from .training import split_tokens_to_files - -LOGGER = logging.getLogger(__name__) - - -def _local_structured_dataset_paths(config: DatasetConfig) -> list[ - tuple[Path, str, str]]: - """Return configured local structured dataset paths. - - Args: - config: Dataset configuration. - - Returns: - Tuples of path, document kind, and progress label. - """ - - items: list[tuple[Path, str, str]] = [] - seen: set[tuple[str, str]] = set() - for path in [config.conversation_dataset_path, - *config.conversation_dataset_paths]: - if path is None or not str(path).strip(): - continue - key = ("conversation", str(Path(path))) - if key not in seen: - seen.add(key) - items.append((Path(path), "conversation", "local conversation")) - for path in [config.instruction_dataset_path, - *config.instruction_dataset_paths]: - if path is None or not str(path).strip(): - continue - key = ("instruction", str(Path(path))) - if key not in seen: - seen.add(key) - items.append((Path(path), "instruction", "local instruction")) - return items - - -@dataclass -class DatasetBuildResult: - """Result returned after dataset preparation. - - Attributes: - output_dir: Prepared dataset folder. - tokenizer_path: Path to tokenizer JSON. - document_count: Number of loaded samples. - token_count: Total encoded tokens. - train_window_count: Number of sliding training windows. - val_window_count: Number of sliding validation windows. - sequence_token_stats: Approximate min/avg/median/max source token lengths. - vocab_size: Final tokenizer vocabulary size. - character_count: Total corpus characters. - suggested_vocab_size: Automatically estimated vocabulary size. - warning: Optional dataset quality warning. - code_sample_count: Number of code samples. - prose_sample_count: Number of prose samples. - conversation_sample_count: Number of conversation/instruction samples. - cached_file_count: Number of unchanged source files reused from cache. - processed_file_count: Number of source files extracted this run. - skipped_file_count: Number of files with no readable text. - failed_file_count: Number of files that failed extraction. - dataset_version_id: Unique dataset version identifier. - dataset_version_number: One-based dataset version number. - mixture_report: Per-source family sampling report. - quality_score: Dataset quality score from 0 to 100. - quality_stars: Dataset quality rating from 0 to 5. - quality_label: Human-readable dataset quality label. - quality_reasons: Short reasons behind the quality score. - duplicate_block_count: Number of repeated blocks in the written corpus. - unique_block_count: Number of unique blocks in the written corpus. - corpus_block_count: Number of non-empty blocks inspected in the written corpus. - duplicate_block_ratio: Fraction of repeated text blocks in the written corpus. - unique_block_ratio: Fraction of unique text blocks in the written corpus. - """ - - output_dir: Path - tokenizer_path: Path - document_count: int - token_count: int - vocab_size: int - character_count: int - suggested_vocab_size: int - train_window_count: int = 0 - val_window_count: int = 0 - sequence_token_stats: dict[str, float] = field(default_factory=dict) - warning: Optional[str] = None - code_sample_count: int = 0 - prose_sample_count: int = 0 - conversation_sample_count: int = 0 - cached_file_count: int = 0 - processed_file_count: int = 0 - skipped_file_count: int = 0 - failed_file_count: int = 0 - dataset_version_id: str = "" - dataset_version_number: int = 0 - mixture_report: dict[str, Any] = field(default_factory=dict) - quality_score: float = 0.0 - quality_stars: float = 0.0 - quality_label: str = "Not rated" - quality_reasons: list[str] = field(default_factory=list) - duplicate_block_count: int = 0 - unique_block_count: int = 0 - corpus_block_count: int = 0 - duplicate_block_ratio: float = 0.0 - unique_block_ratio: float = 1.0 - - -def _emit(progress: Optional[Callable[[Any], None]], message: str, - percent: Optional[int] = None) -> None: - """Emit a progress event if a callback is available. - - Args: - progress: Optional callback for progress dictionaries. - message: Human-readable progress message. - percent: Optional progress percentage. - """ - - LOGGER.info(message) - if progress: - progress({"message": message, "percent": percent}) - - -def estimate_vocab_size(character_count: int, unique_word_count: int) -> int: - """Estimate a reasonable tokenizer vocabulary size. - - Args: - character_count: Number of corpus characters. - unique_word_count: Approximate number of unique whitespace words. - - Returns: - Suggested vocabulary size. - """ - - if character_count < 20_000: - ceiling = 1_000 - elif character_count < 100_000: - ceiling = 4_000 - elif character_count < 500_000: - ceiling = 8_000 - elif character_count < 2_000_000: - ceiling = 16_000 - else: - ceiling = 32_000 - - desired = max(512, int(unique_word_count * 1.7), int(character_count / 45)) - return max(256, min(ceiling, desired)) - - -def content_warning(character_count: int) -> Optional[str]: - """Return a corpus-size warning when the dataset is small. - - Args: - character_count: Number of corpus characters. - - Returns: - Warning text, or ``None`` when the corpus is large enough. - """ - - if character_count < 10_000: - return "The corpus is very small. Training can run, but the model will only be useful for smoke tests." - if character_count < 100_000: - return "The corpus is modest. Use more text for better generations and reasoning behavior." - return None - - -def _resolve_tokenizer_strategy(config: DatasetConfig, tokenizer_path: Path) -> \ -tuple[str, bool]: - """Resolve tokenizer strategy into an executable mode. - - Args: - config: Dataset configuration. - tokenizer_path: Dataset tokenizer output path. - - Returns: - Strategy name and whether the dataset tokenizer should be reused. - """ - - strategy = config.tokenizer_strategy or "auto" - if strategy == "auto": - return strategy, config.prepare_mode == "incremental" and tokenizer_path.exists() - if strategy == "reuse_dataset": - if not tokenizer_path.exists(): - raise FileNotFoundError( - f"Cannot reuse dataset tokenizer because tokenizer.json was not found in {config.output_dir}." - ) - return strategy, True - if strategy in {"train_new", "import_tokenizer"}: - return strategy, False - raise ValueError(f"Unsupported tokenizer strategy: {strategy}") - - -def _load_or_create_tokenizer( - config: DatasetConfig, - corpus_path: Path, - tokenizer_path: Path, - selected_vocab_size: int, - progress: Optional[Callable[[Any], None]], - should_stop: Optional[Callable[[], bool]], -) -> tuple[Any, bool, bool, Optional[str]]: - """Load, import, or train a tokenizer for the prepared corpus. - - Args: - config: Dataset configuration. - corpus_path: Normalized training corpus path. - tokenizer_path: Dataset tokenizer output path. - selected_vocab_size: Vocabulary size used when training a new tokenizer. - progress: Optional progress callback. - should_stop: Optional cancellation callback. - - Returns: - Tokenizer, reused flag, imported flag, and optional source path. - """ - - strategy, reuse_tokenizer = _resolve_tokenizer_strategy(config, - tokenizer_path) - imported = False - source_path: Optional[str] = None - - if reuse_tokenizer: - _emit(progress, "Reusing existing dataset tokenizer.json...", 62) - return load_tokenizer(tokenizer_path), True, imported, source_path - - if strategy == "import_tokenizer": - if config.tokenizer_path is None: - raise ValueError( - "Choose a tokenizer.json file when tokenizer strategy is Import tokenizer.json.") - import_path = Path(config.tokenizer_path) - if not import_path.exists(): - raise FileNotFoundError( - f"Tokenizer import file not found: {import_path}") - _emit(progress, f"Importing tokenizer from {import_path}...", 62) - tokenizer_path.parent.mkdir(parents=True, exist_ok=True) - if import_path.resolve() != tokenizer_path.resolve(): - shutil.copy2(import_path, tokenizer_path) - return load_tokenizer(tokenizer_path), False, True, str(import_path) - - corpus_size_bytes = corpus_path.stat().st_size - training_mb = corpus_size_bytes / (1024 * 1024) - max_training_bytes = ( - int(config.tokenizer_training_max_gb * 1024**3) - if config.tokenizer_training_max_gb > 0 - else None - ) - if max_training_bytes is not None and corpus_size_bytes > max_training_bytes: - _emit( - progress, - ( - f"Training tokenizer on a {config.tokenizer_training_max_gb:.1f} GiB sample of the " - f"{training_mb:.1f} MB corpus (tokenizer_training_max_gb)..." - ), - 62, - ) - else: - _emit( - progress, - f"Training tokenizer on the full {training_mb:.1f} MB corpus...", - 62, - ) - tokenizer = train_tokenizer( - corpus_path, - tokenizer_path, - vocab_size=selected_vocab_size, - min_frequency=config.min_frequency, - should_stop=should_stop, - max_training_bytes=max_training_bytes, - ) - return tokenizer, False, imported, source_path - - -def _cache_key(config: DatasetConfig) -> str: - """Return a cache key for extraction-affecting options. - - Args: - config: Dataset configuration. - - Returns: - Cache key string. - """ - - return json.dumps( - { - "lowercase": config.lowercase, - "code_training_mode": config.code_training_mode, - "include_prose": config.include_prose, - "include_source_code": config.include_source_code, - "extract_code_blocks": config.extract_code_blocks, - "preserve_indentation": config.preserve_indentation, - "generate_instruction_samples": config.generate_instruction_samples, - "reasoning_sample_mode": config.reasoning_sample_mode, - "dataset_stage": config.dataset_stage, - "conversation_datasets": config.conversation_datasets, - "conversation_sample_limit": config.conversation_sample_limit, - "conversation_dataset_path": str( - config.conversation_dataset_path or ""), - "instruction_dataset_path": str( - config.instruction_dataset_path or ""), - "conversation_dataset_paths": [str(path) for path in - config.conversation_dataset_paths], - "instruction_dataset_paths": [str(path) for path in - config.instruction_dataset_paths], - "default_data_paths": [str(path) for path in - config.default_data_paths], - }, - sort_keys=True, - ) - - -@dataclass -class _CorpusBuildStats: - """Streaming accumulator for corpus-wide statistics. - - Every field here is either a small counter, a hash, or a length-capped - example list -- never full document text. This is what keeps - :class:`_StreamingCorpusBuilder` bounded in memory regardless of how - large the source corpus is. - """ - - character_count: int = 0 - unique_words: set[str] = field(default_factory=set) - code_sample_count: int = 0 - prose_sample_count: int = 0 - conversation_sample_count: int = 0 - accepted_document_count: int = 0 - document_char_lengths: list[int] = field(default_factory=list) - source_files: list[str] = field(default_factory=list) - source_files_truncated: bool = False - exact_duplicates_removed: int = 0 - exact_duplicate_examples: list[dict[str, str]] = field(default_factory=list) - low_diversity_removed: int = 0 - low_diversity_removed_characters: int = 0 - low_diversity_examples: list[dict[str, Any]] = field(default_factory=list) - block_counts: Counter = field(default_factory=Counter) - block_examples: dict[str, str] = field(default_factory=dict) - block_total: int = 0 - block_ignored: int = 0 - - -class _StreamingCorpusBuilder: - """Filters and writes the training corpus one document at a time. - - Replaces the previous pipeline of "load every document into one list, - then run exact-dedup over the whole list, then run a repetition filter - over the whole list, then write the whole list to disk" -- each of which - held the entire prepared corpus in memory at once. Here, each document - is deduplicated, quality-checked, written to ``corpus.txt``, and then - immediately eligible for garbage collection, so at most one document's - text is resident at a time (aside from small bookkeeping state). - """ - - _EXAMPLE_CAP = 50 - _SOURCE_FILE_CAP = 1000 - _BLOCK_EXAMPLE_CAP = 8 - - def __init__( - self, - corpus_path: Path, - code_training_mode: bool, - generate_instruction_samples: bool, - reasoning_sample_mode: str, - ) -> None: - """Open the corpus file for streaming writes. - - Args: - corpus_path: Destination corpus text file. - code_training_mode: Whether to use code/prose tags. - generate_instruction_samples: Whether code samples should include - a simple instruction wrapper. - reasoning_sample_mode: Instruction/reasoning style for code - samples. - """ - - self._code_training_mode = code_training_mode - self._generate_instruction_samples = generate_instruction_samples - self._reasoning_sample_mode = reasoning_sample_mode - self._seen_digests: dict[str, str] = {} - self.stats = _CorpusBuildStats() - corpus_path.parent.mkdir(parents=True, exist_ok=True) - self._file = corpus_path.open("w", encoding="utf-8") - - def submit(self, document: Document) -> None: - """Filter, count, and write one document, then let it be freed. - - Args: - document: Candidate document to evaluate and possibly write. - """ - - canonical = _canonical_corpus_block(document.text) - if not canonical: - return - - digest = hashlib.sha256( - f"{document.kind}\n{document.language or ''}\n{canonical}".encode("utf-8") - ).hexdigest() - original_path = self._seen_digests.get(digest) - if original_path is not None: - self.stats.exact_duplicates_removed += 1 - if len(self.stats.exact_duplicate_examples) < self._EXAMPLE_CAP: - self.stats.exact_duplicate_examples.append( - {"path": str(document.path), "duplicate_of": original_path, "kind": document.kind} - ) - return - self._seen_digests[digest] = str(document.path) - - if self._is_low_diversity(document): - self.stats.low_diversity_removed += 1 - self.stats.low_diversity_removed_characters += len(document.text) - if len(self.stats.low_diversity_examples) < self._EXAMPLE_CAP: - self.stats.low_diversity_examples.append( - {"path": str(document.path), "kind": document.kind} - ) - return - - self._accept(document, canonical) - - @staticmethod - def _is_low_diversity(document: Document) -> bool: - """Return whether a document is dominated by repeated content units. - - Args: - document: Candidate document. - - Returns: - True when the document should be excluded as low-diversity. - """ - - units = _content_units_for_diversity(document) - if len(document.text) < MIN_REPETITION_CHECK_CHARS or len(units) < MIN_REPETITION_CHECK_UNITS: - return False - duplicate_ratio = 1.0 - (len(set(units)) / len(units)) - return duplicate_ratio > MAX_REPETITIVE_UNIT_RATIO - - def _accept(self, document: Document, canonical: str) -> None: - """Record stats for and write one accepted document. - - Args: - document: Accepted document. - canonical: Canonicalized text used for block-duplicate hashing. - """ - - stats = self.stats - stats.accepted_document_count += 1 - stats.character_count += len(document.text) + 1 - stats.unique_words.update(word.lower() for word in document.text.split()) - stats.document_char_lengths.append(len(document.text)) - if document.kind == "code": - stats.code_sample_count += 1 - elif document.kind in {"conversation", "instruction"}: - stats.conversation_sample_count += 1 - else: - stats.prose_sample_count += 1 - if len(stats.source_files) < self._SOURCE_FILE_CAP: - stats.source_files.append(str(document.path)) - else: - stats.source_files_truncated = True - - if len(canonical) >= 12: - block_digest = hashlib.sha256(canonical.encode("utf-8")).hexdigest() - stats.block_counts[block_digest] += 1 - stats.block_examples.setdefault(block_digest, canonical[:240]) - stats.block_total += 1 - else: - stats.block_ignored += 1 - - if self._code_training_mode: - self._file.write( - format_document_for_training( - document, - generate_instruction_samples=self._generate_instruction_samples, - reasoning_sample_mode=self._reasoning_sample_mode, - ) - ) - else: - self._file.write(document.text) - self._file.write("\n\n") - - def close(self) -> dict[str, Any]: - """Flush the corpus file and compute the final duplicate-block report. - - Returns: - Duplicate-block report dictionary, matching the shape previously - produced by scanning the fully written corpus file. - """ - - self._file.close() - stats = self.stats - unique_blocks = len(stats.block_counts) - duplicate_blocks = sum(count - 1 for count in stats.block_counts.values() if count > 1) - duplicate_ratio = duplicate_blocks / max(stats.block_total, 1) - unique_ratio = unique_blocks / max(stats.block_total, 1) - repeated = [ - {"count": count, "sample": stats.block_examples[digest]} - for digest, count in stats.block_counts.most_common(self._BLOCK_EXAMPLE_CAP) - if count > 1 - ] - return { - "block_count": stats.block_total, - "unique_block_count": unique_blocks, - "duplicate_block_count": duplicate_blocks, - "duplicate_block_ratio": duplicate_ratio, - "unique_block_ratio": unique_ratio, - "ignored_block_count": stats.block_ignored, - "truncated": False, - "most_repeated_block_count": repeated[0]["count"] if repeated else 1, - "top_repeated_blocks": repeated, - } - - -def _load_documents_with_cache( - config: DatasetConfig, - corpus_builder: "_StreamingCorpusBuilder", - progress: Optional[Callable[[Any], None]], - should_stop: Optional[Callable[[], bool]], -) -> tuple[Any, int, int, int, int]: - """Load documents using an extraction cache and stream them into the corpus. - - New (non-cached) files are extracted in parallel worker processes (see - ``extract_documents_worker``), bounding peak memory to roughly - ``config.max_workers`` files' worth of text at a time rather than the - whole corpus. Every extracted or cached document is handed to - ``corpus_builder.submit`` and then immediately released -- this function - never accumulates a list of documents itself. - - Args: - config: Dataset configuration. - corpus_builder: Streaming builder that filters, counts, and writes - each document as it arrives. - progress: Optional progress callback. - should_stop: Optional cancellation callback. - - Returns: - Manifest, cached, processed, skipped, and failed file counts. - """ - - manifest_db_path = config.output_dir / "dataset_manifest.sqlite3" - legacy_manifest_path = config.output_dir / "dataset_manifest.json" - cache_dir = config.output_dir / "cache" / "documents" - cache_dir.mkdir(parents=True, exist_ok=True) - manifest = ManifestStore.open(manifest_db_path, - legacy_json_path=legacy_manifest_path) - key = _cache_key(config) - force_reprocess = config.prepare_mode == "force_reprocess" - - local_structured_paths = _local_structured_dataset_paths(config) - selected_default_files = [ - Path(path) - for path in config.default_data_paths - if Path(path).exists() and Path(path).is_file() - ] - input_dir_resolved = config.input_dir.resolve() if config.input_dir.exists() else None - default_files_under_input = bool( - selected_default_files) and input_dir_resolved is not None and all( - input_dir_resolved in candidate.resolve().parents or candidate.resolve() == input_dir_resolved - for candidate in selected_default_files - ) - if config.input_dir.exists() and not default_files_under_input: - source_paths = supported_source_paths( - config.input_dir, - code_training_mode=config.code_training_mode, - include_source_code=config.include_source_code, - ) - elif config.conversation_datasets or local_structured_paths or config.default_data_paths: - source_paths = [] - else: - source_paths = supported_source_paths( - config.input_dir, - code_training_mode=config.code_training_mode, - include_source_code=config.include_source_code, - ) - default_paths = [] - seen_source_paths = {path.resolve() for path in source_paths if - path.exists()} - for candidate in selected_default_files: - if not candidate.exists() or not candidate.is_file(): - _emit(progress, f"Skipped bundled data file: {candidate}") - continue - suffix = candidate.suffix.lower() - if suffix not in SUPPORTED_TEXT_SUFFIXES and suffix not in SUPPORTED_CODE_SUFFIXES and suffix not in { - ".pdf", ".json", ".jsonl"}: - _emit(progress, - f"Skipped unsupported bundled data file: {candidate.name}") - continue - resolved = candidate.resolve() - if resolved in seen_source_paths: - continue - seen_source_paths.add(resolved) - default_paths.append(candidate) - if default_paths: - source_paths.extend(default_paths) - source_paths = sorted(source_paths) - _emit(progress, - f"Bundled starter data enabled: {len(default_paths)} file(s).", - 8) - _emit(progress, - f"Found {len(source_paths)} supported files in {config.input_dir}.", - 8) - cached_count = 0 - processed_count = 0 - skipped_count = 0 - failed_count = 0 - - def _submit_cached(cached_documents: list[Document]) -> None: - for document in cached_documents: - corpus_builder.submit(document) - - # First pass: separate files that can be served from the extraction - # cache (cheap disk read, done inline) from files that need real - # extraction (CPU-heavy, farmed out to worker processes below). - pending_extraction: list[Path] = [] - file_digests: dict[str, str] = {} - file_stats: dict[str, Any] = {} - for path in source_paths: - if should_stop and should_stop(): - raise RuntimeError("Dataset preparation stopped by user.") - stat = path.stat() - digest = file_fingerprint(path, fast=config.fast_scan_mode, - sample_bytes=config.fast_scan_sample_bytes) - file_digests[str(path)] = digest - file_stats[str(path)] = stat - cache_path = cache_dir / f"{digest}.json" - manifest_key = str(path.resolve()) - previous = manifest.get(manifest_key) or {} - can_use_cache = ( - not force_reprocess - and previous.get("sha256") == digest - and previous.get("cache_key") == key - and cache_path.exists() - ) - if not can_use_cache: - pending_extraction.append(path) - continue - - cached_documents = [ - document_from_dict(item) - for item in json.loads(cache_path.read_text(encoding="utf-8")) - ] - cached_extraction_reasons = [] - if path.suffix.lower() == ".pdf": - cached_text = "\n".join(document.text for document in cached_documents) - cached_extraction_reasons = _bad_extraction_reasons( - path, - { - "path": str(path), - "kind": cached_documents[0].kind if cached_documents else "prose", - "language": cached_documents[0].language if cached_documents else "", - "characters": str(len(cached_text)), - "preview": cached_text[:1200], - }, - stat.st_size, - ) - if cached_extraction_reasons: - skipped_count += 1 - reason_text = "; ".join(cached_extraction_reasons) - _emit(progress, f"Skipped cached {path.name}: suspicious PDF extraction ({reason_text}).") - manifest.upsert( - manifest_key, - { - "path": str(path), "sha256": digest, "size": stat.st_size, - "mtime_ns": stat.st_mtime_ns, "cache_key": key, - "status": "skipped_bad_extraction", "reasons": cached_extraction_reasons, - }, - commit=False, - ) - continue - _submit_cached(cached_documents) - del cached_documents - cached_count += 1 - manifest.upsert( - manifest_key, - { - "path": str(path), "sha256": digest, "size": stat.st_size, - "mtime_ns": stat.st_mtime_ns, "cache_key": key, - "cache_file": str(cache_path.relative_to(config.output_dir)), - "status": "cached", - }, - commit=False, - ) - _emit(progress, f"Reused {path.name} from cache.") - - # Second pass: extract new/changed files in parallel worker processes. - # ``max_workers`` bounds how many files' full text can be resident (one - # per in-flight worker) at any moment, regardless of total corpus size. - # Also capped by CPU count: this is CPU-bound work, and each worker - # process is a full Python interpreter, so requesting more workers than - # cores adds contention and (thanks to spawn re-importing this app's - # dependency chain per worker) startup/memory overhead without a - # throughput benefit. - cpu_cap = max(1, os.cpu_count() or 1) - worker_count = ( - max(1, min(config.max_workers, cpu_cap, len(pending_extraction))) - if pending_extraction - else 0 - ) - if pending_extraction: - _emit( - progress, - f"Extracting {len(pending_extraction):,} file(s) with {worker_count} worker process(es)...", - 10, - ) - if worker_count: - # build_dataset() itself typically already runs inside a spawned - # child process (see ui/workers.py's ProcessTaskWorker, used with - # isolate_process=True). This app loads torch/CUDA and Qt, and - # forking a process that may already have CUDA initialized is a - # known source of crashes and hangs -- the app's own worker - # deliberately uses "spawn" for exactly that reason. This pool must - # match that choice explicitly rather than rely on the platform - # default (which is "fork" on Linux). - mp_context = mp.get_context("spawn") - with ProcessPoolExecutor(max_workers=worker_count, mp_context=mp_context) as executor: - future_map = { - executor.submit( - extract_documents_worker, - path, - config.lowercase, - config.code_training_mode, - config.preserve_indentation, - config.include_prose, - config.extract_code_blocks, - ): path - for path in pending_extraction - } - completed = 0 - for future in as_completed(future_map): - if should_stop and should_stop(): - # cancel_futures drops any not-yet-started work - # immediately; already-running extractions in worker - # processes still have to finish their current file - # (there is no safe way to interrupt mid-extraction), - # matching the previous ThreadPoolExecutor's behavior. - executor.shutdown(wait=False, cancel_futures=True) - raise RuntimeError("Dataset preparation stopped by user.") - path = future_map[future] - completed += 1 - percent = 10 + int(32 * completed / max(len(pending_extraction), 1)) - digest = file_digests[str(path)] - stat = file_stats[str(path)] - manifest_key = str(path.resolve()) - cache_path = cache_dir / f"{digest}.json" - try: - result = future.result() - except Exception as exc: # noqa: BLE001 - reported to the user - failed_count += 1 - _emit(progress, f"Failed {path.name}: {exc}", percent) - manifest.upsert( - manifest_key, - { - "path": str(path), "sha256": digest, "size": stat.st_size, - "mtime_ns": stat.st_mtime_ns, "cache_key": key, - "status": "failed", "error": str(exc), - }, - commit=False, - ) - continue - - if result["error"] is not None: - failed_count += 1 - _emit(progress, f"Failed {path.name}: {result['error']}", percent) - manifest.upsert( - manifest_key, - { - "path": str(path), "sha256": digest, "size": stat.st_size, - "mtime_ns": stat.st_mtime_ns, "cache_key": key, - "status": "failed", "error": result["error"], - }, - commit=False, - ) - continue - - if result["bad_extraction_reasons"]: - skipped_count += 1 - reason_text = "; ".join(result["bad_extraction_reasons"]) - _emit(progress, f"Skipped {path.name}: suspicious PDF extraction ({reason_text}).", percent) - manifest.upsert( - manifest_key, - { - "path": str(path), "sha256": digest, "size": stat.st_size, - "mtime_ns": stat.st_mtime_ns, "cache_key": key, - "status": "skipped_bad_extraction", "reasons": result["bad_extraction_reasons"], - }, - commit=False, - ) - continue - - if not result["documents"]: - skipped_count += 1 - _emit(progress, f"Skipped {path.name}: no readable text found.", percent) - manifest.upsert( - manifest_key, - { - "path": str(path), "sha256": digest, "size": stat.st_size, - "mtime_ns": stat.st_mtime_ns, "cache_key": key, - "status": "skipped_empty", - }, - commit=False, - ) - continue - - cache_path.write_text(json.dumps(result["documents"], ensure_ascii=False), encoding="utf-8") - for item in result["documents"]: - corpus_builder.submit(document_from_dict(item)) - processed_count += 1 - _emit(progress, f"Processed {path.name}: {len(result['documents'])} sample(s).", percent) - manifest.upsert( - manifest_key, - { - "path": str(path), "sha256": digest, "size": stat.st_size, - "mtime_ns": stat.st_mtime_ns, "cache_key": key, - "cache_file": str(cache_path.relative_to(config.output_dir)), - "status": "processed", - }, - commit=False, - ) - - for local_path, kind, label in local_structured_paths: - if should_stop and should_stop(): - raise RuntimeError("Dataset preparation stopped by user.") - local_path = Path(local_path) - _emit(progress, f"Loading {label} JSON/JSONL dataset: {local_path}", - 42) - local_documents = load_structured_json_documents(local_path, kind=kind, - lowercase=config.lowercase) - for document in local_documents: - corpus_builder.submit(document) - local_document_count = len(local_documents) - del local_documents - processed_count += 1 - manifest_key = f"local-{kind}://{local_path.resolve()}" - manifest.upsert( - manifest_key, - { - "path": str(local_path), - "kind": kind, - "sample_count": local_document_count, - "cache_key": key, - "status": "processed", - }, - commit=False, - ) - _emit(progress, - f"Loaded {local_document_count:,} {kind} sample(s) from {local_path.name}.", - 43) - - if config.conversation_datasets: - allowed_dataset_ids = set(dataset_ids_for_stage(config.dataset_stage)) - skipped_stage_ids = [dataset_id for dataset_id in - config.conversation_datasets if - dataset_id not in allowed_dataset_ids] - selected_dataset_ids = [dataset_id for dataset_id in - config.conversation_datasets if - dataset_id in allowed_dataset_ids] - if skipped_stage_ids: - skipped_labels = [ - CONVERSATION_DATASET_PRESETS[item].label - for item in skipped_stage_ids - if item in CONVERSATION_DATASET_PRESETS - ] - _emit(progress, - f"Skipping dataset(s) not recommended for {config.dataset_stage}: {', '.join(skipped_labels)}.") - if not selected_dataset_ids: - _emit(progress, - f"No online datasets selected for {config.dataset_stage}; continuing with local sources only.") - config.conversation_datasets = [] - manifest.set_meta("dataset_config", dataclass_to_jsonable(config), - commit=False) - manifest.set_meta("cache_key", key, commit=False) - manifest.commit() - return ( - manifest, - cached_count, - processed_count, - skipped_count, - failed_count, - ) - hf_cache_dir = config.output_dir / "cache" / "huggingface" - labels = [ - CONVERSATION_DATASET_PRESETS[item].label - for item in selected_dataset_ids - if item in CONVERSATION_DATASET_PRESETS - ] - _emit(progress, - f"Online training datasets enabled: {', '.join(labels)}.", 8) - _emit(progress, - f"Online training datasets will be cached in: {hf_cache_dir}", 8) - hf_documents = load_conversation_documents( - selected_dataset_ids, - config.conversation_sample_limit, - hf_cache_dir, - lowercase=config.lowercase, - progress=progress, - should_stop=should_stop, - ) - for document in hf_documents: - corpus_builder.submit(document) - del hf_documents - config.conversation_datasets = selected_dataset_ids - for dataset_id in selected_dataset_ids: - preset = CONVERSATION_DATASET_PRESETS.get(dataset_id) - manifest.upsert( - f"hf://{dataset_id}", - { - "path": f"hf://{dataset_id}", - "dataset": preset.hf_path if preset else dataset_id, - "config_name": preset.config_name if preset else "", - "split": preset.split if preset else "", - "sample_limit": config.conversation_sample_limit, - "cache_key": key, - "status": "processed", - }, - commit=False, - ) - processed_count += len(selected_dataset_ids) - - manifest.set_meta("dataset_config", dataclass_to_jsonable(config), - commit=False) - manifest.set_meta("cache_key", key, commit=False) - manifest.commit() - return ( - manifest, - cached_count, - processed_count, - skipped_count, - failed_count, - ) - - -def build_dataset( - config: DatasetConfig, - progress: Optional[Callable[[Any], None]] = None, - should_stop: Optional[Callable[[], bool]] = None, -) -> DatasetBuildResult: - """Build a tokenizer-ready dataset project. - - Args: - config: Dataset preparation settings. - progress: Optional callback receiving progress event dictionaries. - should_stop: Optional callback returning true when the user requested stop. - - Returns: - Dataset build summary. - - Raises: - ValueError: If no supported documents are found. - """ - - config.output_dir.mkdir(parents=True, exist_ok=True) - _emit(progress, "Scanning source folder...", 3) - corpus_path = config.output_dir / "corpus.txt" - corpus_builder = _StreamingCorpusBuilder( - corpus_path, - code_training_mode=config.code_training_mode, - generate_instruction_samples=config.generate_instruction_samples, - reasoning_sample_mode=config.reasoning_sample_mode, - ) - # Loading, exact-duplicate removal, low-diversity filtering, and corpus - # writing all happen inside this single streaming pass -- each document - # is evaluated and written (or dropped) as it arrives, so at no point is - # the full document set held in memory at once. See - # _StreamingCorpusBuilder / _load_documents_with_cache. - ( - manifest, - cached_file_count, - processed_file_count, - skipped_file_count, - failed_file_count, - ) = _load_documents_with_cache(config, corpus_builder, progress, should_stop) - duplicate_report = corpus_builder.close() - stats = corpus_builder.stats - if should_stop and should_stop(): - raise RuntimeError("Dataset preparation stopped by user.") - if stats.accepted_document_count == 0: - corpus_path.unlink(missing_ok=True) - raise ValueError( - "No supported text, PDF, JSONL, or structured JSON documents were found.") - if stats.exact_duplicates_removed: - _emit( - progress, - f"Removed {stats.exact_duplicates_removed:,} exact duplicate extracted document(s).", - 44, - ) - if stats.low_diversity_removed: - _emit( - progress, - ( - "Excluded " - f"{stats.low_diversity_removed:,} low-diversity document(s) " - f"({stats.low_diversity_removed_characters:,} characters) " - "instead of padding the corpus with repeated templates." - ), - 45, - ) - _emit(progress, "Low-diversity files excluded:", 45) - for excluded in stats.low_diversity_examples: - _emit(progress, f" - {excluded['path']}", 45) - mixture_report = { - "applied": False, - "reason": "Dataset mixture disabled", - } - - character_count = stats.character_count - unique_words = len(stats.unique_words) - suggested_vocab_size = estimate_vocab_size(character_count, unique_words) - selected_vocab_size = config.vocab_size or suggested_vocab_size - warning = content_warning(character_count) - if character_count < 1_000_000: - low_corpus_message = ( - "Prepared corpus is below 1M characters after quality filtering. " - "Add licensed, provenance-tracked sources or select an approved online dataset; " - "the app will not pad training data with synthetic repetition." - ) - warning = f"{warning} {low_corpus_message}" if warning else low_corpus_message - code_sample_count = stats.code_sample_count - conversation_sample_count = stats.conversation_sample_count - prose_sample_count = stats.prose_sample_count - document_count = stats.accepted_document_count - _emit(progress, - f"Content size: {character_count:,} characters across {document_count} files.", - 45) - if config.code_training_mode: - _emit(progress, - f"Code mode: {code_sample_count:,} code samples, {prose_sample_count:,} prose samples.", - 46) - if conversation_sample_count: - _emit(progress, - f"Conversation data: {conversation_sample_count:,} dialogue/instruction samples.", - 46) - if cached_file_count or processed_file_count: - _emit(progress, - f"Cache: reused {cached_file_count:,} file(s), processed {processed_file_count:,} file(s).", - 47) - if skipped_file_count or failed_file_count: - _emit(progress, - f"Quality: skipped {skipped_file_count:,} empty file(s), failed {failed_file_count:,} file(s).", - 48) - _emit(progress, f"Unique word estimate: {unique_words:,}.", 48) - _emit(progress, f"Auto vocabulary size: {selected_vocab_size:,}.", 50) - if warning: - _emit(progress, f"Warning: {warning}") - - _emit(progress, "Corpus written.", 56) - if should_stop and should_stop(): - raise RuntimeError("Dataset preparation stopped by user.") - _emit( - progress, - ( - "Corpus diversity: " - f"{duplicate_report['unique_block_count']:,}/{duplicate_report['block_count']:,} unique blocks, " - f"{duplicate_report['duplicate_block_ratio'] * 100:.1f}% repeated." - ), - 74, - ) - tokenizer_path = config.output_dir / "tokenizer.json" - if should_stop and should_stop(): - raise RuntimeError("Dataset preparation stopped by user.") - tokenizer, reuse_tokenizer, tokenizer_imported, tokenizer_source_path = _load_or_create_tokenizer( - config, - corpus_path, - tokenizer_path, - selected_vocab_size, - progress, - should_stop, - ) - validate_training_tokenizer(tokenizer) - save_tokenizer_package(tokenizer, tokenizer_path, - model_max_length=config.context_length) - - _emit(progress, "Encoding corpus into token IDs...", 78) - if should_stop and should_stop(): - raise RuntimeError("Dataset preparation stopped by user.") - # Encode straight to a memmap-friendly .npy file. Encoding is streamed in - # bounded batches (never a full in-memory token list), and the resulting - # file is then opened read-only as a memmap so the split step below also - # never holds the full token stream in RAM. - token_dtype = token_dtype_for_vocab(tokenizer.get_vocab_size()) - all_tokens_path = config.output_dir / "all_tokens.npy" - token_count = encode_file_to_npy( - tokenizer, corpus_path, all_tokens_path, token_dtype, should_stop=should_stop - ) - _emit(progress, f"Encoded {token_count:,} tokens.", 86) - - token_density = (token_count / max(character_count, 1)) if character_count else 0.0 - document_token_lengths = [max(1, int(round(char_len * token_density))) - for char_len in stats.document_char_lengths if char_len] - if document_token_lengths: - sequence_stats = { - "min": min(document_token_lengths), - "average": sum(document_token_lengths) / len( - document_token_lengths), - "median": statistics.median(document_token_lengths), - "max": max(document_token_lengths), - } - else: - sequence_stats = {"min": 0, "average": 0.0, "median": 0.0, "max": 0} - _emit( - progress, - ( - "Token distribution: " - f"min {int(sequence_stats['min']):,}, " - f"avg {float(sequence_stats['average']):,.0f}, " - f"median {float(sequence_stats['median']):,.0f}, " - f"max {int(sequence_stats['max']):,}." - ), - 88, - ) - if should_stop and should_stop(): - raise RuntimeError("Dataset preparation stopped by user.") - all_tokens = np.load(all_tokens_path, mmap_mode="r") - train_token_count, val_token_count = split_tokens_to_files( - all_tokens, - config.output_dir / "train_tokens.npy", - config.output_dir / "val_tokens.npy", - config.validation_split, - dtype=token_dtype, - should_stop=should_stop, - ) - del all_tokens # release the memmap handle before deleting the backing file - all_tokens_path.unlink(missing_ok=True) - train_window_count = max(0, train_token_count - config.context_length) - val_window_count = max(0, val_token_count - config.context_length) - _emit(progress, - f"Training tokens: {train_token_count:,}; validation tokens: {val_token_count:,}.", - 92) - _emit(progress, - f"Training windows: {train_window_count:,}; validation windows: {val_window_count:,}.", - 92) - quality_report = _dataset_quality_report( - document_count=document_count, - token_count=token_count, - vocab_size=tokenizer.get_vocab_size(), - unique_words=unique_words, - train_window_count=train_window_count, - val_window_count=val_window_count, - code_sample_count=code_sample_count, - prose_sample_count=prose_sample_count, - conversation_sample_count=conversation_sample_count, - skipped_file_count=skipped_file_count, - failed_file_count=failed_file_count, - warning=warning, - sequence_stats=sequence_stats, - duplicate_report=duplicate_report, - ) - _emit( - progress, - f"Dataset rating: {quality_report['stars']:.1f}/5 stars ({quality_report['label']}, score {quality_report['score']:.1f}/100).", - 94, - ) - - summary = { - "dataset_config": dataclass_to_jsonable(config), - "document_count": document_count, - "character_count": character_count, - "token_count": token_count, - "train_token_count": train_token_count, - "val_token_count": val_token_count, - "train_tokens_path": "train_tokens.npy", - "val_tokens_path": "val_tokens.npy", - "token_storage_format": "npy", - "train_window_count": train_window_count, - "val_window_count": val_window_count, - "context_length": config.context_length, - "sequence_token_stats": sequence_stats, - "code_sample_count": code_sample_count, - "prose_sample_count": prose_sample_count, - "conversation_sample_count": conversation_sample_count, - "dataset_stage": config.dataset_stage, - "conversation_datasets": config.conversation_datasets, - "conversation_sample_limit": config.conversation_sample_limit, - "conversation_dataset_path": str( - config.conversation_dataset_path or ""), - "instruction_dataset_path": str(config.instruction_dataset_path or ""), - "conversation_dataset_paths": [str(path) for path in - config.conversation_dataset_paths], - "instruction_dataset_paths": [str(path) for path in - config.instruction_dataset_paths], - "default_data_paths": [str(path) for path in - config.default_data_paths], - "mixture_weights": config.mixture_weights, - "mixture_report": mixture_report, - "exact_duplicate_documents_removed": stats.exact_duplicates_removed, - "exact_duplicate_document_examples": stats.exact_duplicate_examples, - "low_diversity_documents_removed": stats.low_diversity_removed, - "low_diversity_characters_removed": stats.low_diversity_removed_characters, - "low_diversity_duplicate_unit_threshold": MAX_REPETITIVE_UNIT_RATIO, - "low_diversity_document_examples": stats.low_diversity_examples, - "suggested_vocab_size": suggested_vocab_size, - "tokenizer_vocab_size": tokenizer.get_vocab_size(), - "tokenizer_sha256": file_sha256(tokenizer_path), - "warning": warning, - "source_files": stats.source_files, - "source_files_truncated": stats.source_files_truncated, - "cached_file_count": cached_file_count, - "processed_file_count": processed_file_count, - "skipped_file_count": skipped_file_count, - "failed_file_count": failed_file_count, - "source_file_count": manifest.count(), - "prepare_mode": config.prepare_mode, - "tokenizer_strategy": config.tokenizer_strategy, - "reasoning_sample_mode": config.reasoning_sample_mode, - "tokenizer_reused": reuse_tokenizer, - "tokenizer_imported": tokenizer_imported, - "tokenizer_source_path": tokenizer_source_path, - "quality_score": quality_report["score"], - "quality_stars": quality_report["stars"], - "quality_label": quality_report["label"], - "quality_reasons": quality_report["reasons"], - "quality_components": quality_report["components"], - "duplicate_block_count": duplicate_report["duplicate_block_count"], - "unique_block_count": duplicate_report["unique_block_count"], - "corpus_block_count": duplicate_report["block_count"], - "duplicate_block_ratio": duplicate_report["duplicate_block_ratio"], - "unique_block_ratio": duplicate_report["unique_block_ratio"], - "most_repeated_block_count": duplicate_report[ - "most_repeated_block_count"], - "top_repeated_blocks": duplicate_report["top_repeated_blocks"], - } - dataset_version = record_dataset_version(config.output_dir, summary, - manifest) - write_json(config.output_dir / "dataset_summary.json", summary) - manifest.close() - _emit(progress, - f"Dataset version recorded: {dataset_version['version_id']}.", 98) - _emit(progress, f"Dataset ready: {config.output_dir}", 100) - return DatasetBuildResult( - config.output_dir, - tokenizer_path, - document_count, - token_count, - tokenizer.get_vocab_size(), - character_count, - suggested_vocab_size, - train_window_count, - val_window_count, - sequence_stats, - warning, - code_sample_count, - prose_sample_count, - conversation_sample_count, - cached_file_count, - processed_file_count, - skipped_file_count, - failed_file_count, - str(dataset_version["version_id"]), - int(dataset_version["version_number"]), - mixture_report, - float(quality_report["score"]), - float(quality_report["stars"]), - str(quality_report["label"]), - list(quality_report["reasons"]), - int(duplicate_report["duplicate_block_count"]), - int(duplicate_report["unique_block_count"]), - int(duplicate_report["block_count"]), - float(duplicate_report["duplicate_block_ratio"]), - float(duplicate_report["unique_block_ratio"]), - ) - - -def _bounded_ratio(value: float, target: float) -> float: - """Return value/target clamped between 0 and 1. - - Args: - value: Actual metric value. - target: Metric value that should receive full credit. - - Returns: - Clamped ratio. - """ - - if target <= 0: - return 0.0 - return max(0.0, min(1.0, float(value) / float(target))) - - -def _canonical_corpus_block(text: str) -> str: - """Normalize a corpus block for repeated-content checks. - - Args: - text: Raw block text. - - Returns: - Whitespace-normalized lowercase text. - """ - - return re.sub(r"\s+", " ", text).strip().lower() - - -def _dataset_quality_report( - *, - document_count: int, - token_count: int, - vocab_size: int, - unique_words: int, - train_window_count: int, - val_window_count: int, - code_sample_count: int, - prose_sample_count: int, - conversation_sample_count: int, - skipped_file_count: int, - failed_file_count: int, - warning: Optional[str], - sequence_stats: dict[str, float], - duplicate_report: dict[str, Any], -) -> dict[str, Any]: - """Rate a prepared dataset for small-LLM training readiness. - - Args: - document_count: Prepared document/sample count. - token_count: Total token count. - vocab_size: Final tokenizer vocabulary size. - unique_words: Estimated unique words in the corpus. - train_window_count: Number of trainable context windows. - val_window_count: Number of validation context windows. - code_sample_count: Prepared code sample count. - prose_sample_count: Prepared prose sample count. - conversation_sample_count: Prepared conversation/instruction sample count. - skipped_file_count: Empty or unreadable source files skipped. - failed_file_count: Source files that failed extraction. - warning: Size/content warning string. - sequence_stats: Approximate per-document token distribution. - duplicate_report: Repeated text-block report for the written corpus. - - Returns: - Dataset quality dictionary with score, stars, label, and reasons. - """ - - reasons: list[str] = [] - token_score = 30.0 * _bounded_ratio(token_count, 1_000_000) - window_score = 20.0 * _bounded_ratio(train_window_count, 50_000) - vocab_target = max(4_000.0, min(32_000.0, unique_words * 0.8)) - vocab_score = 18.0 * _bounded_ratio(vocab_size, vocab_target) - document_score = 12.0 * _bounded_ratio(document_count, 1_000) - validation_score = 8.0 * _bounded_ratio(val_window_count, 2_000) - families = sum(1 for count in ( - code_sample_count, prose_sample_count, conversation_sample_count) if - count > 0) - diversity_score = 7.0 * _bounded_ratio(families, 3) - average_sequence = float(sequence_stats.get("average", 0.0) or 0.0) - sequence_score = 5.0 * _bounded_ratio(average_sequence, 256) - penalty = min(20.0, failed_file_count * 3.0 + skipped_file_count * 0.5) - duplicate_ratio = float( - duplicate_report.get("duplicate_block_ratio", 0.0) or 0.0) - duplicate_penalty = min(35.0, duplicate_ratio * 70.0) - penalty += duplicate_penalty - if warning and warning != "none": - penalty += 5.0 - score = max( - 0.0, - min( - 100.0, - token_score - + window_score - + vocab_score - + document_score - + validation_score - + diversity_score - + sequence_score - - penalty, - ), - ) - stars = round(score / 20.0 * 2.0) / 2.0 - if score >= 85: - label = "Excellent" - elif score >= 70: - label = "Good" - elif score >= 50: - label = "Usable" - elif score >= 30: - label = "Weak" - else: - label = "Very weak" - if token_count < 250_000: - reasons.append("Token count is low for robust training.") - else: - reasons.append("Token count is sufficient for a small experiment.") - if train_window_count < 5_000: - reasons.append("Few training windows; model may memorize quickly.") - if vocab_size < 4_000: - reasons.append( - "Vocabulary is small; language coverage may be limited.") - elif vocab_size > 50_000: - reasons.append( - "Vocabulary is large; tiny models may spend capacity on tokens.") - else: - reasons.append("Vocabulary size is in a reasonable small-model range.") - if families >= 2: - reasons.append("Dataset includes multiple content families.") - if skipped_file_count or failed_file_count: - reasons.append( - f"Extraction skipped {skipped_file_count} file(s) and failed {failed_file_count} file(s).") - if duplicate_ratio >= 0.5: - reasons.append( - "Prepared corpus is heavily repeated; training may memorize instead of generalize.") - elif duplicate_ratio >= 0.2: - reasons.append( - "Prepared corpus has many repeated blocks; add more varied data or deduplicate.") - elif duplicate_ratio >= 0.05: - reasons.append("Prepared corpus has some repeated blocks.") - else: - reasons.append("Prepared corpus block diversity looks healthy.") - if warning and warning != "none": - reasons.append(str(warning)) - return { - "score": round(score, 1), - "stars": stars, - "label": label, - "reasons": reasons, - "components": { - "tokens": round(token_score, 1), - "windows": round(window_score, 1), - "vocabulary": round(vocab_score, 1), - "documents": round(document_score, 1), - "validation": round(validation_score, 1), - "diversity": round(diversity_score, 1), - "sequence": round(sequence_score, 1), - "duplicate_penalty": round(duplicate_penalty, 1), - "penalty": round(penalty, 1), - }, - } - - -__all__ = [ - "DatasetBuildResult", - "build_dataset", - "estimate_vocab_size", - "content_warning", -] \ No newline at end of file diff --git a/llm_trainer/dataset_mixture.py b/llm_trainer/dataset_mixture.py deleted file mode 100644 index 066936c..0000000 --- a/llm_trainer/dataset_mixture.py +++ /dev/null @@ -1,375 +0,0 @@ -from __future__ import annotations - -import hashlib -import logging -import re -from pathlib import Path -from typing import Any, Callable, Optional - -from .conversation_datasets import CONVERSATION_DATASET_PRESETS -from .data import Document, SUPPORTED_CODE_SUFFIXES - -LOGGER = logging.getLogger(__name__) - -AGGREGATE_MIXTURE_FAMILIES: set[str] = set() -MIXTURE_CHUNK_CHARS = 25_000 -# A default corpus must not gain apparent scale by repeating a tiny template. -# This threshold is deliberately conservative: it only applies once a document -# has enough independently meaningful units to make the measurement useful. -MAX_REPETITIVE_UNIT_RATIO = 0.35 -MIN_REPETITION_CHECK_UNITS = 20 -MIN_REPETITION_CHECK_CHARS = 2_000 - - -def _emit(progress: Optional[Callable[[Any], None]], message: str, percent: Optional[int] = None) -> None: - LOGGER.info(message) - if progress: - progress({"message": message, "percent": percent}) - - -def _canonical_corpus_block(text: str) -> str: - return re.sub(r"\s+", " ", text).strip().lower() - - -def _slugify_category(value: str) -> str: - slug = re.sub(r"[^a-z0-9]+", "_", value.lower()).strip("_") - return slug or "general_prose" - - -def _mixture_label(category: str) -> str: - return category.replace("_", " ").title() - - -def _default_data_category(path: Path) -> Optional[str]: - parts_lower = [part.lower() for part in path.parts] - data_roots = ("default_data", "training_data") - root_indices = [parts_lower.index(root) for root in data_roots if root in parts_lower] - if root_indices: - default_index = max(root_indices) - relative_parts = path.parts[default_index + 1 :] - else: - # Prepared documents can come from any user-selected root. In that - # case the containing folder is the only reliable folder metadata. - relative_parts = path.parts - # Categories are configured by directory layout. The first directory - # under the configured training-data root owns all nested files. - if len(relative_parts) > 1: - return _slugify_category(relative_parts[0]) - return None - - -def _deduplicate_documents(documents: list[Document]) -> tuple[list[Document], dict[str, Any]]: - unique_documents: list[Document] = [] - seen: dict[str, Document] = {} - duplicates: list[dict[str, str]] = [] - for document in documents: - canonical_text = _canonical_corpus_block(document.text) - if not canonical_text: - unique_documents.append(document) - continue - digest = hashlib.sha256( - f"{document.kind}\n{document.language or ''}\n{canonical_text}".encode("utf-8") - ).hexdigest() - original = seen.get(digest) - if original is not None: - duplicates.append( - { - "path": str(document.path), - "duplicate_of": str(original.path), - "kind": document.kind, - } - ) - continue - seen[digest] = document - unique_documents.append(document) - return unique_documents, { - "removed_documents": len(duplicates), - "duplicates": duplicates[:50], - } - - -def _content_units_for_diversity(document: Document) -> list[str]: - """Return comparable content units for a repetition-quality check. - - Prose sources are often whitespace-normalised during ingestion, so using - source lines would miss repeated sentences. Code remains line-oriented; - prose and chat are instead split at sentence and turn boundaries. - """ - - text = document.text.replace("\r\n", "\n").replace("\r", "\n") - if document.kind == "code": - raw_units = text.split("\n") - else: - raw_units = re.split(r"(?<=[.!?])\s+|\n+(?=(?:User|Assistant|System|Instruction|Response):)", text) - return [ - _canonical_corpus_block(unit) - for unit in raw_units - if len(_canonical_corpus_block(unit)) >= 24 - ] - - -def _filter_repetitive_documents(documents: list[Document]) -> tuple[list[Document], dict[str, Any]]: - """Remove documents dominated by exact repeated content units. - - This is a quality gate, not a substitute for semantic deduplication. It - catches generated padding such as the old bundled curriculum files before - it can dominate token counts and make a small corpus look large. - """ - - accepted: list[Document] = [] - rejected: list[dict[str, Any]] = [] - for document in documents: - units = _content_units_for_diversity(document) - if len(document.text) < MIN_REPETITION_CHECK_CHARS or len(units) < MIN_REPETITION_CHECK_UNITS: - accepted.append(document) - continue - duplicate_ratio = 1.0 - (len(set(units)) / len(units)) - if duplicate_ratio > MAX_REPETITIVE_UNIT_RATIO: - rejected.append( - { - "path": str(document.path), - "kind": document.kind, - "unit_count": len(units), - "duplicate_unit_ratio": round(duplicate_ratio, 4), - } - ) - continue - accepted.append(document) - rejected_paths = {item["path"] for item in rejected} - return accepted, { - "removed_documents": len(rejected), - "removed_characters": sum( - len(document.text) - for document in documents - if str(document.path) in rejected_paths - ), - "threshold": MAX_REPETITIVE_UNIT_RATIO, - "examples": rejected[:50], - } - - -def _document_mixture_family(document: Document) -> str: - default_category = _default_data_category(document.path) - if default_category: - return default_category - if document.kind == "code": - return "source_code" - if document.kind == "instruction": - return "instruction" - if document.kind == "conversation": - return "conversation" - dataset_id = str(document.language or "") - preset = CONVERSATION_DATASET_PRESETS.get(dataset_id) - if preset and preset.stage == "base": - return "online_base" - if "__hf_datasets__" in document.path.parts: - for part in document.path.parts: - preset = CONVERSATION_DATASET_PRESETS.get(part) - if preset and preset.stage == "base": - return "online_base" - return "local_prose" - - -def _stable_document_sort_key(document: Document) -> str: - text_digest = hashlib.sha256(document.text[:4096].encode("utf-8", errors="ignore")).hexdigest() - key = f"{document.path}|{document.kind}|{document.language or ''}|{len(document.text)}|{text_digest}" - return hashlib.sha256(key.encode("utf-8")).hexdigest() - - -def _chunk_document_for_mixture(document: Document, chunk_chars: int = MIXTURE_CHUNK_CHARS) -> list[Document]: - text = document.text - if len(text) <= chunk_chars: - return [document] - chunks: list[Document] = [] - start = 0 - while start < len(text): - end = min(len(text), start + chunk_chars) - if end < len(text): - boundary = text.rfind("\n\n", start, end) - if boundary <= start + int(chunk_chars * 0.5): - boundary = text.rfind("\n", start, end) - if boundary > start + int(chunk_chars * 0.5): - end = boundary - chunk_text = text[start:end].strip() - if chunk_text: - chunks.append( - Document( - path=document.path, - text=chunk_text, - kind=document.kind, - language=document.language, - ) - ) - start = max(end, start + 1) - return chunks or [document] - - -def _chunk_documents_for_mixture(documents: list[Document]) -> list[Document]: - chunks: list[Document] = [] - for document in documents: - chunks.extend(_chunk_document_for_mixture(document)) - return chunks - - -def _empty_mixture_report(weights: dict[str, float], documents: list[Document], applied: bool, reason: str = "") -> dict[str, Any]: - document_families = {_document_mixture_family(document) for document in documents} - families_to_report = sorted({*MIXTURE_LABELS, *weights, *document_families}) - by_family: dict[str, list[Document]] = {key: [] for key in families_to_report} - for document in documents: - by_family.setdefault(_document_mixture_family(document), []).append(document) - total_chars = sum(len(document.text) for document in documents) - families = {} - for family in families_to_report: - available = by_family.get(family, []) - available_chars = sum(len(document.text) for document in available) - families[family] = { - "label": _mixture_label(family), - "requested_weight": float(weights.get(family, 0.0) or 0.0), - "available_documents": len(available), - "available_characters": available_chars, - "selected_documents": len(available) if not applied else 0, - "selected_characters": available_chars if not applied else 0, - "actual_percent": (available_chars * 100.0 / total_chars) if total_chars else 0.0, - "dropped_documents": 0, - "dropped_characters": 0, - } - return { - "applied": applied, - "reason": reason, - "total_available_documents": len(documents), - "total_selected_documents": len(documents) if not applied else 0, - "total_available_characters": total_chars, - "total_selected_characters": total_chars if not applied else 0, - "families": families, - } - - -def _apply_dataset_mixture( - documents: list[Document], - weights: dict[str, float], - progress: Optional[Callable[[Any], None]], -) -> tuple[list[Document], dict[str, Any]]: - """ - Apply Dataset Blueprint percentages independently. - - Unlike the previous implementation, percentages are NOT normalized. - - 100% means: - Include every document in that category. - - 50% means: - Include approximately half of the documents from that category. - - Categories never reduce one another. - """ - - original_document_count = len(documents) - - documents = _chunk_documents_for_mixture(documents) - - if len(documents) != original_document_count: - _emit( - progress, - f"Dataset mixture: split {original_document_count:,} source file(s) into " - f"{len(documents):,} sampling chunk(s).", - 49, - ) - - # --------------------------------------------------------- - # Group documents by category - # --------------------------------------------------------- - - families: dict[str, list[Document]] = {} - - for doc in documents: - family = _document_mixture_family(doc) - families.setdefault(family, []).append(doc) - - selected_documents: list[Document] = [] - - report = { - "applied": True, - "reason": "", - "total_available_documents": len(documents), - "total_selected_documents": 0, - "total_available_characters": sum(len(d.text) for d in documents), - "total_selected_characters": 0, - "families": {}, - } - - # --------------------------------------------------------- - # Process each category independently - # --------------------------------------------------------- - - selected_documents = [] - - for family in sorted(families.keys()): - - docs = families[family] - - docs.sort(key=_stable_document_sort_key) - - percentage = float(weights.get(family, 100.0)) - - percentage = max(0.0, min(100.0, percentage)) - - available_documents = len(docs) - available_characters = sum(len(d.text) for d in docs) - - if percentage >= 100.0: - - chosen = docs - - elif percentage <= 0.0: - - chosen = [] - - else: - - keep = round(available_documents * percentage / 100.0) - - chosen = docs[:keep] - - selected_documents.extend(chosen) - - selected_characters = sum(len(d.text) for d in chosen) - - report["families"][family] = { - "label": _mixture_label(family), - "requested_weight": percentage, - "available_documents": available_documents, - "available_characters": available_characters, - "selected_documents": len(chosen), - "selected_characters": selected_characters, - "effective_requested_percent": percentage, - "actual_percent": ( - len(chosen) * 100.0 / available_documents - if available_documents - else 0.0 - ), - "dropped_documents": available_documents - len(chosen), - "dropped_characters": available_characters - selected_characters, - } - - report["total_selected_documents"] = len(selected_documents) - report["total_selected_characters"] = sum( - len(d.text) for d in selected_documents - ) - - _emit( - progress, - f"Dataset mixture selected " - f"{len(selected_documents):,} of {len(documents):,} document chunks.", - 50, - ) - - return selected_documents, report - -__all__ = [ - "MIXTURE_LABELS", - "AGGREGATE_MIXTURE_FAMILIES", - "MIXTURE_CHUNK_CHARS", - "_apply_dataset_mixture", - "_filter_repetitive_documents", - "MAX_REPETITIVE_UNIT_RATIO", -] diff --git a/llm_trainer/dataset_preview.py b/llm_trainer/dataset_preview.py deleted file mode 100644 index 089208d..0000000 --- a/llm_trainer/dataset_preview.py +++ /dev/null @@ -1,729 +0,0 @@ -from __future__ import annotations - -import hashlib -import json -import logging -import os -import re -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any, Callable, Optional - -import PyPDF2 -import torch - -from .config import DatasetConfig -from .conversation_datasets import CONVERSATION_DATASET_PRESETS -from .data import ( - SUPPORTED_CODE_SUFFIXES, - SUPPORTED_TEXT_SUFFIXES, - file_fingerprint, - load_structured_json_documents, - supported_source_paths, -) -from .dataset_build import _local_structured_dataset_paths -from .lineage import read_json, write_json - -LOGGER = logging.getLogger(__name__) - - -@dataclass -class ProjectHealthResult: - status: str - checks: list[dict[str, str]] - summary: str - - -@dataclass -class DatasetPreviewResult: - source_file_count: int - prepared: bool - total_bytes: int - suffix_counts: dict[str, int] - sample_previews: list[dict[str, str]] - issues: list[str] - summary: dict[str, Any] - duplicate_count: int = 0 - duplicate_groups: list[dict[str, Any]] = field(default_factory=list) - bad_extraction_count: int = 0 - bad_extraction_files: list[dict[str, str]] = field(default_factory=list) - code_preview_count: int = 0 - prose_preview_count: int = 0 - balance_label: str = "Unknown" - readiness_score: int = 0 - readiness_label: str = "Unknown" - readiness_reasons: list[str] = field(default_factory=list) - - -def _emit(progress: Optional[Callable[[Any], None]], message: str, percent: Optional[int] = None) -> None: - LOGGER.info(message) - if progress: - progress({"message": message, "percent": percent}) - - -def _health_check(name: str, status: str, detail: str) -> dict[str, str]: - return {"name": name, "status": status, "detail": detail} - - -def check_project_health( - input_dir: Path, - dataset_dir: Path, - model_dir: Path, - export_dir: Path, - gguf_path: Optional[Path], - llama_cpp_dir: Optional[Path], - training_device: str, - progress: Optional[Callable[[Any], None]] = None, - should_stop: Optional[Callable[[], bool]] = None, -) -> ProjectHealthResult: - checks: list[dict[str, str]] = [] - _emit(progress, "Checking project paths...", 10) - if should_stop and should_stop(): - raise RuntimeError("Project health check stopped by user.") - - if input_dir.exists() and input_dir.is_dir(): - try: - source_count = len(supported_source_paths(input_dir, code_training_mode=True, include_source_code=True)) - if source_count: - checks.append(_health_check("Source vault", "ok", f"{source_count:,} supported file(s) found.")) - else: - checks.append(_health_check("Source vault", "warning", "Folder exists, but no supported files were found.")) - except Exception as exc: - checks.append(_health_check("Source vault", "error", str(exc))) - else: - checks.append(_health_check("Source vault", "warning", f"Folder not found: {input_dir}")) - - _emit(progress, "Checking prepared dataset artifacts...", 30) - required_dataset = ("tokenizer.json", "dataset_summary.json") - missing_dataset = [name for name in required_dataset if not (dataset_dir / name).exists()] - if not ( - (dataset_dir / "train_tokens.npy").exists() - and (dataset_dir / "val_tokens.npy").exists() - or (dataset_dir / "train_tokens.json").exists() - and (dataset_dir / "val_tokens.json").exists() - ): - missing_dataset.append("train_tokens.(npy/json), val_tokens.(npy/json)") - if not dataset_dir.exists(): - checks.append(_health_check("Dataset core", "warning", f"Dataset folder not found yet: {dataset_dir}")) - elif missing_dataset: - checks.append(_health_check("Dataset core", "warning", f"Missing prepared artifact(s): {', '.join(missing_dataset)}")) - else: - summary = read_json(dataset_dir / "dataset_summary.json", default={}) or {} - tokens = int(summary.get("token_count", 0) or 0) - vocab = int(summary.get("tokenizer_vocab_size", 0) or 0) - train_windows = int(summary.get("train_window_count", 0) or 0) - val_windows = int(summary.get("val_window_count", 0) or 0) - window_text = f", {train_windows:,}/{val_windows:,} train/val window(s)" if train_windows or val_windows else "" - checks.append(_health_check("Dataset core", "ok", f"Prepared with {tokens:,} token(s), vocab {vocab:,}{window_text}.")) - - _emit(progress, "Checking model artifacts...", 50) - if not model_dir.exists(): - checks.append(_health_check("Model output", "warning", f"Model folder not found yet: {model_dir}")) - elif (model_dir / "final_model.pt").exists(): - checks.append(_health_check("Model output", "ok", "final_model.pt found.")) - elif (model_dir / "checkpoints").exists() and any((model_dir / "checkpoints").glob("*.pt")): - checks.append(_health_check("Model output", "warning", "Checkpoints found, but final_model.pt is not present yet.")) - else: - checks.append(_health_check("Model output", "warning", "No trained model or checkpoint found yet.")) - - _emit(progress, "Checking export and GGUF paths...", 65) - if export_dir.exists(): - artifact_count = sum(1 for item in export_dir.iterdir()) - checks.append(_health_check("Export bay", "ok" if artifact_count else "warning", f"{artifact_count:,} item(s) in export folder.")) - else: - checks.append(_health_check("Export bay", "warning", f"Export folder not found yet: {export_dir}")) - - if gguf_path and str(gguf_path).strip(): - if gguf_path.exists(): - checks.append(_health_check("GGUF chat model", "ok", f"Found {gguf_path.name}.")) - else: - checks.append(_health_check("GGUF chat model", "warning", f"GGUF file not found: {gguf_path}")) - else: - checks.append(_health_check("GGUF chat model", "warning", "No GGUF model selected for chat.")) - - if llama_cpp_dir and str(llama_cpp_dir).strip(): - converter = llama_cpp_dir / "convert_hf_to_gguf.py" - checks.append( - _health_check( - "llama.cpp", - "ok" if converter.exists() else "warning", - "convert_hf_to_gguf.py found." if converter.exists() else f"Converter not found in {llama_cpp_dir}.", - ) - ) - - _emit(progress, "Checking hardware/runtime...", 85) - if training_device == "cuda": - if torch.cuda.is_available(): - checks.append(_health_check("Training device", "ok", f"CUDA ready: {torch.cuda.get_device_name(0)}.")) - else: - checks.append(_health_check("Training device", "error", "CUDA selected but PyTorch cannot use CUDA.")) - else: - checks.append(_health_check("Training device", "ok", "CPU selected. Training will be slower but compatible.")) - - statuses = [check["status"] for check in checks] - if "error" in statuses: - status = "error" - elif "warning" in statuses: - status = "warning" - else: - status = "ok" - summary = f"{sum(1 for item in statuses if item == 'ok')} ok, {sum(1 for item in statuses if item == 'warning')} warning, {sum(1 for item in statuses if item == 'error')} error." - _emit(progress, f"Project health check complete: {summary}", 100) - return ProjectHealthResult(status=status, checks=checks, summary=summary) - - -def _supported_source_paths_cancellable( - input_dir: Path, - code_training_mode: bool, - include_source_code: bool, - should_stop: Optional[Callable[[], bool]], -) -> list[Path]: - if not input_dir.exists(): - raise FileNotFoundError(f"Input folder does not exist: {input_dir}") - allowed = set(SUPPORTED_TEXT_SUFFIXES) | {".pdf", ".jsonl"} - if code_training_mode and include_source_code: - allowed |= set(SUPPORTED_CODE_SUFFIXES) - paths: list[Path] = [] - for root, dirs, files in os.walk(input_dir): - if should_stop and should_stop(): - raise RuntimeError("Dataset preview stopped by user.") - dirs[:] = [ - name - for name in dirs - if name not in {".git", "__pycache__", ".venv", "venv", "node_modules"} - ] - for filename in files: - if should_stop and should_stop(): - raise RuntimeError("Dataset preview stopped by user.") - path = Path(root) / filename - if path.suffix.lower() in allowed: - paths.append(path) - return sorted(paths) - - -def _preview_supported_document( - path: Path, - config: DatasetConfig, - should_stop: Optional[Callable[[], bool]], - max_chars: int = 1200, -) -> Optional[dict[str, str]]: - if should_stop and should_stop(): - raise RuntimeError("Dataset preview stopped by user.") - suffix = path.suffix.lower() - kind = "prose" - language = "" - text = "" - if config.code_training_mode and suffix in SUPPORTED_CODE_SUFFIXES: - kind = "code" - language = SUPPORTED_CODE_SUFFIXES[suffix] - with path.open("rb") as file: - text = file.read(128 * 1024).decode("utf-8", errors="ignore") - text = text.replace("\x00", "").replace("\r\n", "\n").replace("\r", "\n").strip() - elif suffix in SUPPORTED_TEXT_SUFFIXES: - with path.open("rb") as file: - text = file.read(128 * 1024).decode("utf-8", errors="ignore") - text = re.sub(r"\s+", " ", text.replace("\x00", " ")).strip() - elif suffix == ".jsonl": - chunks: list[str] = [] - with path.open("r", encoding="utf-8", errors="ignore") as file: - for index, line in enumerate(file): - if should_stop and should_stop(): - raise RuntimeError("Dataset preview stopped by user.") - if index >= 40 or sum(len(chunk) for chunk in chunks) >= max_chars: - break - if not line.strip(): - continue - try: - value = json.loads(line) - except json.JSONDecodeError: - continue - if isinstance(value, str): - chunks.append(value) - elif isinstance(value, dict): - for key in ("text", "content", "prompt", "completion"): - if value.get(key): - chunks.append(str(value[key])) - break - text = "\n".join(chunks).strip() - elif suffix == ".pdf": - chunks: list[str] = [] - with path.open("rb") as file: - reader = PyPDF2.PdfReader(file) - for page in reader.pages[:3]: - if should_stop and should_stop(): - raise RuntimeError("Dataset preview stopped by user.") - chunks.append(page.extract_text() or "") - if sum(len(chunk) for chunk in chunks) >= max_chars: - break - text = re.sub(r"\s+", " ", "\n".join(chunks).replace("\x00", " ")).strip() - if config.lowercase: - text = text.lower() - if not text: - return None - return { - "path": str(path), - "kind": kind, - "language": language, - "characters": str(len(text)), - "preview": text[:max_chars], - } - - -def _file_sha256_cancellable(path: Path, should_stop: Optional[Callable[[], bool]]) -> str: - digest = hashlib.sha256() - with path.open("rb") as file: - for chunk in iter(lambda: file.read(1024 * 1024), b""): - if should_stop and should_stop(): - raise RuntimeError("Dataset preview stopped by user.") - digest.update(chunk) - return digest.hexdigest() - - -def _has_prepared_token_artifacts(dataset_dir: Path) -> bool: - if not (dataset_dir / "tokenizer.json").exists(): - return False - has_npy_tokens = (dataset_dir / "train_tokens.npy").exists() and (dataset_dir / "val_tokens.npy").exists() - has_json_tokens = (dataset_dir / "train_tokens.json").exists() and (dataset_dir / "val_tokens.json").exists() - return has_npy_tokens or has_json_tokens - - -def _preview_fingerprint(text: str) -> str: - normalized = re.sub(r"\s+", " ", text.lower()).strip() - normalized = re.sub(r"\d+", "0", normalized) - return hashlib.sha1(normalized[:4000].encode("utf-8", errors="ignore")).hexdigest() - - -def _bad_extraction_reasons(path: Path, preview: Optional[dict[str, str]], size: int) -> list[str]: - suffix = path.suffix.lower() - if preview is None: - return ["no readable preview text"] - text = str(preview.get("preview", "")) - if not text.strip(): - return ["empty preview text"] - reasons: list[str] = [] - visible = [char for char in text if not char.isspace()] - if len(text) < 80 and suffix == ".pdf": - reasons.append("very little text extracted from PDF preview") - if size > 250_000 and suffix == ".pdf" and len(text) < 200: - reasons.append("large PDF produced very little readable text") - if visible: - alpha_ratio = sum(char.isalpha() for char in visible) / len(visible) - symbol_ratio = sum(not char.isalnum() for char in visible) / len(visible) - if alpha_ratio < 0.25 and str(preview.get("kind")) != "code": - reasons.append("low alphabetic text ratio") - if symbol_ratio > 0.45: - reasons.append("high symbol/noise ratio") - if re.search(r"(.)\1{18,}", text): - reasons.append("long repeated character run") - if text.count("\ufffd") >= 3 or "Ã" in text[:500]: - reasons.append("encoding artifacts detected") - words = re.findall(r"[A-Za-z]{2,}", text) - if suffix in {".pdf", ".txt", ".md", ".text"} and len(set(words)) < 8 and len(text) > 200: - reasons.append("very low word variety") - return reasons - - -def _balance_label(code_count: int, prose_count: int, code_training_mode: bool) -> str: - total = code_count + prose_count - if total <= 0: - return "Unknown" - code_ratio = code_count / total - if not code_training_mode: - return "Prose focused" - if code_ratio < 0.2: - return "Prose heavy" - if code_ratio > 0.8: - return "Code heavy" - return "Balanced code/prose" - - -def _readiness_report( - source_file_count: int, - total_bytes: int, - prepared: bool, - summary: dict[str, Any], - duplicate_count: int, - bad_extraction_count: int, - code_count: int, - prose_count: int, - code_training_mode: bool, -) -> tuple[int, str, list[str]]: - score = 100 - reasons: list[str] = [] - token_count = int(summary.get("token_count", 0) or 0) - if prepared and token_count: - if token_count < 50_000: - score -= 35 - reasons.append("Prepared token count is very small.") - elif token_count < 250_000: - score -= 18 - reasons.append("Prepared token count is modest.") - else: - reasons.append("Prepared token count looks usable.") - else: - score -= 15 - reasons.append("Dataset is not prepared yet; score is based on source preview.") - if total_bytes < 250_000: - score -= 25 - reasons.append("Source size is small for meaningful training.") - elif total_bytes < 2_000_000: - score -= 10 - reasons.append("Source size is modest.") - - if source_file_count == 0: - score -= 60 - reasons.append("No supported source files were found.") - duplicate_ratio = duplicate_count / max(source_file_count, 1) - if duplicate_ratio >= 0.25: - score -= 25 - reasons.append("Duplicate ratio is high.") - elif duplicate_ratio >= 0.1: - score -= 12 - reasons.append("Some likely duplicate files were found.") - - bad_ratio = bad_extraction_count / max(source_file_count, 1) - if bad_ratio >= 0.2: - score -= 25 - reasons.append("Many files have suspicious extraction quality.") - elif bad_ratio > 0: - score -= 10 - reasons.append("Some files have suspicious extraction quality.") - - total_samples = code_count + prose_count - if code_training_mode and total_samples: - code_ratio = code_count / total_samples - if code_ratio < 0.1: - score -= 12 - reasons.append("Code training mode is enabled but very little code was detected.") - elif code_ratio > 0.95 and prose_count == 0: - score -= 5 - reasons.append("Dataset is almost entirely code; explanations may be weak.") - - score = max(0, min(100, score)) - if score >= 80: - label = "Ready" - elif score >= 60: - label = "Usable with warnings" - elif score >= 35: - label = "Needs cleanup" - else: - label = "Not ready" - return score, label, reasons - - -def scan_dataset_preview( - config: DatasetConfig, - sample_limit: int = 8, - progress: Optional[Callable[[Any], None]] = None, - should_stop: Optional[Callable[[], bool]] = None, -) -> DatasetPreviewResult: - _emit(progress, "Scanning supported source files...", 5) - local_structured_paths = _local_structured_dataset_paths(config) - if config.input_dir.exists(): - paths = _supported_source_paths_cancellable( - config.input_dir, - config.code_training_mode, - config.include_source_code, - should_stop, - ) - elif config.conversation_datasets or local_structured_paths: - paths = [] - else: - paths = _supported_source_paths_cancellable( - config.input_dir, - config.code_training_mode, - config.include_source_code, - should_stop, - ) - suffix_counts: dict[str, int] = {} - total_bytes = 0 - for path in paths: - if should_stop and should_stop(): - raise RuntimeError("Dataset preview stopped by user.") - suffix_counts[path.suffix.lower() or ""] = suffix_counts.get(path.suffix.lower() or "", 0) + 1 - try: - total_bytes += path.stat().st_size - except OSError: - pass - for path, _, _ in local_structured_paths: - if should_stop and should_stop(): - raise RuntimeError("Dataset preview stopped by user.") - path = Path(path) - if path.exists() and path.is_file(): - suffix_counts[path.suffix.lower() or ""] = suffix_counts.get(path.suffix.lower() or "", 0) + 1 - try: - total_bytes += path.stat().st_size - except OSError: - pass - - summary = read_json(config.output_dir / "dataset_summary.json", default={}) or {} - prepared = _has_prepared_token_artifacts(config.output_dir) - issues: list[str] = [] - if not paths and not config.conversation_datasets and not local_structured_paths: - issues.append("No supported source files found.") - if config.conversation_dataset_path: - issues.append(f"Local conversation JSON selected: {config.conversation_dataset_path}.") - if config.instruction_dataset_path: - issues.append(f"Local instruction JSON selected: {config.instruction_dataset_path}.") - if config.conversation_datasets: - labels = [ - CONVERSATION_DATASET_PRESETS[item].label - for item in config.conversation_datasets - if item in CONVERSATION_DATASET_PRESETS - ] - issues.append(f"Conversation datasets selected: {', '.join(labels)}.") - if total_bytes < 100_000: - issues.append("Source content appears small for meaningful LLM training.") - if prepared and summary: - token_count = int(summary.get("token_count", 0) or 0) - if token_count < 50_000: - issues.append("Prepared token count is low; expect smoke-test quality only.") - if summary.get("warning"): - issues.append(str(summary["warning"])) - elif config.output_dir.exists(): - issues.append("Dataset folder exists but does not contain a complete prepared dataset.") - - _emit(progress, "Reading a few preview samples...", 30) - sample_previews: list[dict[str, str]] = [] - all_previews: list[dict[str, str]] = [] - bad_extraction_files: list[dict[str, str]] = [] - content_fingerprints: dict[str, list[str]] = {} - preview_cache_path = config.output_dir / "preview_scan_cache.json" - preview_cache = read_json(preview_cache_path, default={}) or {} - cached_files = preview_cache.get("files") if isinstance(preview_cache.get("files"), dict) else {} - updated_cache: dict[str, Any] = {} - readable = 0 - scan_limit = min(len(paths), max(sample_limit * 8, 80)) - for path in paths[:scan_limit]: - if should_stop and should_stop(): - raise RuntimeError("Dataset preview stopped by user.") - try: - stat = path.stat() - size = stat.st_size - mtime_ns = stat.st_mtime_ns - except OSError: - size = 0 - mtime_ns = 0 - cache_key = str(path.resolve()) - cached_entry = cached_files.get(cache_key, {}) if isinstance(cached_files, dict) else {} - cached_preview = cached_entry.get("preview") - cached_reasons = cached_entry.get("bad_extraction_reasons") - if ( - cached_entry.get("size") == size - and cached_entry.get("mtime_ns") == mtime_ns - and isinstance(cached_preview, dict) - and isinstance(cached_reasons, list) - ): - preview = cached_preview - reasons = [str(item) for item in cached_reasons] - else: - try: - preview = _preview_supported_document(path, config, should_stop) - except RuntimeError: - raise - except Exception as exc: - issues.append(f"Could not preview {path.name}: {exc}") - continue - reasons = _bad_extraction_reasons(path, preview, size) - updated_cache[cache_key] = { - "size": size, - "mtime_ns": mtime_ns, - "preview": preview, - "bad_extraction_reasons": reasons, - "duplicate_digest": cached_entry.get("duplicate_digest"), - "duplicate_digest_mode": cached_entry.get("duplicate_digest_mode"), - "strict_duplicate_digest": cached_entry.get("strict_duplicate_digest"), - } - if reasons: - bad_extraction_files.append( - { - "path": str(path), - "reasons": "; ".join(reasons), - "size": str(size), - } - ) - if preview is None: - issues.append(f"{path.name} has no readable text.") - continue - readable += 1 - all_previews.append(preview) - preview_text = preview.get("preview", "") - if len(preview_text) >= 120: - content_fingerprints.setdefault(_preview_fingerprint(preview_text), []).append(str(path)) - if len(sample_previews) < sample_limit: - sample_previews.append(preview) - percent = 30 + int(45 * len(sample_previews) / max(sample_limit, 1)) - _emit(progress, f"Previewed {path.name}.", percent) - - for local_path, kind, _ in local_structured_paths: - if len(sample_previews) >= sample_limit: - continue - if should_stop and should_stop(): - raise RuntimeError("Dataset preview stopped by user.") - try: - local_documents = load_structured_json_documents(Path(local_path), kind=kind, lowercase=config.lowercase) - except Exception as exc: - issues.append(f"Could not preview {kind} JSON dataset: {exc}") - continue - for document in local_documents[: max(1, sample_limit - len(sample_previews))]: - preview = { - "path": str(document.path), - "kind": document.kind, - "language": document.language or "", - "characters": str(len(document.text)), - "preview": document.text[:1200], - } - all_previews.append(preview) - sample_previews.append(preview) - _emit(progress, f"Previewed {kind} JSON sample.", 75) - if len(sample_previews) >= sample_limit: - break - - if readable == 0 and paths: - issues.append("Supported files were found, but none produced readable preview text.") - _emit(progress, "Checking duplicate files and repeated extracted text...", 85) - size_groups: dict[int, list[Path]] = {} - for path in paths: - if should_stop and should_stop(): - raise RuntimeError("Dataset preview stopped by user.") - try: - size_groups.setdefault(path.stat().st_size, []).append(path) - except OSError: - continue - exact_hashes: dict[str, list[str]] = {} - for same_size_paths in size_groups.values(): - if len(same_size_paths) < 2: - continue - for path in same_size_paths: - cache_key = str(path.resolve()) - entry = updated_cache.get(cache_key) - try: - stat = path.stat() - size = stat.st_size - mtime_ns = stat.st_mtime_ns - except OSError: - continue - if entry is None: - cached_entry = cached_files.get(cache_key, {}) if isinstance(cached_files, dict) else {} - entry = { - "size": size, - "mtime_ns": mtime_ns, - "preview": cached_entry.get("preview"), - "bad_extraction_reasons": cached_entry.get("bad_extraction_reasons"), - "duplicate_digest": cached_entry.get("duplicate_digest"), - "duplicate_digest_mode": cached_entry.get("duplicate_digest_mode"), - "strict_duplicate_digest": cached_entry.get("strict_duplicate_digest"), - } - digest_mode = "fast" if config.fast_scan_mode else "full" - if ( - entry.get("size") == size - and entry.get("mtime_ns") == mtime_ns - and entry.get("duplicate_digest_mode") == digest_mode - and isinstance(entry.get("duplicate_digest"), str) - ): - digest = str(entry["duplicate_digest"]) - else: - if config.fast_scan_mode: - digest = file_fingerprint(path, fast=True, sample_bytes=config.fast_scan_sample_bytes) - else: - digest = _file_sha256_cancellable(path, should_stop) - entry["duplicate_digest"] = digest - entry["duplicate_digest_mode"] = digest_mode - entry["size"] = size - entry["mtime_ns"] = mtime_ns - updated_cache[cache_key] = entry - exact_hashes.setdefault(digest, []).append(str(path)) - if config.fast_scan_mode and config.strict_duplicate_verification: - verified_hashes: dict[str, list[str]] = {} - for fast_group in exact_hashes.values(): - if len(fast_group) < 2: - continue - for path_str in fast_group: - path = Path(path_str) - cache_key = str(path.resolve()) - entry = updated_cache.get(cache_key, {}) - try: - stat = path.stat() - size = stat.st_size - mtime_ns = stat.st_mtime_ns - except OSError: - continue - if ( - entry.get("size") == size - and entry.get("mtime_ns") == mtime_ns - and isinstance(entry.get("strict_duplicate_digest"), str) - ): - strict_digest = str(entry["strict_duplicate_digest"]) - else: - strict_digest = _file_sha256_cancellable(path, should_stop) - entry["strict_duplicate_digest"] = strict_digest - entry["size"] = size - entry["mtime_ns"] = mtime_ns - updated_cache[cache_key] = entry - verified_hashes.setdefault(strict_digest, []).append(path_str) - exact_hashes = verified_hashes - duplicate_groups: list[dict[str, Any]] = [] - for group in exact_hashes.values(): - if len(group) > 1: - duplicate_groups.append({"type": "exact file", "count": len(group), "files": group[:8]}) - for group in content_fingerprints.values(): - unique_group = sorted(set(group)) - if len(unique_group) > 1: - duplicate_groups.append({"type": "similar extracted text", "count": len(unique_group), "files": unique_group[:8]}) - duplicate_files = { - file_path - for group in duplicate_groups - for file_path in group.get("files", []) - } - duplicate_count = len(duplicate_files) - if duplicate_groups: - issues.append(f"Found {len(duplicate_groups)} likely duplicate group(s) involving {duplicate_count} file entries.") - if bad_extraction_files: - issues.append(f"Found {len(bad_extraction_files)} file(s) with suspicious extraction quality.") - summary_code_count = int(summary.get("code_sample_count", 0) or 0) - summary_prose_count = int(summary.get("prose_sample_count", 0) or 0) - code_preview_count = summary_code_count or sum(1 for preview in all_previews if preview.get("kind") == "code") - prose_preview_count = summary_prose_count or sum(1 for preview in all_previews if preview.get("kind") != "code") - balance = _balance_label(code_preview_count, prose_preview_count, config.code_training_mode) - effective_source_count = len(paths) + len(config.conversation_datasets) + len(local_structured_paths) - readiness_score, readiness_label, readiness_reasons = _readiness_report( - source_file_count=effective_source_count, - total_bytes=total_bytes, - prepared=prepared, - summary=summary, - duplicate_count=duplicate_count, - bad_extraction_count=len(bad_extraction_files), - code_count=code_preview_count, - prose_count=prose_preview_count, - code_training_mode=config.code_training_mode, - ) - _emit(progress, "Dataset preview complete.", 100) - if config.output_dir.exists(): - write_json(preview_cache_path, {"version": 1, "files": updated_cache}) - return DatasetPreviewResult( - source_file_count=effective_source_count, - prepared=prepared, - total_bytes=total_bytes, - suffix_counts=dict(sorted(suffix_counts.items())), - sample_previews=sample_previews, - issues=issues, - summary=summary, - duplicate_count=duplicate_count, - duplicate_groups=duplicate_groups, - bad_extraction_count=len(bad_extraction_files), - bad_extraction_files=bad_extraction_files[:30], - code_preview_count=code_preview_count, - prose_preview_count=prose_preview_count, - balance_label=balance, - readiness_score=readiness_score, - readiness_label=readiness_label, - readiness_reasons=readiness_reasons, - ) - - -__all__ = [ - "ProjectHealthResult", - "DatasetPreviewResult", - "check_project_health", - "scan_dataset_preview", -] diff --git a/llm_trainer/document_extraction.py b/llm_trainer/document_extraction.py deleted file mode 100644 index a8f34b5..0000000 --- a/llm_trainer/document_extraction.py +++ /dev/null @@ -1,147 +0,0 @@ -"""Standalone document extraction helpers, safe for worker processes. - -This module intentionally imports only from :mod:`.data`. ``dataset_build.py`` -(via ``.config``, ``.tokenizer``, ``.training``) pulls in ``torch`` and other -heavy dependencies; when a function defined in ``dataset_build.py`` is used as -a ``ProcessPoolExecutor`` target under the ``spawn`` start method, every -worker process has to re-import that entire chain just to resolve the -function, which is slow and memory-heavy for something that only needs to -read and clean one text file. Keeping the actual worker function here instead -means worker processes only pay for what they use. -""" - -from __future__ import annotations - -import re -from pathlib import Path -from typing import Any, Optional - -from .data import ( - Document, - document_to_dict, - expand_code_documents, - read_supported_document, -) - - -def bad_extraction_reasons( - path: Path, - preview: Optional[dict[str, str]], - size: int, -) -> list[str]: - """Return quality reasons when extracted preview text looks suspicious. - - Args: - path: Source file path. - preview: Preview dictionary with at least ``preview`` and ``kind``. - size: Source file size in bytes. - - Returns: - Human-readable list of extraction quality concerns, empty if none. - """ - - suffix = path.suffix.lower() - if preview is None: - return ["no readable preview text"] - text = str(preview.get("preview", "")) - if not text.strip(): - return ["empty preview text"] - reasons: list[str] = [] - visible = [char for char in text if not char.isspace()] - if len(text) < 80 and suffix == ".pdf": - reasons.append("very little text extracted from PDF preview") - if size > 250_000 and suffix == ".pdf" and len(text) < 200: - reasons.append("large PDF produced very little readable text") - if visible: - alpha_ratio = sum(char.isalpha() for char in visible) / len(visible) - symbol_ratio = sum(not char.isalnum() for char in visible) / len(visible) - if alpha_ratio < 0.25 and str(preview.get("kind")) != "code": - reasons.append("low alphabetic text ratio") - if symbol_ratio > 0.45: - reasons.append("high symbol/noise ratio") - if re.search(r"(.)\1{18,}", text): - reasons.append("long repeated character run") - if text.count("\ufffd") >= 3 or "Ã" in text[:500]: - reasons.append("encoding artifacts detected") - words = re.findall(r"[A-Za-z]{2,}", text) - if suffix in {".pdf", ".txt", ".md", ".text"} and len(set(words)) < 8 and len(text) > 200: - reasons.append("very low word variety") - return reasons - - -def extract_documents_worker( - path: Path, - lowercase: bool, - code_training_mode: bool, - preserve_indentation: bool, - include_prose: bool, - extract_code_blocks: bool, -) -> dict[str, Any]: - """Extract and expand one source file into documents. - - Intended as a ``ProcessPoolExecutor`` target: runs inside a worker - process so CPU-bound extraction work (PDF parsing, regex text cleaning, - code-block detection) for multiple large files genuinely runs in - parallel across CPU cores, rather than serialized behind the GIL as with - threads. Only one source file's text is ever resident in a given worker - process at a time. - - Args: - path: Source file path. - lowercase: Whether to lowercase loaded content. - code_training_mode: Enables code-aware loading and expansion. - preserve_indentation: Keeps code formatting where possible. - include_prose: Keeps prose documents in code-aware mode. - extract_code_blocks: Extracts code-like blocks from prose documents. - - Returns: - Dict with the source path, extracted documents (as plain dicts, for - safe pickling back to the main process), an optional error message, - and any bad-extraction-quality reasons. - """ - - try: - source_doc = read_supported_document( - path, - lowercase=lowercase, - code_training_mode=code_training_mode, - preserve_indentation=preserve_indentation, - ) - except Exception as exc: # noqa: BLE001 - reported back to the main process - return {"path": str(path), "documents": [], "error": str(exc), "bad_extraction_reasons": []} - - if source_doc is None: - return {"path": str(path), "documents": [], "error": None, "bad_extraction_reasons": []} - - reasons = bad_extraction_reasons( - path, - { - "path": str(path), - "kind": source_doc.kind, - "language": source_doc.language or "", - "characters": str(len(source_doc.text)), - "preview": source_doc.text[:1200], - }, - path.stat().st_size, - ) - if path.suffix.lower() == ".pdf" and reasons: - return {"path": str(path), "documents": [], "error": None, "bad_extraction_reasons": reasons} - - source_documents: list[Document] = [source_doc] - if code_training_mode: - source_documents = expand_code_documents( - source_documents, - include_prose=include_prose, - extract_code_blocks=extract_code_blocks, - preserve_indentation=preserve_indentation, - should_stop=None, - ) - return { - "path": str(path), - "documents": [document_to_dict(doc) for doc in source_documents], - "error": None, - "bad_extraction_reasons": [], - } - - -__all__ = ["bad_extraction_reasons", "extract_documents_worker"] \ No newline at end of file diff --git a/llm_trainer/evaluation.py b/llm_trainer/evaluation.py deleted file mode 100644 index e33ff9a..0000000 --- a/llm_trainer/evaluation.py +++ /dev/null @@ -1,168 +0,0 @@ -from __future__ import annotations - -import json -from dataclasses import dataclass -from pathlib import Path -from time import perf_counter -from typing import Any, Callable, Optional - -import torch - -from .generation import load_model_from_checkpoint -from .lineage import read_json, utc_timestamp, write_json -from .tokenizer import EOS_TOKEN, load_tokenizer, token_id - - -DEFAULT_BENCHMARK_PROMPTS = [ - "Explain what a Python function is and give a tiny example.", - "Write a Python function that adds two numbers.", - "Review this code and explain any issue:\n```python\ndef add(a, b):\nprint(a + b)\n```", -] - - -@dataclass -class BenchmarkResult: - """Result returned after running benchmark prompts. - - Attributes: - output_path: JSON file containing benchmark outputs. - prompt_count: Number of prompts evaluated. - total_seconds: Total elapsed time. - total_generated_tokens: Total generated tokens across prompts. - tokens_per_second: Average generated-token throughput. - """ - - output_path: Path - prompt_count: int - total_seconds: float - total_generated_tokens: int = 0 - tokens_per_second: float = 0.0 - - -def normalize_prompts(raw_prompts: str) -> list[str]: - """Split raw benchmark prompt text into prompts. - - Args: - raw_prompts: Text containing prompts separated by blank lines. - - Returns: - Non-empty prompt list. - """ - - prompts = [part.strip() for part in raw_prompts.replace("\r\n", "\n").split("\n\n") if part.strip()] - return prompts or DEFAULT_BENCHMARK_PROMPTS - - -def evaluate_checkpoint( - model_dir: Path, - prompts: list[str], - output_dir: Optional[Path] = None, - max_new_tokens: int = 128, - temperature: float = 0.7, - top_k: int = 50, - device: Optional[str] = None, - use_kv_cache: bool = True, - progress: Optional[Callable[[Any], None]] = None, - should_stop: Optional[Callable[[], bool]] = None, -) -> BenchmarkResult: - """Evaluate a trained checkpoint with benchmark prompts. - - Args: - model_dir: Folder containing ``final_model.pt`` and ``tokenizer.json``. - prompts: Benchmark prompts. - output_dir: Optional output folder for benchmark JSON. - max_new_tokens: Maximum new tokens per prompt. - temperature: Sampling temperature. - top_k: Top-k sampling cutoff. - device: Optional device override. - use_kv_cache: Whether to reuse key/value cache during generation. - progress: Optional progress callback. - should_stop: Optional cancellation callback. - - Returns: - Benchmark result summary. - - Raises: - FileNotFoundError: If checkpoint or tokenizer is missing. - """ - - model_dir = Path(model_dir) - checkpoint_path = model_dir / "final_model.pt" - tokenizer_path = model_dir / "tokenizer.json" - if not checkpoint_path.exists(): - raise FileNotFoundError(f"Final checkpoint not found: {checkpoint_path}") - if not tokenizer_path.exists(): - raise FileNotFoundError(f"Tokenizer not found: {tokenizer_path}") - - output_dir = output_dir or model_dir / "benchmarks" - output_dir.mkdir(parents=True, exist_ok=True) - device = device or ("cuda" if torch.cuda.is_available() else "cpu") - lineage = read_json(model_dir / "model_lineage.json", default={}) or {} - tokenizer = load_tokenizer(tokenizer_path) - model = load_model_from_checkpoint(checkpoint_path, device=device) - eos_id = token_id(tokenizer, EOS_TOKEN) - - started = perf_counter() - outputs: list[dict[str, Any]] = [] - for index, prompt in enumerate(prompts, start=1): - if should_stop and should_stop(): - raise RuntimeError("Benchmark stopped by user.") - if progress: - progress({"message": f"Benchmark prompt {index}/{len(prompts)}...", "percent": int(90 * (index - 1) / max(len(prompts), 1))}) - prompt_started = perf_counter() - input_ids = tokenizer.encode(prompt).ids - context = torch.tensor([input_ids], dtype=torch.long, device=device) - generated = model.generate( - context, - max_new_tokens, - temperature=temperature, - top_k=top_k, - use_kv_cache=use_kv_cache, - ) - output_ids = generated[0].tolist() - generated_token_count = max(len(output_ids) - len(input_ids), 0) - if eos_id in output_ids[len(input_ids) :]: - eos_index = output_ids.index(eos_id, len(input_ids)) - output_ids = output_ids[:eos_index] - generated_token_count = max(len(output_ids) - len(input_ids), 0) - text = tokenizer.decode(output_ids) - elapsed = perf_counter() - prompt_started - outputs.append( - { - "index": index, - "prompt": prompt, - "output": text, - "elapsed_seconds": elapsed, - "generated_tokens": generated_token_count, - "tokens_per_second": generated_token_count / max(elapsed, 1e-9), - "characters": len(text), - } - ) - - total = perf_counter() - started - total_generated_tokens = sum(int(item.get("generated_tokens", 0)) for item in outputs) - tokens_per_second = total_generated_tokens / max(total, 1e-9) - payload = { - "schema": "micro_llm_benchmark", - "version": 1, - "created_at": utc_timestamp(), - "model_dir": str(model_dir), - "checkpoint_path": str(checkpoint_path), - "tokenizer_path": str(tokenizer_path), - "device": device, - "max_new_tokens": max_new_tokens, - "temperature": temperature, - "top_k": top_k, - "use_kv_cache": use_kv_cache, - "model_lineage": lineage, - "prompt_count": len(prompts), - "total_seconds": total, - "total_generated_tokens": total_generated_tokens, - "tokens_per_second": tokens_per_second, - "results": outputs, - } - output_path = output_dir / f"benchmark_{utc_timestamp()}.json" - write_json(output_path, payload) - if progress: - progress({"message": f"Benchmark saved: {output_path}", "percent": 100}) - return BenchmarkResult(output_path, len(prompts), total, total_generated_tokens, tokens_per_second) diff --git a/llm_trainer/export.py b/llm_trainer/export.py deleted file mode 100644 index 34814e1..0000000 --- a/llm_trainer/export.py +++ /dev/null @@ -1,406 +0,0 @@ -from __future__ import annotations - -import shutil -import subprocess -import json -from pathlib import Path -from typing import Optional - -import torch - -from .tokenizer import DEFAULT_CHAT_TEMPLATE - - -class ExportError(RuntimeError): - """Raised when model export or quantization cannot be completed.""" - - pass - - -def export_project_bundle(project_dir: Path, output_dir: Path) -> Path: - """Create a portable model bundle. - - Args: - project_dir: Trained model project folder. - output_dir: Destination export folder. - - Returns: - Export folder path. - - Raises: - ExportError: If required model artifacts are missing. - """ - output_dir.mkdir(parents=True, exist_ok=True) - required = ["final_model.pt", "tokenizer.json", "training_summary.json"] - for name in required: - source = project_dir / name - if not source.exists(): - raise ExportError(f"Missing required file for export: {source}") - shutil.copy2(source, output_dir / name) - for metadata_name in ("tokenizer_config.json", "special_tokens_map.json"): - source = project_dir / metadata_name - if source.exists(): - shutil.copy2(source, output_dir / metadata_name) - optional_files = ["model_lineage.json", "dataset_summary.json"] - for name in optional_files: - source = project_dir / name - if source.exists(): - shutil.copy2(source, output_dir / name) - benchmarks = project_dir / "benchmarks" - if benchmarks.exists(): - shutil.copytree(benchmarks, output_dir / "benchmarks", dirs_exist_ok=True) - manifest = { - "schema": "micro_llm_export_bundle", - "project_dir": str(project_dir), - "output_dir": str(output_dir), - "files": sorted(path.name for path in output_dir.iterdir()), - } - (output_dir / "export_summary.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8") - return output_dir - - -def quantize_checkpoint(checkpoint_path: Path, output_path: Path, mode: str = "fp16") -> Path: - """Create a smaller inference checkpoint. - - Args: - checkpoint_path: Source PyTorch checkpoint path. - output_path: Destination quantized checkpoint path. - mode: Quantization mode. Currently only ``fp16`` is supported. - - Returns: - Quantized checkpoint path. - - Raises: - ExportError: If the checkpoint is missing or mode is unsupported. - """ - checkpoint_path = Path(checkpoint_path) - output_path = Path(output_path) - if not checkpoint_path.exists(): - raise ExportError(f"Checkpoint not found: {checkpoint_path}") - - mode = mode.lower() - if mode not in {"fp16", "float16"}: - raise ExportError("Only FP16 checkpoint quantization is currently supported for MicroGPT.") - - checkpoint = torch.load(checkpoint_path, map_location="cpu") - state_dict = checkpoint.get("model_state_dict") - if not state_dict: - raise ExportError("Checkpoint does not contain model_state_dict.") - - checkpoint["model_state_dict"] = { - key: value.half() if torch.is_floating_point(value) else value - for key, value in state_dict.items() - } - checkpoint["quantization"] = { - "mode": "fp16", - "source": str(checkpoint_path), - } - output_path.parent.mkdir(parents=True, exist_ok=True) - torch.save(checkpoint, output_path) - return output_path - - -def export_hf_microgpt_package(project_dir: Path, output_dir: Optional[Path] = None) -> Path: - """Export a MicroGPT checkpoint as an HF-style local model package. - - Args: - project_dir: Trained model project folder. - output_dir: Optional destination folder. Defaults to ``project_dir/hf_model``. - - Returns: - HF-style package folder. - - Raises: - ExportError: If required files are missing. - """ - - project_dir = Path(project_dir) - output_dir = Path(output_dir) if output_dir else project_dir / "hf_model" - checkpoint_path = project_dir / "final_model.pt" - tokenizer_path = project_dir / "tokenizer.json" - summary_path = project_dir / "training_summary.json" - for path in (checkpoint_path, tokenizer_path, summary_path): - if not path.exists(): - raise ExportError(f"Missing required file for HF package: {path}") - - checkpoint = torch.load(checkpoint_path, map_location="cpu") - model_config = checkpoint.get("model_config") - state_dict = checkpoint.get("model_state_dict") - if not isinstance(model_config, dict) or not state_dict: - raise ExportError("Checkpoint must contain model_config and model_state_dict.") - - output_dir.mkdir(parents=True, exist_ok=True) - torch.save(state_dict, output_dir / "pytorch_model.bin") - shutil.copy2(tokenizer_path, output_dir / "tokenizer.json") - shutil.copy2(summary_path, output_dir / "training_summary.json") - for optional in ("model_lineage.json", "dataset_summary.json"): - source = project_dir / optional - if source.exists(): - shutil.copy2(source, output_dir / optional) - - config = { - "model_type": "microgpt", - "architectures": ["MicroGPTForCausalLM"], - "library_name": "micro-llm-creator", - "llama_cpp_convertible": False, - "vocab_size": model_config.get("vocab_size"), - "n_positions": model_config.get("context_length"), - "n_ctx": model_config.get("context_length"), - "n_embd": model_config.get("embedding_size"), - "n_head": model_config.get("head_count"), - "n_layer": model_config.get("layer_count"), - "dropout": model_config.get("dropout"), - "bias": model_config.get("bias"), - "norm_type": model_config.get("norm_type", "layernorm"), - "position_encoding": model_config.get("position_encoding", "learned"), - "mlp_type": model_config.get("mlp_type", "gelu"), - "rope_theta": model_config.get("rope_theta", 10000.0), - } - (output_dir / "config.json").write_text(json.dumps(config, indent=2), encoding="utf-8") - (output_dir / "generation_config.json").write_text( - json.dumps( - { - "max_new_tokens": 128, - "temperature": 0.7, - "top_k": 50, - "do_sample": True, - }, - indent=2, - ), - encoding="utf-8", - ) - (output_dir / "special_tokens_map.json").write_text( - json.dumps( - { - "bos_token": "", - "eos_token": "", - "unk_token": "", - "pad_token": "", - }, - indent=2, - ), - encoding="utf-8", - ) - (output_dir / "tokenizer_config.json").write_text( - json.dumps( - { - "tokenizer_class": "PreTrainedTokenizerFast", - "tokenizer_file": "tokenizer.json", - "bos_token": "", - "eos_token": "", - "unk_token": "", - "pad_token": "", - "model_max_length": model_config.get("context_length"), - "chat_template": DEFAULT_CHAT_TEMPLATE, - }, - indent=2, - ), - encoding="utf-8", - ) - (output_dir / "README.md").write_text(_hf_readme(config), encoding="utf-8") - return output_dir - - -def export_llama_adapter_package(project_dir: Path, output_dir: Optional[Path] = None) -> Path: - """Export a Llama-compatible state dict when the trained architecture matches Llama. - - This is a real tensor-name/layout adapter, not a relabelled MicroGPT - checkpoint. Classic-GPT models are rejected because their learned - positions, LayerNorm, or GELU MLP cannot be represented faithfully by - Llama/Qwen/Mistral loaders. - """ - - project_dir = Path(project_dir) - output_dir = Path(output_dir) if output_dir else project_dir / "llama_model" - checkpoint_path = project_dir / "final_model.pt" - tokenizer_path = project_dir / "tokenizer.json" - if not checkpoint_path.exists() or not tokenizer_path.exists(): - raise ExportError("Llama export requires final_model.pt and tokenizer.json.") - checkpoint = torch.load(checkpoint_path, map_location="cpu") - config = checkpoint.get("model_config") - state = checkpoint.get("model_state_dict") - if not isinstance(config, dict) or not isinstance(state, dict): - raise ExportError("Checkpoint must contain model_config and model_state_dict.") - required = {"norm_type": "rmsnorm", "position_encoding": "rope", "mlp_type": "swiglu"} - incompatible = [f"{key}={config.get(key)!r}" for key, expected in required.items() if config.get(key) != expected] - if config.get("bias", True): - incompatible.append("bias=True") - if int(config.get("attention_window", 0) or 0) != 0: - incompatible.append("attention_window is enabled") - if incompatible: - raise ExportError( - "Llama adapter requires the Llama-like architecture (RoPE, RMSNorm, SwiGLU, no bias, full attention). " - "This checkpoint is incompatible: " + ", ".join(incompatible) - ) - hidden = int(config["embedding_size"]) - heads = int(config["head_count"]) - kv_heads = int(config.get("kv_head_count") or heads) - if config.get("attention_type") == "mqa": - kv_heads = 1 - elif config.get("attention_type") == "gqa" and not config.get("kv_head_count"): - kv_heads = max(1, heads // 2) - head_dim = hidden // heads - kv_hidden = kv_heads * head_dim - adapted: dict[str, torch.Tensor] = { - "model.embed_tokens.weight": state["token_embedding.weight"], - "model.norm.weight": state["ln_f.weight"], - "lm_head.weight": state["lm_head.weight"], - } - for layer in range(int(config["layer_count"])): - source = f"blocks.{layer}" - target = f"model.layers.{layer}" - qkv = state[f"{source}.attn.c_attn.weight"] - q, k, v = qkv.split((hidden, kv_hidden, kv_hidden), dim=0) - adapted.update({ - f"{target}.input_layernorm.weight": state[f"{source}.ln_1.weight"], - f"{target}.self_attn.q_proj.weight": q, - f"{target}.self_attn.k_proj.weight": k, - f"{target}.self_attn.v_proj.weight": v, - f"{target}.self_attn.o_proj.weight": state[f"{source}.attn.c_proj.weight"], - f"{target}.post_attention_layernorm.weight": state[f"{source}.ln_2.weight"], - f"{target}.mlp.gate_proj.weight": state[f"{source}.mlp.w1.weight"], - f"{target}.mlp.down_proj.weight": state[f"{source}.mlp.w2.weight"], - f"{target}.mlp.up_proj.weight": state[f"{source}.mlp.w3.weight"], - }) - output_dir.mkdir(parents=True, exist_ok=True) - torch.save(adapted, output_dir / "pytorch_model.bin") - shutil.copy2(tokenizer_path, output_dir / "tokenizer.json") - for name in ("tokenizer_config.json", "special_tokens_map.json"): - source = project_dir / name - if source.exists(): - shutil.copy2(source, output_dir / name) - llama_config = { - "model_type": "llama", "architectures": ["LlamaForCausalLM"], - "vocab_size": int(config["vocab_size"]), "hidden_size": hidden, - "intermediate_size": hidden * 4, "num_hidden_layers": int(config["layer_count"]), - "num_attention_heads": heads, "num_key_value_heads": kv_heads, - "max_position_embeddings": int(config["context_length"]), - "rope_theta": float(config.get("rope_theta", 10000.0)), - "rms_norm_eps": 1e-5, "hidden_act": "silu", "tie_word_embeddings": True, - "bos_token_id": 2, "eos_token_id": 3, "pad_token_id": 0, - } - (output_dir / "config.json").write_text(json.dumps(llama_config, indent=2), encoding="utf-8") - (output_dir / "README.md").write_text( - "# Llama-compatible MicroGPT export\n\n" - "Weights were structurally adapted from a RoPE/RMSNorm/SwiGLU MicroGPT checkpoint. " - "Load with Transformers `LlamaForCausalLM` or compatible Llama harnesses.\n", - encoding="utf-8", - ) - return output_dir - - -def _hf_readme(config: dict[str, object]) -> str: - """Create README text for a MicroGPT HF-style package. - - Args: - config: Exported config dictionary. - - Returns: - README Markdown. - """ - - return ( - "# MicroGPT HF-Style Package\n\n" - "This folder was exported by Micro LLM Creator. It uses a Hugging Face-style " - "layout (`config.json`, `pytorch_model.bin`, `tokenizer.json`) for portability, " - "but `model_type` is `microgpt` and it is not directly convertible by llama.cpp as " - "a Llama/Mistral/Gemma model.\n\n" - "Load it with Micro LLM Creator's MicroGPT code, or build a custom Transformers " - "model class that understands this config and tensor naming.\n\n" - f"- Block style: {config.get('norm_type')} / {config.get('position_encoding')} / {config.get('mlp_type')}\n" - f"- Context length: {config.get('n_ctx')}\n" - f"- Embedding size: {config.get('n_embd')}\n" - f"- Layers: {config.get('n_layer')}\n" - f"- Heads: {config.get('n_head')}\n" - ) - - -def find_llama_cpp_converter(llama_cpp_dir: Path) -> Path: - """Find the llama.cpp Hugging Face to GGUF converter. - - Args: - llama_cpp_dir: llama.cpp checkout folder or direct converter path. - - Returns: - Path to the converter script. - - Raises: - ExportError: If no converter is found. - """ - - llama_cpp_dir = Path(llama_cpp_dir) - if llama_cpp_dir.is_file() and llama_cpp_dir.name == "convert_hf_to_gguf.py": - return llama_cpp_dir - if not str(llama_cpp_dir).strip() or str(llama_cpp_dir) == ".": - raise ExportError( - "Choose your local llama.cpp folder first. It should contain convert_hf_to_gguf.py." - ) - if not llama_cpp_dir.exists(): - raise ExportError(f"llama.cpp path does not exist: {llama_cpp_dir}") - if not llama_cpp_dir.is_dir(): - raise ExportError(f"llama.cpp path is not a folder: {llama_cpp_dir}") - - candidates = [ - llama_cpp_dir / "convert_hf_to_gguf.py", - llama_cpp_dir / "convert" / "convert_hf_to_gguf.py", - llama_cpp_dir / "examples" / "convert_hf_to_gguf.py", - ] - for candidate in candidates: - if candidate.exists(): - return candidate - - matches = sorted(llama_cpp_dir.rglob("convert_hf_to_gguf.py")) - if matches: - return matches[0] - - searched = "\n".join(f"- {candidate}" for candidate in candidates) - raise ExportError( - "Could not find llama.cpp converter script.\n" - "Expected a recent llama.cpp checkout containing convert_hf_to_gguf.py.\n" - f"Searched:\n{searched}" - ) - - -def export_gguf_with_llama_cpp(project_dir: Path, llama_cpp_dir: Path, output_path: Path, outtype: str = "f16") -> Path: - """Export a Hugging Face-compatible model through llama.cpp. - - Args: - project_dir: Model project containing an ``hf_model`` folder. - llama_cpp_dir: Local llama.cpp checkout folder. - output_path: Destination GGUF file path. - outtype: llama.cpp converter output type, usually f16 or f32. - - Returns: - GGUF output path. - - Raises: - ExportError: If converter or HF model folder is missing. - """ - converter = find_llama_cpp_converter(llama_cpp_dir) - - hf_dir = project_dir / "hf_model" - if not hf_dir.exists(): - raise ExportError( - "GGUF export needs model_core/hf_model. Use Export HF Package first, " - "but note MicroGPT packages are not llama.cpp-convertible unless llama.cpp " - "has a matching MicroGPT converter/model implementation." - ) - config_path = hf_dir / "config.json" - if config_path.exists(): - config = json.loads(config_path.read_text(encoding="utf-8")) - if config.get("model_type") == "microgpt": - raise ExportError( - "This hf_model folder is a MicroGPT package, not a llama.cpp-supported " - "Llama/Mistral/Gemma model. Real GGUF export needs a supported HF model " - "architecture or a custom llama.cpp MicroGPT converter." - ) - - output_path.parent.mkdir(parents=True, exist_ok=True) - if outtype not in {"f16", "f32", "bf16", "q8_0"}: - raise ExportError(f"Unsupported GGUF outtype for converter: {outtype}") - - subprocess.run( - ["python", str(converter), str(hf_dir), "--outfile", str(output_path), "--outtype", outtype], - check=True, - ) - return output_path diff --git a/llm_trainer/external_dataset.py b/llm_trainer/external_dataset.py deleted file mode 100644 index 8f252fd..0000000 --- a/llm_trainer/external_dataset.py +++ /dev/null @@ -1,304 +0,0 @@ -"""Download and manage datasets published as GitHub release assets.""" - -from __future__ import annotations - -import hashlib -import json -import logging -import shutil -import tempfile -import time -import urllib.request -import urllib.error -import zipfile -from dataclasses import dataclass -from pathlib import Path -from typing import Callable, Optional - - -LOGGER = logging.getLogger(__name__) -DEFAULT_MANIFEST_URL = ( - "https://github.com/drunkenbot-ai/dataset/releases/latest/download/manifest.json" -) - - -@dataclass(frozen=True) -class DatasetCategory: - """Metadata for one downloadable dataset category.""" - - name: str - archive: str - size_bytes: int - file_count: int - sha256: str - - -@dataclass(frozen=True) -class DatasetManifest: - """Validated metadata published for an external dataset.""" - - dataset_id: str - version: str - categories: tuple[DatasetCategory, ...] - - @classmethod - def from_json(cls, payload: object) -> "DatasetManifest": - """Build a manifest from decoded JSON. - - Args: - payload: Decoded manifest object. - - Returns: - Validated dataset manifest. - - Raises: - ValueError: If required fields are missing or invalid. - """ - if not isinstance(payload, dict): - raise ValueError("Dataset manifest must be a JSON object") - dataset_id = payload.get("dataset_id") - version = payload.get("version") - raw_categories = payload.get("categories") - if not isinstance(dataset_id, str) or not dataset_id.strip(): - raise ValueError("Dataset manifest has no dataset_id") - if not isinstance(version, str) or not version.strip(): - raise ValueError("Dataset manifest has no version") - if not isinstance(raw_categories, list): - raise ValueError("Dataset manifest categories must be a list") - - categories = [] - seen_names: set[str] = set() - for raw in raw_categories: - if not isinstance(raw, dict): - raise ValueError("Dataset category must be an object") - values = {key: raw.get(key) for key in ("name", "archive", "sha256")} - if not all(isinstance(value, str) and value.strip() for value in values.values()): - raise ValueError("Dataset category has invalid name, archive, or sha256") - name = values["name"] - if name in seen_names: - raise ValueError(f"Duplicate dataset category: {name}") - if Path(values["archive"]).name != values["archive"]: - raise ValueError(f"Unsafe archive name: {values['archive']}") - if not isinstance(raw.get("size_bytes"), int) or raw["size_bytes"] < 0: - raise ValueError(f"Invalid size for dataset category: {name}") - if not isinstance(raw.get("file_count"), int) or raw["file_count"] < 0: - raise ValueError(f"Invalid file count for dataset category: {name}") - if len(values["sha256"]) != 64: - raise ValueError(f"Invalid SHA-256 for dataset category: {name}") - categories.append(DatasetCategory( - name=name, - archive=values["archive"], - size_bytes=raw["size_bytes"], - file_count=raw["file_count"], - sha256=values["sha256"].lower(), - )) - seen_names.add(name) - return cls(dataset_id=dataset_id, version=version.strip(), categories=tuple(categories)) - - -def parse_version(version: str) -> tuple[int, ...]: - """Parse a numeric dotted dataset version. - - Args: - version: Version such as ``"2.0.0"``. - - Returns: - Numeric version components. - - Raises: - ValueError: If the version is not numeric and dotted. - """ - parts = version.strip().split(".") - if not parts or any(not part.isdigit() for part in parts): - raise ValueError(f"Invalid dataset version: {version!r}") - return tuple(int(part) for part in parts) - - -def is_newer_version(remote: str, local: Optional[str]) -> bool: - """Return whether a remote version is newer than a local version. - - Args: - remote: Remote dataset version. - local: Installed dataset version, or ``None``. - - Returns: - True when the remote version is greater than the local version. - """ - return local is None or parse_version(remote) > parse_version(local) - - -def _download(url: str, destination: Path, progress: Optional[Callable[[int, int], None]]) -> None: - """Download a URL to a file while reporting byte progress.""" - request = urllib.request.Request(url, headers={"User-Agent": "DrunkenBot-LLM-IDE"}) - for attempt in range(3): - try: - with urllib.request.urlopen(request, timeout=45) as response, destination.open("wb") as output: - total = int(response.headers.get("Content-Length") or 0) - downloaded = 0 - while True: - block = response.read(1024 * 1024) - if not block: - break - output.write(block) - downloaded += len(block) - if progress: - progress(downloaded, total) - return - except (TimeoutError, OSError, urllib.error.URLError): - if attempt == 2: - raise - LOGGER.warning("Download attempt %d failed for %s; retrying", attempt + 1, url) - time.sleep(2 ** attempt) - - -def _sha256_file(path: Path) -> str: - """Calculate a file checksum without loading it all into memory.""" - digest = hashlib.sha256() - with path.open("rb") as stream: - for block in iter(lambda: stream.read(1024 * 1024), b""): - digest.update(block) - return digest.hexdigest() - - -def _extract_archive(archive: Path, destination: Path) -> None: - """Extract an archive while rejecting paths outside the staging folder.""" - with zipfile.ZipFile(archive) as package: - root = destination.resolve() - for member in package.infolist(): - target = (destination / member.filename).resolve() - if target != root and root not in target.parents: - raise OSError(f"Unsafe path in dataset archive: {member.filename}") - package.extractall(destination) - - -def load_manifest(url: str = DEFAULT_MANIFEST_URL) -> DatasetManifest: - """Download and validate the current remote dataset manifest. - - Args: - url: Manifest URL. - - Returns: - Validated remote manifest. - """ - with urllib.request.urlopen( - urllib.request.Request(url, headers={"User-Agent": "DrunkenBot-LLM-IDE"}), - timeout=60, - ) as response: - payload = json.loads(response.read().decode("utf-8")) - manifest = DatasetManifest.from_json(payload) - LOGGER.info("Loaded dataset manifest %s version %s", manifest.dataset_id, manifest.version) - return manifest - - -def install_categories( - manifest: DatasetManifest, - destination: Path, - categories: Optional[list[str]] = None, - progress: Optional[Callable[[int, int], None]] = None, - manifest_url: str = DEFAULT_MANIFEST_URL, -) -> Path: - """Download, verify, and atomically install selected dataset categories. - - Args: - manifest: Validated dataset manifest to install. - destination: Managed dataset root. - categories: Category names to install, or all categories when omitted. - progress: Optional callback receiving downloaded and total bytes. - manifest_url: URL used to resolve release asset URLs. - - Returns: - Destination dataset root. - - Raises: - ValueError: If a requested category is not in the manifest. - OSError: If download, extraction, or replacement fails. - zipfile.BadZipFile: If an archive is invalid. - """ - selected_names = set(categories) if categories is not None else { - category.name for category in manifest.categories if category.file_count > 0 - } - by_name = {category.name: category for category in manifest.categories} - unknown = selected_names - by_name.keys() - if unknown: - raise ValueError(f"Unknown dataset categories: {', '.join(sorted(unknown))}") - release_root = manifest_url.rsplit("/", 1)[0] - destination = Path(destination) - destination.parent.mkdir(parents=True, exist_ok=True) - - with tempfile.TemporaryDirectory(prefix="dataset-install-", dir=destination.parent) as temp_name: - staging = Path(temp_name) / "dataset" - staging.mkdir() - if destination.is_dir(): - shutil.copytree(destination, staging, dirs_exist_ok=True) - ordered_names = sorted(selected_names) - for category_index, name in enumerate(ordered_names, start=1): - category = by_name[name] - archive = staging / category.archive - LOGGER.info("Downloading dataset category %s", name) - archive_progress = ( - (lambda downloaded, total: progress({ - "message": f"Downloaded {category.name}: {downloaded} bytes", - "button_text": f"Downloading {category.archive} dataset ({category_index} of {len(ordered_names)})", - "percent": int(downloaded / total * 100) if total else 0, - })) - if progress - else None - ) - _download(f"{release_root}/{category.archive}", archive, archive_progress) - if archive.stat().st_size != category.size_bytes: - raise OSError(f"Size mismatch for dataset category {name}") - digest = _sha256_file(archive) - if digest != category.sha256: - raise OSError(f"SHA-256 mismatch for dataset category {name}") - _extract_archive(archive, staging) - archive.unlink() - - (staging / "version.txt").write_text(manifest.version + "\n", encoding="utf-8") - (staging / "manifest.json").write_text( - json.dumps({ - "dataset_id": manifest.dataset_id, - "version": manifest.version, - "categories": [category.__dict__ for category in manifest.categories], - }, indent=2) + "\n", - encoding="utf-8", - ) - backup = destination.with_name(destination.name + ".previous") - if backup.exists(): - shutil.rmtree(backup) - if destination.exists(): - destination.replace(backup) - staging.replace(destination) - if backup.exists(): - shutil.rmtree(backup) - LOGGER.info("Installed dataset version %s at %s", manifest.version, destination) - return destination - - -def download_latest_dataset( - destination: Path, - categories: Optional[list[str]] = None, - progress: Optional[Callable[[int, int], None]] = None, - should_stop: Optional[Callable[[], bool]] = None, - manifest_url: str = DEFAULT_MANIFEST_URL, -) -> DatasetManifest: - """Download and install all categories from the latest release. - - Args: - destination: Managed dataset root. - progress: Optional callback receiving downloaded and total bytes. - should_stop: Optional cooperative cancellation callback. - - Returns: - The installed release manifest. - """ - if should_stop and should_stop(): - raise RuntimeError("Dataset download stopped by user.") - manifest = load_manifest(manifest_url) - install_categories( - manifest, - destination, - categories=categories, - progress=progress, - manifest_url=manifest_url, - ) - return manifest diff --git a/llm_trainer/fine_tuning_service.py b/llm_trainer/fine_tuning_service.py deleted file mode 100644 index c0211e5..0000000 --- a/llm_trainer/fine_tuning_service.py +++ /dev/null @@ -1,119 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from pathlib import Path -from typing import Any, Callable, Optional - -from .config import ModelConfig, TrainingConfig -from .training import TrainingResult -from .training_service import LocalTrainerService, TrainingJobRequest, TrainingService - - -ProgressCallback = Callable[[Any], None] -StopCallback = Callable[[], bool] - - -@dataclass -class FineTuningJobRequest: - """Request payload for a fine-tuning job. - - Args: - dataset_dir: Prepared instruction/conversation/domain dataset folder. - model_config: Compatible model architecture settings. - training_config: Fine-tuning optimizer/runtime settings. - stage: Fine-tuning stage label: instruction, conversation, or domain. - metadata: Optional metadata persisted with the training job. - """ - - dataset_dir: Path - model_config: ModelConfig - training_config: TrainingConfig - stage: str = "domain" - metadata: Optional[dict[str, Any]] = None - - def to_training_request(self) -> TrainingJobRequest: - """Convert to the generic training service request. - - Returns: - Generic training request. - """ - - self.training_config.training_mode = "fine_tune" - metadata = dict(self.metadata or {}) - metadata["fine_tune_stage"] = self.stage - return TrainingJobRequest( - dataset_dir=self.dataset_dir, - model_config=self.model_config, - training_config=self.training_config, - metadata=metadata, - ) - - -class FineTuningService: - """Fine-tuning API boundary used by the desktop UI. - - The service delegates to the generic training service but keeps - fine-tuning orchestration separate from GUI code. - """ - - def __init__(self, training_service: Optional[TrainingService] = None) -> None: - """Create the fine-tuning service. - - Args: - training_service: Optional training service implementation. - """ - - self.training_service = training_service or LocalTrainerService() - - def run( - self, - request: FineTuningJobRequest, - progress: Optional[ProgressCallback] = None, - should_stop: Optional[StopCallback] = None, - ) -> TrainingResult: - """Run a fine-tuning request. - - Args: - request: Fine-tuning request. - progress: Optional progress callback. - should_stop: Optional cooperative cancellation callback. - - Returns: - Training result. - """ - - return self.training_service.run( - request.to_training_request(), - progress=progress, - should_stop=should_stop, - ) - - -def run_fine_tuning_job( - dataset_dir: Path, - model_config: ModelConfig, - training_config: TrainingConfig, - stage: str = "domain", - progress: Optional[ProgressCallback] = None, - should_stop: Optional[StopCallback] = None, -) -> TrainingResult: - """Run a fine-tuning job through the fine-tuning service. - - Args: - dataset_dir: Prepared fine-tuning dataset directory. - model_config: Compatible model architecture settings. - training_config: Fine-tuning settings. - stage: Fine-tuning stage label. - progress: Optional progress callback. - should_stop: Optional cooperative cancellation callback. - - Returns: - Training result. - """ - - service = FineTuningService() - return service.run( - FineTuningJobRequest(dataset_dir, model_config, training_config, stage=stage), - progress=progress, - should_stop=should_stop, - ) diff --git a/llm_trainer/generate_identity_data.py b/llm_trainer/generate_identity_data.py deleted file mode 100644 index 0ea6e4c..0000000 --- a/llm_trainer/generate_identity_data.py +++ /dev/null @@ -1,328 +0,0 @@ -"""Generate a base-pretraining identity corpus file for a Micro LLM project. - -Reads the model's name and creation date straight from the project's -``project.json`` (written by the app itself), combines them with fixed facts -about the creator, and writes a plain-text file of many distinctly-phrased -sentences about who the model is. Designed to be dropped straight into a -project's ``training_data`` folder as an ordinary source file. - -Why sentences must be genuinely distinct, not just recombined: - The dataset-build pipeline drops documents whose sentences are - dominated by exact repetition (see ``dataset_mixture.MAX_REPETITIVE_UNIT_RATIO``, - checked per document against the *sentence* level, not the paragraph - level) and removes exact duplicate documents outright. Shuffling a small - fixed sentence pool into different paragraph groupings does NOT create - new sentences from the filter's point of view -- it still sees the same - handful of sentences repeated over and over and will exclude the whole - file. This script instead builds sentences combinatorially (varying verb - choice, phrase order, and which facts are mentioned) so the pool of - distinct sentences is large, and writes each one only once. - -Why file size does not need to be huge: - Training runs for multiple epochs (see ``TrainingConfig.epochs``), so - every sentence in this file is seen again on every epoch automatically. - A few hundred distinct sentences already gives real, repeated exposure - across a training run without this file dominating a multi-hundred-MB - corpus. - -Usage: - python generate_identity_data.py /path/to/project_dir \ - [--output training_data/identity/identity_facts.txt] \ - [--sentence-count 500] \ - [--creator "DrunkenBot"] \ - [--maker "Nilesh Jadhav"] \ - [--role "AI assistant"] - -If ``--output`` is a relative path, it is resolved against the project -directory. The default output path lands under the project's -``training_data`` folder, in its own subfolder, so it is picked up by the -normal source-vault scan alongside every other category. -""" - -from __future__ import annotations - -import argparse -import itertools -import json -import re -import random -from datetime import datetime -from pathlib import Path - -# Mirrors dataset_mixture.py's thresholds exactly, so this script's -# self-check reflects what the real pipeline will do with this file. -MAX_REPETITIVE_UNIT_RATIO = 0.35 -MIN_REPETITION_CHECK_UNITS = 20 -MIN_REPETITION_CHECK_CHARS = 2_000 - - -def load_project_facts(project_dir: Path) -> dict[str, str]: - """Read the project name and creation date from project.json. - - Args: - project_dir: Project folder containing ``project.json``. - - Returns: - Dict with ``name`` and ``created`` (a human-readable date string). - - Raises: - FileNotFoundError: If ``project.json`` does not exist. - ValueError: If ``project.json`` does not contain a project name. - """ - - project_file = project_dir / "project.json" - if not project_file.exists(): - raise FileNotFoundError( - f"Could not find {project_file}. Pass the folder that contains your project's project.json." - ) - data = json.loads(project_file.read_text(encoding="utf-8")) - name = str(data.get("project_name", "")).strip() - if not name: - raise ValueError(f"{project_file} has no project_name set.") - - raw_created = str(data.get("created_at") or data.get("saved_at") or "").strip() - created = _format_date(raw_created) - return {"name": name, "created": created} - - -def _format_date(raw: str) -> str: - """Format an ISO-ish timestamp into a human-readable month/year. - - Args: - raw: Timestamp string, typically ``datetime.isoformat()`` output. - - Returns: - A "Month YYYY" string, or "an unknown date" if parsing fails. - """ - - if not raw: - return "an unknown date" - try: - return datetime.fromisoformat(raw).strftime("%B %Y") - except ValueError: - return "an unknown date" - - -# Independent slot dimensions. itertools.product over these (crossed with -# several sentence "shapes" below) is what makes the combinatorial pool -# large without hand-writing hundreds of sentences. -CREATE_VERBS = ["created", "built", "made", "developed", "brought to life"] -ROLE_PHRASES = ["an {role}", "a helpful {role}", "an {role} you can talk to"] -CONNECTORS = ["", " -- a project by {maker}", ", led by {maker}"] -DATE_PHRASES = ["", " in {created}", ", first developed in {created}"] - - -def _render_shapes(values: dict[str, str]) -> set[str]: - """Render every template shape across every slot combination. - - Args: - values: Base fact values (``name``, ``creator``, ``maker``, ``role``, - ``created``). - - Returns: - Set of distinct rendered sentences. - """ - - sentences: set[str] = set() - for verb, role_phrase, connector, date_phrase in itertools.product( - CREATE_VERBS, ROLE_PHRASES, CONNECTORS, DATE_PHRASES - ): - role_text = role_phrase.format(role=values["role"]) - connector_text = connector.format(maker=values["maker"]) - date_text = date_phrase.format(created=values["created"]) - - sentences.add( - f"{values['name']} is {role_text} {verb} by {values['creator']}{connector_text}{date_text}.".replace( - " ", " " - ) - ) - sentences.add( - f"{values['creator']}{connector_text} {verb} {values['name']}, {role_text}{date_text}.".replace( - " ", " " - ) - ) - sentences.add( - f"{values['name']}, {role_text}, was {verb} by {values['creator']}{connector_text}{date_text}.".replace( - " ", " " - ) - ) - - # A modest number of naturally-phrased question/answer lines, written as - # prose rather than a structured instruction format (this file stays - # inside the base-pretraining corpus, not a fine-tuning dataset). - qa_lines = [ - "What is your name? {name}.", - "Who made you? {creator}, a project by {maker}, created me.", - "Who created you? I was created by {creator}.", - "Are you an AI? Yes, {name} is an {role}.", - "When were you created? Around {created}.", - "Who is your creator? {creator}, founded by {maker}.", - "Do you belong to another company? No, {name} was made by {creator}, not any other company.", - ] - sentences.update(line.format(**values) for line in qa_lines) - return sentences - - -def build_sentence_pool( - facts: dict[str, str], - creator: str, - maker: str, - role: str, - target_count: int, - seed: int, -) -> list[str]: - """Build a pool of distinct sentences, sized to ``target_count``. - - Args: - facts: Project facts (``name``, ``created``). - creator: Creator/organization name. - maker: Person who made the creator. - role: What the model is (e.g. "AI assistant"). - target_count: Desired number of sentences. - seed: Random seed used when sampling down to ``target_count``. - - Returns: - List of distinct sentences, shuffled, at most ``target_count`` long. - """ - - values = { - "name": facts["name"], - "created": facts["created"], - "creator": creator, - "maker": maker, - "role": role, - } - all_sentences = sorted(_render_shapes(values)) - rng = random.Random(seed) - rng.shuffle(all_sentences) - if target_count >= len(all_sentences): - print( - f"Note: requested {target_count} sentences, but the combinatorial " - f"template pool only has {len(all_sentences)} distinct options. " - "Using all of them. Add more verbs/phrases/connectors to the " - "CREATE_VERBS / ROLE_PHRASES / CONNECTORS / DATE_PHRASES lists " - "above to raise this ceiling." - ) - return all_sentences - return all_sentences[:target_count] - - -def group_into_paragraphs(sentences: list[str], sentences_per_paragraph: tuple[int, int], seed: int) -> list[str]: - """Group sentences into paragraphs, each sentence used exactly once. - - Args: - sentences: Distinct sentences to group (already shuffled/ordered). - sentences_per_paragraph: Inclusive (min, max) sentence count range - per generated paragraph. - seed: Random seed for paragraph-size choices. - - Returns: - List of paragraph strings covering every input sentence exactly once. - """ - - rng = random.Random(seed) - low, high = sentences_per_paragraph - paragraphs: list[str] = [] - index = 0 - while index < len(sentences): - size = rng.randint(low, high) - chunk = sentences[index : index + size] - paragraphs.append(" ".join(chunk)) - index += size - return paragraphs - - -def _canonical_block(text: str) -> str: - """Match dataset_mixture._canonical_corpus_block exactly.""" - - return re.sub(r"\s+", " ", text).strip().lower() - - -def self_check_diversity(full_text: str) -> tuple[int, float, bool]: - """Reproduce the pipeline's low-diversity check on the generated text. - - Args: - full_text: The full generated file content. - - Returns: - Tuple of (unit_count, duplicate_ratio, would_be_excluded). - """ - - raw_units = re.split( - r"(?<=[.!?])\s+|\n+(?=(?:User|Assistant|System|Instruction|Response):)", full_text - ) - units = [_canonical_block(unit) for unit in raw_units if len(_canonical_block(unit)) >= 24] - if len(full_text) < MIN_REPETITION_CHECK_CHARS or len(units) < MIN_REPETITION_CHECK_UNITS: - return len(units), 0.0, False - duplicate_ratio = 1.0 - (len(set(units)) / len(units)) - return len(units), duplicate_ratio, duplicate_ratio > MAX_REPETITIVE_UNIT_RATIO - - -def main() -> None: - """Parse arguments, generate the identity corpus, and write it to disk.""" - - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("project_dir", type=Path, help="Project folder containing project.json") - parser.add_argument( - "--output", - type=Path, - default=Path("training_data/identity/identity_facts.txt"), - help="Output path, resolved against project_dir if relative (default: %(default)s)", - ) - parser.add_argument("--creator", default="DrunkenBot", help="Creator/organization name") - parser.add_argument("--maker", default="Nilesh Jadhav", help="Person who made the creator") - parser.add_argument("--role", default="AI assistant", help="What the model is") - parser.add_argument( - "--sentence-count", - type=int, - default=500, - help="Target number of distinct sentences (default: %(default)s)", - ) - parser.add_argument("--seed", type=int, default=1337, help="Random seed for reproducibility") - args = parser.parse_args() - - project_dir = args.project_dir.resolve() - facts = load_project_facts(project_dir) - sentence_pool = build_sentence_pool( - facts, args.creator, args.maker, args.role, args.sentence_count, args.seed - ) - paragraphs = group_into_paragraphs(sentence_pool, sentences_per_paragraph=(2, 4), seed=args.seed) - full_text = "\n\n".join(paragraphs) + "\n" - - unit_count, duplicate_ratio, would_be_excluded = self_check_diversity(full_text) - - output_path = args.output if args.output.is_absolute() else project_dir / args.output - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_text(full_text, encoding="utf-8") - - print(f"Project name: {facts['name']}") - print(f"Created: {facts['created']}") - print(f"Creator / maker: {args.creator} / {args.maker}") - print(f"Distinct sentences: {len(sentence_pool)}") - print(f"Paragraphs written: {len(paragraphs)}") - print(f"Output size: {len(full_text):,} characters") - print(f"Written to: {output_path}") - print() - print("Diversity self-check (same rule the pipeline applies):") - print(f" sentence units: {unit_count}") - print(f" duplicate ratio: {duplicate_ratio:.1%} (excluded if over 35%)") - print(f" would be excluded: {would_be_excluded}") - if would_be_excluded: - print( - " WARNING: this file would be dropped by the low-diversity filter. " - "Increase --sentence-count so the distinct-sentence pool covers more of the file." - ) - print() - print("Next steps:") - print(" 1. Re-run dataset preparation (this is a new file, so it will") - print(" be extracted automatically -- no need for force reprocess).") - print(" 2. Train for multiple epochs; this file is re-seen every epoch,") - print(" which is what gives a small model repeated exposure to it.") - print(" 3. Recall from a small from-scratch model is probabilistic, not") - print(" guaranteed -- if you need reliable identity answers, pair") - print(" this with a system-prompt/prepended-context mechanism at") - print(" inference time if your chat interface supports one.") - - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/llm_trainer/generation.py b/llm_trainer/generation.py deleted file mode 100644 index 0765589..0000000 --- a/llm_trainer/generation.py +++ /dev/null @@ -1,72 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -from typing import Optional - -import torch - -from .config import ModelConfig -from .model import MicroGPT -from .tokenizer import EOS_TOKEN, load_tokenizer, token_id - - -def load_model_from_checkpoint(checkpoint_path: Path, device: Optional[str] = None) -> MicroGPT: - """Load a trained MicroGPT checkpoint. - - Args: - checkpoint_path: Path to a saved model checkpoint. - device: Optional device override. - - Returns: - Loaded model in evaluation mode. - """ - - device = device or ("cuda" if torch.cuda.is_available() else "cpu") - checkpoint = torch.load(checkpoint_path, map_location=device) - config = ModelConfig(**checkpoint["model_config"]) - model = MicroGPT(config) - model.load_state_dict(checkpoint["model_state_dict"]) - model.to(device) - model.eval() - return model - - -@torch.no_grad() -def generate_text( - checkpoint_path: Path, - tokenizer_path: Path, - prompt: str, - max_new_tokens: int = 100, - temperature: float = 0.8, - top_k: Optional[int] = 50, - device: Optional[str] = None, - use_kv_cache: bool = True, -) -> str: - """Generate text from a trained checkpoint. - - Args: - checkpoint_path: Path to model checkpoint. - tokenizer_path: Path to tokenizer JSON. - prompt: Prompt text. - max_new_tokens: Maximum tokens to sample. - temperature: Sampling temperature. - top_k: Optional top-k sampling cutoff. - device: Optional device override. - use_kv_cache: Whether to use key/value cache during generation. - - Returns: - Decoded generated text. - """ - - device = device or ("cuda" if torch.cuda.is_available() else "cpu") - tokenizer = load_tokenizer(tokenizer_path) - model = load_model_from_checkpoint(checkpoint_path, device=device) - input_ids = tokenizer.encode(prompt).ids - context = torch.tensor([input_ids], dtype=torch.long, device=device) - generated = model.generate(context, max_new_tokens, temperature=temperature, top_k=top_k, use_kv_cache=use_kv_cache) - eos_id = token_id(tokenizer, EOS_TOKEN) - output_ids = generated[0].tolist() - if eos_id in output_ids[len(input_ids) :]: - eos_index = output_ids.index(eos_id, len(input_ids)) - output_ids = output_ids[:eos_index] - return tokenizer.decode(output_ids) diff --git a/llm_trainer/license_client.py b/llm_trainer/license_client.py deleted file mode 100644 index f44a876..0000000 --- a/llm_trainer/license_client.py +++ /dev/null @@ -1,391 +0,0 @@ -"""IDE license validation: online-first, with a signed offline grace cache. - -Calls the DrunkenBot cloud service's ``POST /license/validate`` at launch. -On success, caches a short-lived signed "grace receipt" locally so the app -can still launch for a limited window if the server is unreachable next -time (see ``LICENSE_SERVER_URL`` and the cloud-service README for the full -design). This module never trusts anything it did not itself verify with -:data:`LICENSE_PUBLIC_KEY_PEM` -- an unsigned or badly-signed receipt is -treated exactly like "no receipt at all." - -IMPORTANT: :data:`LICENSE_PUBLIC_KEY_PEM` below is a placeholder generated -for scaffolding only. Before this ships, replace it with the real public -key printed by ``cloud-service/scripts/generate_keypair.py`` when the -production signing keypair is generated. Shipping the placeholder means -this app will only ever trust receipts signed by a throwaway key nobody -else has -- i.e. license validation will never succeed for real customers. -""" - -from __future__ import annotations - -import json -import platform -import time -import urllib.error -import urllib.request -import uuid -from dataclasses import dataclass -from datetime import datetime, timezone -from pathlib import Path -from typing import Optional -import traceback -import ssl -import certifi - -from cryptography.exceptions import InvalidSignature -from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey -from cryptography.hazmat.primitives.serialization import load_pem_public_key - -# --- REPLACE BEFORE SHIPPING: see module docstring. --- -LICENSE_PUBLIC_KEY_PEM = """-----BEGIN PUBLIC KEY----- -MCowBQYDK2VwAyEAledx+Yhz/kvDTBFfBscicAMUIcwwG2jI2/zrwK6VFwI= ------END PUBLIC KEY----- -""" - -LICENSE_DIR = Path.home() / ".drunkenbot_ide" / "license" -LICENSE_KEY_FILE = LICENSE_DIR / "license_key.txt" -GRACE_CACHE_FILE = LICENSE_DIR / "grace_receipt.json" -MACHINE_ID_FILE = LICENSE_DIR / "machine_id.txt" - -_REQUEST_TIMEOUT_SECONDS = 15.0 -_ONLINE_VALIDATION_ATTEMPTS = 5 - - -@dataclass -class LicenseCheckResult: - """Outcome of a license check. - - Attributes: - valid: Whether the app is licensed to launch. - reason: Human-readable explanation, especially when not valid. - used_offline_grace: Whether this result came from a cached grace - receipt rather than a live server response. - version_ceiling: Highest app version the license currently covers, - when known. - grace_period_until: ISO timestamp of any temporary grace extension, - when known. - """ - - valid: bool - reason: str - used_offline_grace: bool = False - version_ceiling: Optional[str] = None - grace_period_until: Optional[str] = None - - -def _get_or_create_machine_id() -> str: - """Return this machine's pseudonymous ID, generating one if needed. - - Deliberately a random value with no relationship to any real hardware - or OS identifier (not a disk serial, not a MAC address, not tied to a - Windows/OS username) -- it exists only to distinguish installs and spot - abuse patterns, not to fingerprint a specific physical device or person. - - Returns: - Stable machine ID, persisted locally after first generation. - """ - - if MACHINE_ID_FILE.exists(): - existing = MACHINE_ID_FILE.read_text(encoding="utf-8").strip() - if existing: - return existing - machine_id = uuid.uuid4().hex - LICENSE_DIR.mkdir(parents=True, exist_ok=True) - MACHINE_ID_FILE.write_text(machine_id, encoding="utf-8") - return machine_id - - -def _collect_telemetry() -> dict: - """Collect the minimal, privacy-conscious launch telemetry payload. - - Deliberately excludes anything directly identifying: no OS username, no - hostname, no hardware identifiers. ISP/geolocation is derived - server-side from the request's source IP, not reported here. - - Returns: - Telemetry dict matching the server's ``LaunchTelemetry`` schema. - """ - - return { - "machine_id": _get_or_create_machine_id(), - "os": platform.system() or None, - "os_version": platform.release() or None, - } - - -def _parse_version(version: str) -> tuple[int, int, int]: - """Parse a version string into a comparable tuple. - - Mirrors the cloud service's ``app/versioning.py`` exactly -- kept as an - independent copy rather than a shared import, since LLM-IDE and - cloud-service are separate applications/repos with no shared package. - - Args: - version: Version string, e.g. ``"2.1.0"``. - - Returns: - ``(major, minor, patch)`` tuple. Missing components default to 0. - """ - - core = version.strip().split("-", 1)[0].split("+", 1)[0] - parts = core.split(".") - numbers = [int(part) for part in parts[:3] if part.isdigit()] - while len(numbers) < 3: - numbers.append(0) - return numbers[0], numbers[1], numbers[2] - - -def _is_version_within_ceiling(app_version: str, version_ceiling: str) -> bool: - """Return whether an app version is covered by a license's ceiling. - - Args: - app_version: Version of the running app. - version_ceiling: Highest version the license entitles the holder to. - - Returns: - True if ``app_version <= version_ceiling``. - """ - - return _parse_version(app_version) <= _parse_version(version_ceiling) - - -def _load_public_key() -> Ed25519PublicKey: - """Load the embedded Ed25519 public key. - - Returns: - Loaded public key object. - """ - - key = load_pem_public_key(LICENSE_PUBLIC_KEY_PEM.encode("utf-8")) - if not isinstance(key, Ed25519PublicKey): - raise ValueError("Embedded license public key is not Ed25519.") - return key - - -def _verify_receipt(receipt: str, signature_b64: str) -> dict: - """Verify a signed receipt and return its parsed payload. - - Args: - receipt: Exact canonical JSON string that was signed. - signature_b64: Base64-encoded Ed25519 signature over ``receipt``. - - Returns: - Parsed receipt payload. - - Raises: - InvalidSignature: If the signature does not match. - ValueError: If the receipt is not valid JSON. - """ - - import base64 - - public_key = _load_public_key() - public_key.verify(base64.b64decode(signature_b64), receipt.encode("utf-8")) - return json.loads(receipt) - - -def load_stored_license_key() -> Optional[str]: - """Return the previously activated license key, if any. - - Returns: - Stored license key, or ``None`` if never activated. - """ - - if not LICENSE_KEY_FILE.exists(): - return None - key = LICENSE_KEY_FILE.read_text(encoding="utf-8").strip() - return key or None - - -def store_license_key(license_key: str) -> None: - """Persist an activated license key for future launches. - - Args: - license_key: License key entered during activation. - """ - - LICENSE_DIR.mkdir(parents=True, exist_ok=True) - LICENSE_KEY_FILE.write_text(license_key.strip(), encoding="utf-8") - - -def _load_cached_receipt() -> Optional[tuple[str, str]]: - """Load the last cached grace receipt and its signature, if present. - - Returns: - ``(receipt, signature)`` tuple, or ``None`` if no cache exists or it - is unreadable. - """ - - if not GRACE_CACHE_FILE.exists(): - return None - try: - cached = json.loads(GRACE_CACHE_FILE.read_text(encoding="utf-8")) - return cached["receipt"], cached["signature"] - except Exception: - return None - - -def _store_cached_receipt(receipt: str, signature: str) -> None: - """Cache a validated grace receipt and its signature to disk. - - Args: - receipt: Canonical receipt JSON string. - signature: Base64-encoded signature over ``receipt``. - """ - - LICENSE_DIR.mkdir(parents=True, exist_ok=True) - GRACE_CACHE_FILE.write_text(json.dumps({"receipt": receipt, "signature": signature}), encoding="utf-8") - - -def _clear_cached_receipt() -> None: - """Delete any cached grace receipt. - - Called whenever the server gives an explicit, live rejection (revoked, - version no longer covered, etc.) so that result cannot be bypassed on a - later launch by simply blocking network access and falling back to a - stale cached grace receipt that predates the rejection. - """ - - GRACE_CACHE_FILE.unlink(missing_ok=True) - - -def _validate_online(license_key: str, app_version: str, server_url: str) -> Optional[dict]: - """Call the cloud service's validation endpoint. - - Args: - license_key: License key to validate. - app_version: Running app's version. - server_url: Base URL of the DrunkenBot cloud service. - - Returns: - Parsed JSON response, or ``None`` if the server could not be - reached at all (caller should fall back to the offline grace - cache in that case). A reachable server that responds with - ``valid: false`` is NOT a network failure -- that is returned - normally so the caller treats it as an authoritative rejection. - """ - - body = json.dumps( - { - "license_key": license_key, - "app_version": app_version, - "telemetry": _collect_telemetry(), - } - ).encode("utf-8") - request = urllib.request.Request( - f"{server_url.rstrip('/')}/license/validate", - data=body, - headers={ - "Content-Type": "application/json", - "Accept": "application/json, text/plain, */*", - "User-Agent": "Mozilla/5.0", - }, - method="POST", - ) - context = ssl.create_default_context(cafile=certifi.where()) - for attempt in range(1, _ONLINE_VALIDATION_ATTEMPTS + 1): - try: - with urllib.request.urlopen( - request, - timeout=_REQUEST_TIMEOUT_SECONDS, - context=context, - ) as response: - return json.loads(response.read().decode("utf-8")) - except urllib.error.HTTPError as e: - print("HTTP Error:", e.code) - print(e.read().decode()) - return None - except (urllib.error.URLError, TimeoutError, OSError) as e: - print(f"License server attempt {attempt}/{_ONLINE_VALIDATION_ATTEMPTS} failed: {e}") - if attempt < _ONLINE_VALIDATION_ATTEMPTS: - time.sleep(min(2.0 * attempt, 8.0)) - except Exception: - traceback.print_exc() - return None - return None - - -def check_license_at_launch(app_version: str, server_url: str) -> LicenseCheckResult: - """Validate the license at app startup: online-first, offline-graceful. - - Order of operations: - 1. No stored license key at all -> not licensed, ask the user to - activate. - 2. Server reachable -> its answer is authoritative. A live "invalid" - response also clears any cached grace receipt (see - :func:`_clear_cached_receipt`), so a revoked license cannot be - revived later just by cutting network access. - 3. Server unreachable -> fall back to a cached grace receipt, if one - exists, is correctly signed, and has not expired. - - Args: - app_version: Version of the running app (compared against the - license's version ceiling). - server_url: Base URL of the DrunkenBot cloud service. - - Returns: - License check result. - """ - - license_key = load_stored_license_key() - if not license_key: - return LicenseCheckResult(valid=False, reason="No license activated on this machine.") - - response = _validate_online(license_key, app_version, server_url) - if response is not None: - if response.get("valid"): - receipt = response.get("receipt") - signature = response.get("signature") - if receipt and signature: - _store_cached_receipt(receipt, signature) - return LicenseCheckResult( - valid=True, - reason="Validated online.", - version_ceiling=response.get("version_ceiling"), - grace_period_until=response.get("grace_period_until"), - ) - # Authoritative, live rejection -- do not let a stale cache override this. - _clear_cached_receipt() - return LicenseCheckResult( - valid=False, - reason=response.get("reason", "License is not valid."), - version_ceiling=response.get("version_ceiling"), - grace_period_until=response.get("grace_period_until"), - ) - - # Server unreachable: fall back to a cached grace receipt, if any. - cached = _load_cached_receipt() - if cached is None: - return LicenseCheckResult( - valid=False, - reason="Could not reach the license server and no cached grace period is available. " - "Please connect to the internet once to validate your license.", - ) - receipt_json, signature = cached - try: - payload = _verify_receipt(receipt_json, signature) - except (InvalidSignature, ValueError, KeyError): - return LicenseCheckResult( - valid=False, - reason="Cached license grace data is corrupt or invalid. " - "Please connect to the internet once to re-validate your license.", - ) - - valid_until = datetime.fromisoformat(payload["valid_until"]) - if datetime.now(timezone.utc) > valid_until: - return LicenseCheckResult( - valid=False, - reason="Your offline grace period has expired. Please connect to the internet to re-validate.", - ) - if not _is_version_within_ceiling(app_version, payload["version_ceiling"]): - return LicenseCheckResult( - valid=False, - reason=f"This license covers up to version {payload['version_ceiling']}.", - version_ceiling=payload["version_ceiling"], - ) - return LicenseCheckResult( - valid=True, - reason="Validated from cached offline grace receipt.", - used_offline_grace=True, - version_ceiling=payload["version_ceiling"], - ) diff --git a/llm_trainer/lineage.py b/llm_trainer/lineage.py deleted file mode 100644 index 139a39d..0000000 --- a/llm_trainer/lineage.py +++ /dev/null @@ -1,244 +0,0 @@ -from __future__ import annotations - -import hashlib -import json -import shutil -from datetime import datetime -from pathlib import Path -from typing import Any, Optional -from uuid import uuid4 - -from .manifest_store import ManifestStore - - -def utc_timestamp() -> str: - """Return a compact UTC timestamp for artifact identifiers. - - Returns: - Timestamp string. - """ - - return datetime.utcnow().strftime("%Y%m%dT%H%M%SZ") - - -def stable_json_hash(value: Any) -> str: - """Return a deterministic short hash for JSON-serializable data. - - Uses streaming hashing so extremely large manifests do not need to be - converted into one huge JSON string first. - - Args: - value: Any JSON-serializable object. - - Returns: - Twelve-character SHA-256 prefix. - """ - - digest = hashlib.sha256() - - def _update(obj: Any) -> None: - if obj is None: - digest.update(b"null") - - elif isinstance(obj, bool): - digest.update(b"true" if obj else b"false") - - elif isinstance(obj, (int, float)): - digest.update(str(obj).encode("utf-8")) - - elif isinstance(obj, str): - digest.update(obj.encode("utf-8")) - - elif isinstance(obj, dict): - digest.update(b"{") - for key in sorted(obj.keys(), key=str): - digest.update(str(key).encode("utf-8")) - digest.update(b":") - _update(obj[key]) - digest.update(b",") - digest.update(b"}") - - elif isinstance(obj, (list, tuple)): - digest.update(b"[") - for item in obj: - _update(item) - digest.update(b",") - digest.update(b"]") - - elif isinstance(obj, set): - digest.update(b"") - for item in sorted(obj, key=str): - _update(item) - digest.update(b",") - - else: - digest.update(str(obj).encode("utf-8")) - - _update(value) - - return digest.hexdigest()[:12] - - -def read_json(path: Path, default: Optional[Any] = None) -> Any: - """Read JSON from disk. - - Args: - path: JSON file path. - default: Value returned when the file does not exist or cannot be read. - - Returns: - Parsed JSON or default. - """ - - if not path.exists(): - return default - try: - return json.loads(path.read_text(encoding="utf-8")) - except Exception: - return default - - -def write_json(path: Path, data: Any) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - - with path.open("w", encoding="utf-8") as file: - json.dump( - data, - file, - indent=2, - ensure_ascii=False, - ) - - -def next_version_number(lineage: dict[str, Any]) -> int: - """Return the next dataset version number. - - Args: - lineage: Existing lineage dictionary. - - Returns: - Next one-based version number. - """ - - versions = lineage.get("versions", []) - return len(versions) + 1 - - -def ensure_dataset_lineage(output_dir: Path) -> dict[str, Any]: - """Load or create dataset lineage metadata. - - Args: - output_dir: Dataset output folder. - - Returns: - Dataset lineage dictionary. - """ - - lineage_path = output_dir / "dataset_lineage.json" - lineage = read_json(lineage_path, default=None) - if isinstance(lineage, dict) and lineage.get("dataset_id"): - lineage.setdefault("versions", []) - return lineage - return { - "schema": "micro_llm_dataset_lineage", - "version": 1, - "dataset_id": f"ds_{uuid4().hex[:12]}", - "created_at": utc_timestamp(), - "versions": [], - } - - -def record_dataset_version(output_dir: Path, summary: dict[str, Any], manifest_store: ManifestStore) -> dict[str, Any]: - """Record a new dataset version and snapshot its metadata. - - Args: - output_dir: Dataset output folder. - summary: Dataset summary dictionary. - manifest_store: Open manifest store tracking every source file. - - Returns: - Version metadata appended to lineage. - """ - - lineage = ensure_dataset_lineage(output_dir) - version_number = next_version_number(lineage) - - # ------------------------------------------------------------------ - # Build a memory-efficient fingerprint by streaming rows straight from - # SQLite (already returned in manifest_key order) instead of loading - # every file's metadata into one big dict or JSON string first. This - # keeps fingerprinting cheap in RAM no matter how many source files a - # project has. - # ------------------------------------------------------------------ - - digest = hashlib.sha256() - - for manifest_key, info in manifest_store.iter_files(): - digest.update(manifest_key.encode("utf-8")) - digest.update(str(info.get("sha256", "")).encode("utf-8")) - digest.update(str(info.get("size", "")).encode("utf-8")) - digest.update(str(info.get("mtime_ns", "")).encode("utf-8")) - - digest.update( - json.dumps( - summary.get("dataset_config", {}), - sort_keys=True, - ensure_ascii=False, - default=str, - ).encode("utf-8") - ) - - digest.update(str(summary.get("tokenizer_vocab_size")).encode("utf-8")) - digest.update(str(summary.get("tokenizer_sha256")).encode("utf-8")) - digest.update(str(summary.get("tokenizer_strategy")).encode("utf-8")) - - source_fingerprint = digest.hexdigest()[:12] - - # ------------------------------------------------------------------ - - version_id = f"v{version_number:03d}_{utc_timestamp()}_{source_fingerprint}" - - version = { - "version_number": version_number, - "version_id": version_id, - "created_at": utc_timestamp(), - "source_fingerprint": source_fingerprint, - "document_count": summary.get("document_count"), - "character_count": summary.get("character_count"), - "token_count": summary.get("token_count"), - "tokenizer_vocab_size": summary.get("tokenizer_vocab_size"), - "tokenizer_sha256": summary.get("tokenizer_sha256"), - "code_sample_count": summary.get("code_sample_count"), - "prose_sample_count": summary.get("prose_sample_count"), - "prepare_mode": summary.get("prepare_mode"), - "tokenizer_strategy": summary.get("tokenizer_strategy"), - "summary_path": "dataset_summary.json", - "manifest_path": "dataset_manifest.sqlite3", - "snapshot_dir": f"versions/{version_id}", - "file_count": manifest_store.count(), - } - - lineage["updated_at"] = utc_timestamp() - lineage.setdefault("versions", []).append(version) - - summary["dataset_id"] = lineage["dataset_id"] - summary["dataset_version"] = version - - manifest_store.set_meta("dataset_id", lineage["dataset_id"]) - manifest_store.set_meta("dataset_version", version) - - snapshot_dir = output_dir / "versions" / version_id - snapshot_dir.mkdir(parents=True, exist_ok=True) - - write_json(snapshot_dir / "dataset_summary.json", summary) - # The manifest snapshot is a plain file copy of the SQLite database, - # not a JSON export -- for a very large file count, serializing every - # row to JSON here would hit exactly the same MemoryError this was - # built to avoid. A copy is fast regardless of size and needs no - # in-memory reconstruction of the data at all. - manifest_store.commit() - if manifest_store.db_path is not None and manifest_store.db_path.exists(): - shutil.copy2(manifest_store.db_path, snapshot_dir / "dataset_manifest.sqlite3") - write_json(output_dir / "dataset_lineage.json", lineage) - - return version \ No newline at end of file diff --git a/llm_trainer/llama_chat.py b/llm_trainer/llama_chat.py deleted file mode 100644 index 53cfc13..0000000 --- a/llm_trainer/llama_chat.py +++ /dev/null @@ -1,369 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -from threading import Lock -from time import perf_counter -from typing import Any, Callable, Optional - - -class LlamaChatSession: - """Persistent GGUF chat session backed by llama-cpp-python.""" - - def __init__(self, model_path: Path, n_ctx: int = 2048, n_threads: int = 4, n_gpu_layers: int = -1) -> None: - """Load a GGUF model once for repeated chat prompts. - - Args: - model_path: Path to a GGUF model file. - n_ctx: Context window used by llama.cpp. - n_threads: CPU thread count. - n_gpu_layers: Number of layers to offload to GPU. - - Raises: - ImportError: If llama-cpp-python is not installed. - FileNotFoundError: If the model path does not exist. - """ - - if not Path(model_path).exists(): - raise FileNotFoundError(f"GGUF model not found: {model_path}") - try: - import llama_cpp - from llama_cpp import Llama - except ImportError as exc: - raise ImportError("Install llama-cpp-python to load GGUF models.") from exc - - self.gpu_offload_supported = bool( - getattr(llama_cpp, "llama_supports_gpu_offload", lambda: False)() - ) - self.requested_gpu_layers = n_gpu_layers - if n_gpu_layers != 0 and not self.gpu_offload_supported: - raise RuntimeError( - "This llama-cpp-python install is CPU-only. GPU layers were requested, " - "but llama.cpp reports GPU offload support is unavailable. Reinstall " - "llama-cpp-python with CUDA, Metal, Vulkan, or another GPU backend, " - "or set GPU layers to 0 for CPU loading." - ) - - self.model_path = Path(model_path) - self._lock = Lock() - self._messages: list[dict[str, str]] = [] - try: - self._llm = Llama( - model_path=str(self.model_path), - n_ctx=n_ctx, - n_threads=n_threads, - n_gpu_layers=n_gpu_layers, - offload_kqv=n_gpu_layers != 0, - verbose=False, - ) - except ValueError as exc: - file_size = self.model_path.stat().st_size if self.model_path.exists() else 0 - hints = [] - if file_size == 0: - hints.append("The file is empty — the download may have failed.") - elif file_size < 1_000_000: - hints.append(f"The file is very small ({file_size:,} bytes) — it may be incomplete or corrupted.") - if any(c > 127 for c in str(self.model_path).encode("utf-8", errors="replace")): - hints.append("The path contains non-ASCII characters — try moving the model to a simple path.") - detail = " ".join(hints) if hints else "Verify the file is a valid GGUF model and is fully downloaded." - raise ValueError( - f"Failed to load GGUF model: {self.model_path}\n{detail}" - ) from exc - - @property - def runtime_summary(self) -> str: - """Return a short runtime summary. - - Returns: - Runtime summary text. - """ - - if self.requested_gpu_layers == 0: - return "Runtime: CPU" - if self.gpu_offload_supported: - return f"Runtime: GPU offload requested ({self.requested_gpu_layers} layers)" - return "Runtime: CPU-only llama.cpp build" - - def reset(self) -> None: - """Clear conversation history while keeping the model loaded.""" - - with self._lock: - self._messages = [] - - def generate( - self, - prompt: str, - system_prompt: str = "", - max_tokens: int = 512, - temperature: float = 0.7, - top_p: float = 0.9, - repeat_penalty: float = 1.1, - reasoning_effort: str = "Balanced", - thinking_enabled: bool = True, - ) -> str: - """Generate one assistant reply. - - Args: - prompt: User message. - system_prompt: Optional system instruction. - max_tokens: Maximum new tokens to generate. - temperature: Sampling temperature. - top_p: Nucleus sampling value. - repeat_penalty: Repetition penalty. - reasoning_effort: User-facing effort mode. - thinking_enabled: Whether to add reasoning-style system guidance. - - Returns: - Assistant reply text. - """ - - effort_instruction = self._effort_instruction(reasoning_effort) if thinking_enabled else self._plain_instruction() - messages = [] - if system_prompt.strip() or effort_instruction: - messages.append({"role": "system", "content": "\n".join(part for part in (system_prompt.strip(), effort_instruction) if part)}) - - with self._lock: - messages.extend(self._messages) - messages.append({"role": "user", "content": prompt}) - response = self._llm.create_chat_completion( - messages=messages, - max_tokens=max_tokens, - temperature=temperature, - top_p=top_p, - repeat_penalty=repeat_penalty, - ) - reply = response["choices"][0]["message"]["content"].strip() - self._messages.append({"role": "user", "content": prompt}) - self._messages.append({"role": "assistant", "content": reply}) - return reply - - def generate_stream( - self, - prompt: str, - system_prompt: str = "", - max_tokens: int = 512, - temperature: float = 0.7, - top_p: float = 0.9, - repeat_penalty: float = 1.1, - reasoning_effort: str = "Balanced", - thinking_enabled: bool = True, - progress: Optional[Callable[[Any], None]] = None, - should_stop: Optional[Callable[[], bool]] = None, - ) -> dict[str, Any]: - """Stream one assistant reply and report timing metrics. - - Args: - prompt: User message. - system_prompt: Optional system instruction. - max_tokens: Maximum new tokens to generate. - temperature: Sampling temperature. - top_p: Nucleus sampling value. - repeat_penalty: Repetition penalty. - reasoning_effort: User-facing effort mode. - thinking_enabled: Whether to add reasoning-style system guidance. - progress: Optional callback receiving stream events. - should_stop: Optional callback returning true when generation should stop. - - Returns: - Reply text and generation metrics. - """ - - effort_instruction = self._effort_instruction(reasoning_effort) if thinking_enabled else self._plain_instruction() - messages = [] - if system_prompt.strip() or effort_instruction: - messages.append({"role": "system", "content": "\n".join(part for part in (system_prompt.strip(), effort_instruction) if part)}) - - started_at = perf_counter() - reply_parts: list[str] = [] - chunk_count = 0 - with self._lock: - messages.extend(self._messages) - messages.append({"role": "user", "content": prompt}) - stream = self._llm.create_chat_completion( - messages=messages, - max_tokens=max_tokens, - temperature=temperature, - top_p=top_p, - repeat_penalty=repeat_penalty, - stream=True, - ) - for chunk in stream: - if should_stop and should_stop(): - break - delta = chunk["choices"][0].get("delta", {}).get("content", "") - if not delta: - continue - reply_parts.append(delta) - chunk_count += 1 - elapsed = max(perf_counter() - started_at, 0.001) - streamed_text = "".join(reply_parts).strip() - token_count = self._count_tokens(streamed_text) if streamed_text else 0 - if progress: - progress( - { - "type": "chat_delta", - "content": delta, - "elapsed_seconds": elapsed, - "chunk_count": chunk_count, - "token_count": token_count, - "tokens_per_second": token_count / elapsed if token_count else 0.0, - } - ) - - reply = "".join(reply_parts).strip() - token_count = self._count_tokens(reply) if reply else 0 - elapsed = max(perf_counter() - started_at, 0.001) - if reply: - self._messages.append({"role": "user", "content": prompt}) - self._messages.append({"role": "assistant", "content": reply}) - - return { - "reply": reply, - "elapsed_seconds": elapsed, - "token_count": token_count, - "tokens_per_second": token_count / elapsed if token_count else 0.0, - "stopped": bool(should_stop and should_stop()), - } - - def _count_tokens(self, text: str) -> int: - """Count generated tokens using the loaded llama tokenizer. - - Args: - text: Generated text. - - Returns: - Token count. - """ - - try: - return len(self._llm.tokenize(text.encode("utf-8"))) - except Exception: - return max(1, len(text.split())) - - @staticmethod - def _effort_instruction(reasoning_effort: str) -> str: - """Translate a UI effort label into a system instruction. - - Args: - reasoning_effort: Selected effort label. - - Returns: - Instruction text. - """ - - if reasoning_effort == "Fast": - return "Answer concisely and prioritize speed. Put code inside fenced Markdown code blocks with language labels." - if reasoning_effort == "Deep": - return "Think carefully, reason through the problem, and provide a detailed answer when useful. Put code inside fenced Markdown code blocks with language labels." - return "Use balanced reasoning and answer clearly. Put code inside fenced Markdown code blocks with language labels." - - @staticmethod - def _plain_instruction() -> str: - """Return the non-thinking chat formatting instruction. - - Returns: - Plain response instruction. - """ - - return "Answer directly. Put code inside fenced Markdown code blocks with language labels." - - -def load_llama_chat_session(model_path: Path, n_ctx: int, n_threads: int, n_gpu_layers: int) -> LlamaChatSession: - """Load a GGUF chat session. - - Args: - model_path: Path to a GGUF model file. - n_ctx: Context window. - n_threads: CPU thread count. - n_gpu_layers: GPU offload layer count. - - Returns: - Loaded chat session. - """ - - return LlamaChatSession(model_path, n_ctx=n_ctx, n_threads=n_threads, n_gpu_layers=n_gpu_layers) - - -def generate_chat_reply( - session: LlamaChatSession, - prompt: str, - system_prompt: str, - max_tokens: int, - temperature: float, - top_p: float, - repeat_penalty: float, - reasoning_effort: str, - thinking_enabled: bool = True, -) -> str: - """Generate a reply from a loaded chat session. - - Args: - session: Loaded llama.cpp chat session. - prompt: User message. - system_prompt: Optional system instruction. - max_tokens: Maximum new tokens. - temperature: Sampling temperature. - top_p: Nucleus sampling value. - repeat_penalty: Repetition penalty. - reasoning_effort: Effort mode label. - thinking_enabled: Whether reasoning-style guidance is enabled. - - Returns: - Assistant reply. - """ - - return session.generate( - prompt, - system_prompt=system_prompt, - max_tokens=max_tokens, - temperature=temperature, - top_p=top_p, - repeat_penalty=repeat_penalty, - reasoning_effort=reasoning_effort, - thinking_enabled=thinking_enabled, - ) - - -def stream_chat_reply( - session: LlamaChatSession, - prompt: str, - system_prompt: str, - max_tokens: int, - temperature: float, - top_p: float, - repeat_penalty: float, - reasoning_effort: str, - thinking_enabled: bool = True, - progress: Optional[Callable[[Any], None]] = None, - should_stop: Optional[Callable[[], bool]] = None, -) -> dict[str, Any]: - """Stream a reply from a loaded chat session. - - Args: - session: Loaded llama.cpp chat session. - prompt: User message. - system_prompt: Optional system instruction. - max_tokens: Maximum new tokens. - temperature: Sampling temperature. - top_p: Nucleus sampling value. - repeat_penalty: Repetition penalty. - reasoning_effort: Effort mode label. - thinking_enabled: Whether reasoning-style guidance is enabled. - progress: Optional stream event callback. - should_stop: Optional cancellation callback. - - Returns: - Reply text and timing metrics. - """ - - return session.generate_stream( - prompt, - system_prompt=system_prompt, - max_tokens=max_tokens, - temperature=temperature, - top_p=top_p, - repeat_penalty=repeat_penalty, - reasoning_effort=reasoning_effort, - thinking_enabled=thinking_enabled, - progress=progress, - should_stop=should_stop, - ) diff --git a/llm_trainer/manifest_store.py b/llm_trainer/manifest_store.py deleted file mode 100644 index 2b45e61..0000000 --- a/llm_trainer/manifest_store.py +++ /dev/null @@ -1,261 +0,0 @@ -from __future__ import annotations - -import json -import sqlite3 -from pathlib import Path -from typing import Any, Iterator, Optional - - -# How many upserts to batch into one SQLite transaction/commit. Committing -# after every single file would be safest against a crash mid-scan, but -# each commit costs a real fsync; batching keeps that cost bounded while -# still checkpointing progress regularly during a very large scan (e.g. -# millions of source files) instead of risking the whole scan's progress -# on one final commit. -COMMIT_BATCH_SIZE = 1000 - - -class ManifestStore: - """SQLite-backed replacement for the old single-JSON-file manifest. - - The previous manifest was one JSON file holding a dict with one entry - per source file (or per online dataset). For a project with a very - large number of source files (hundreds of thousands to millions), that - dict -- and the JSON string built from it -- had to be held entirely in - memory, which could raise ``MemoryError`` during dataset preparation. - - This store keeps one row per file in a small SQLite database instead, - with each row's entry stored as its own small JSON blob. Entries for - local files and for online (``hf://...``) datasets have different - shapes (different fields), so a generic per-row blob is simpler and - more flexible than a rigid fixed-column schema, while still keeping - the one property that actually matters here: every lookup or update - touches exactly one row, so memory use stays flat no matter how many - files are tracked, and there is never a large in-memory dict or JSON - string covering the whole manifest. - """ - - def __init__(self, connection: sqlite3.Connection, db_path: Optional[Path] = None) -> None: - """Wrap an open SQLite connection. - - Args: - connection: Open SQLite connection with the manifest schema - already created. - db_path: Path to the underlying database file, if known. Used - by callers that need to copy the raw file (for example, to - snapshot it into a dataset version folder). - """ - - self._connection = connection - self._pending_since_commit = 0 - self.db_path = db_path - - @classmethod - def open(cls, db_path: Path, legacy_json_path: Optional[Path] = None) -> "ManifestStore": - """Open (creating if needed) a manifest database. - - If ``db_path`` does not exist yet but ``legacy_json_path`` does, - the old JSON manifest is imported once, so existing projects - prepared before this change keep their file cache instead of - starting from scratch. - - Args: - db_path: SQLite database file path. - legacy_json_path: Old ``dataset_manifest.json`` path to migrate - from, if present. - - Returns: - Opened manifest store. - """ - - db_path.parent.mkdir(parents=True, exist_ok=True) - is_new = not db_path.exists() - connection = sqlite3.connect(str(db_path)) - connection.execute("PRAGMA journal_mode=WAL") - connection.execute( - """ - CREATE TABLE IF NOT EXISTS files ( - manifest_key TEXT PRIMARY KEY, - entry_json TEXT NOT NULL - ) - """ - ) - connection.execute( - """ - CREATE TABLE IF NOT EXISTS meta ( - key TEXT PRIMARY KEY, - value TEXT - ) - """ - ) - connection.commit() - store = cls(connection, db_path=db_path) - if is_new and legacy_json_path is not None and legacy_json_path.exists(): - store._migrate_from_json(legacy_json_path) - return store - - def _migrate_from_json(self, legacy_json_path: Path) -> None: - """One-time import of an old JSON manifest into this store. - - Reads the legacy file once (the same memory cost the old code - always had for that one file), writes every entry into SQLite, and - does not touch the legacy file itself -- it is left in place as a - harmless leftover the user can delete later. - - Args: - legacy_json_path: Old ``dataset_manifest.json`` path. - """ - - try: - legacy = json.loads(legacy_json_path.read_text(encoding="utf-8")) - except Exception: - return - if not isinstance(legacy, dict): - return - files = legacy.get("files", {}) - if isinstance(files, dict): - for manifest_key, entry in files.items(): - if isinstance(entry, dict): - self.upsert(manifest_key, entry, commit=False) - for meta_key in ("dataset_config", "cache_key", "dataset_id", "dataset_version"): - if meta_key in legacy: - self.set_meta(meta_key, legacy[meta_key], commit=False) - self.commit() - - def get(self, manifest_key: str) -> Optional[dict[str, Any]]: - """Return one file's manifest entry, or ``None`` if not tracked. - - Args: - manifest_key: Resolved source path (or ``hf://dataset_id``) - used as the row's primary key. - - Returns: - Entry dictionary, or ``None``. - """ - - row = self._connection.execute( - "SELECT entry_json FROM files WHERE manifest_key = ?", - (manifest_key,), - ).fetchone() - if row is None: - return None - try: - return json.loads(row[0]) - except Exception: - return None - - def upsert(self, manifest_key: str, entry: dict[str, Any], commit: bool = True) -> None: - """Insert or update one file's manifest entry. - - Args: - manifest_key: Resolved source path (or ``hf://dataset_id``). - entry: Entry fields. Any JSON-serializable shape is accepted -- - local files and online datasets store different fields. - commit: Whether to commit immediately. Callers doing many - upserts in a row (e.g. scanning thousands of files) can - pass ``False`` and let the internal batch-size safety net - (or an explicit call to ``commit()``) checkpoint instead, - to amortize the fsync cost of committing. - """ - - self._connection.execute( - """ - INSERT INTO files (manifest_key, entry_json) VALUES (?, ?) - ON CONFLICT(manifest_key) DO UPDATE SET entry_json=excluded.entry_json - """, - (manifest_key, json.dumps(entry, ensure_ascii=False, default=str)), - ) - self._pending_since_commit += 1 - if commit: - self.commit() - elif self._pending_since_commit >= COMMIT_BATCH_SIZE: - # Safety net for callers doing many commit=False upserts in a - # row (a large file scan): checkpoint periodically so a crash - # partway through doesn't lose the entire scan's progress, even - # though the caller hasn't explicitly called commit() yet. - self.commit() - - def iter_files(self) -> Iterator[tuple[str, dict[str, Any]]]: - """Yield every tracked file's ``(manifest_key, entry)`` pair. - - Rows are streamed directly from SQLite in primary-key order, one at - a time, rather than being loaded into one big list or dict first -- - this is what makes fingerprinting and export safe for very large - file counts. - - Yields: - Pairs of manifest key and entry dictionary. - """ - - cursor = self._connection.execute("SELECT manifest_key, entry_json FROM files ORDER BY manifest_key") - for manifest_key, entry_json in cursor: - try: - entry = json.loads(entry_json) - except Exception: - entry = {} - yield manifest_key, entry - - def count(self) -> int: - """Return how many files are tracked. - - Returns: - Row count. - """ - - row = self._connection.execute("SELECT COUNT(*) FROM files").fetchone() - return int(row[0]) if row else 0 - - def get_meta(self, key: str, default: Any = None) -> Any: - """Return a small top-level metadata value (not per-file). - - Args: - key: Metadata key, such as ``"dataset_config"`` or - ``"cache_key"``. - default: Value returned if the key is not set. - - Returns: - Parsed JSON value, or ``default``. - """ - - row = self._connection.execute("SELECT value FROM meta WHERE key = ?", (key,)).fetchone() - if row is None: - return default - try: - return json.loads(row[0]) - except Exception: - return default - - def set_meta(self, key: str, value: Any, commit: bool = True) -> None: - """Set a small top-level metadata value (not per-file). - - Args: - key: Metadata key. - value: JSON-serializable value. - commit: Whether to commit immediately. - """ - - self._connection.execute( - "INSERT INTO meta (key, value) VALUES (?, ?) " - "ON CONFLICT(key) DO UPDATE SET value=excluded.value", - (key, json.dumps(value, ensure_ascii=False, default=str)), - ) - if commit: - self.commit() - - def commit(self) -> None: - """Commit any pending writes.""" - - self._connection.commit() - self._pending_since_commit = 0 - - def close(self) -> None: - """Commit and close the underlying connection.""" - - self.commit() - self._connection.close() - - def __enter__(self) -> "ManifestStore": - return self - - def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: - self.close() \ No newline at end of file diff --git a/llm_trainer/microgpt_chat.py b/llm_trainer/microgpt_chat.py deleted file mode 100644 index f971b12..0000000 --- a/llm_trainer/microgpt_chat.py +++ /dev/null @@ -1,316 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -from threading import Lock -from time import perf_counter -from typing import Any, Callable, Optional - -import torch -import torch.nn.functional as F - -from llm_trainer.config import ModelConfig -from llm_trainer.model import MicroGPT -from llm_trainer.tokenizer import EOS_TOKEN, load_tokenizer, token_id - - -class MicroGPTChatSession: - """Persistent chat session backed by a native MicroGPT checkpoint.""" - - def __init__(self, model_path: Path, device: str = "auto") -> None: - """Load a native MicroGPT checkpoint for repeated prompts. - - Args: - model_path: Model folder or checkpoint path. - device: Device selector: auto, cuda, or cpu. - - Raises: - FileNotFoundError: If checkpoint or tokenizer files are missing. - ValueError: If the checkpoint is not a MicroGPT checkpoint. - """ - - self.model_path = _resolve_model_checkpoint(model_path) - self.model_dir = self.model_path.parent - tokenizer_path = self.model_dir / "tokenizer.json" - if not tokenizer_path.exists(): - # Training checkpoints are stored in a child checkpoints folder, - # while the tokenizer is copied to the training output directory. - output_tokenizer = self.model_dir.parent / "tokenizer.json" - if output_tokenizer.exists(): - tokenizer_path = output_tokenizer - - if not tokenizer_path.exists(): - raise FileNotFoundError( - "Tokenizer not found beside checkpoint or its training output folder: " - f"{self.model_path}. Expected {self.model_dir / 'tokenizer.json'} " - f"or {self.model_dir.parent / 'tokenizer.json'}." - ) - requested_device = device.lower().strip() - self.device = "cuda" if requested_device == "auto" and torch.cuda.is_available() else requested_device - if self.device == "cuda" and not torch.cuda.is_available(): - self.device = "cpu" - if self.device not in {"cuda", "cpu"}: - self.device = "cpu" - - checkpoint = torch.load(self.model_path, map_location=self.device) - config_data = checkpoint.get("model_config") - state_dict = checkpoint.get("model_state_dict") - if not isinstance(config_data, dict) or not state_dict: - raise ValueError("Checkpoint must contain model_config and model_state_dict.") - self.config = ModelConfig(**config_data) - self.model = MicroGPT(self.config).to(self.device) - self.model.load_state_dict(state_dict) - self.model.eval() - self.tokenizer = load_tokenizer(tokenizer_path) - self.eos_id = token_id(self.tokenizer, EOS_TOKEN) - self._lock = Lock() - self._messages: list[dict[str, str]] = [] - - @property - def runtime_summary(self) -> str: - """Return a short runtime summary. - - Returns: - Runtime summary text. - """ - - return ( - f"Runtime: native MicroGPT on {self.device.upper()} | " - f"{self.config.layer_count} layers, {self.config.embedding_size} hidden, ctx {self.config.context_length}" - ) - - def reset(self) -> None: - """Clear conversation history while keeping the model loaded.""" - - with self._lock: - self._messages = [] - - def generate_stream( - self, - prompt: str, - system_prompt: str = "", - max_tokens: int = 512, - temperature: float = 0.7, - top_p: float = 0.9, - repeat_penalty: float = 1.1, - reasoning_effort: str = "Balanced", - thinking_enabled: bool = True, - progress: Optional[Callable[[Any], None]] = None, - should_stop: Optional[Callable[[], bool]] = None, - ) -> dict[str, Any]: - """Stream one assistant reply and report timing metrics. - - Args: - prompt: User message. - system_prompt: Optional system instruction. - max_tokens: Maximum new tokens to generate. - temperature: Sampling temperature. - top_p: Nucleus sampling value. - repeat_penalty: Penalty for generated token repetition. - reasoning_effort: Effort mode label. - thinking_enabled: Whether to add reasoning guidance. - progress: Optional progress callback. - should_stop: Optional cancellation callback. - - Returns: - Reply text and timing metrics. - """ - - started_at = perf_counter() - reply_parts: list[str] = [] - generated_ids: list[int] = [] - with self._lock, torch.no_grad(): - prompt_text = self._render_prompt(prompt, system_prompt, reasoning_effort, thinking_enabled) - input_ids = self.tokenizer.encode(prompt_text).ids[-self.config.context_length :] - ids = torch.tensor([input_ids], dtype=torch.long, device=self.device) - emitted_text = "" - for _ in range(max_tokens): - if should_stop and should_stop(): - break - idx_cond = ids[:, -self.config.context_length :] - logits = self.model(idx_cond)[:, -1, :] - logits = self._apply_repeat_penalty(logits, generated_ids, repeat_penalty) - next_id = self._sample_next_token(logits, temperature, top_p) - if next_id == self.eos_id and generated_ids: - break - ids = torch.cat((ids, torch.tensor([[next_id]], dtype=torch.long, device=self.device)), dim=1) - generated_ids.append(next_id) - # Decoding one token at a time can split a multi-byte - # character (emoji, accented letters, CJK text) across - # token boundaries, since each token only covers part of - # its UTF-8 bytes -- decoding it alone then produces the - # Unicode replacement character ("\ufffd") rather than the - # real text. Decoding the whole sequence generated so far - # recovers the correct character once enough tokens have - # arrived, but a naive length-based diff against what was - # already shown can still get stuck: a stale replacement - # character and the character that later replaces it are - # often the same length (one code point each), so slicing - # by length alone would silently keep showing the wrong - # one forever. Withholding emission whenever the decoded - # text currently ends in a replacement character -- and - # only emitting once it resolves -- avoids that. - full_text = self.tokenizer.decode(generated_ids, skip_special_tokens=True) - if full_text.endswith("\ufffd"): - continue - piece = full_text[len(emitted_text) :] - if not piece: - continue - emitted_text = full_text - reply_parts.append(piece) - elapsed = max(perf_counter() - started_at, 0.001) - if progress: - progress( - { - "type": "chat_delta", - "content": piece, - "elapsed_seconds": elapsed, - "token_count": len(generated_ids), - "tokens_per_second": len(generated_ids) / elapsed, - } - ) - - reply = self.tokenizer.decode(generated_ids, skip_special_tokens=True).strip() if generated_ids else "" - # If generation stopped (should_stop or max_tokens) exactly - # mid-way through a multi-byte character, the trailing bytes - # are incomplete and decode to a replacement character. That's - # a truncation artifact, not real content, so trim it rather - # than showing it to the user. - reply = reply.rstrip("\ufffd") - if reply: - self._messages.append({"role": "user", "content": prompt}) - self._messages.append({"role": "assistant", "content": reply}) - elapsed = max(perf_counter() - started_at, 0.001) - return { - "reply": reply, - "elapsed_seconds": elapsed, - "token_count": len(generated_ids), - "tokens_per_second": len(generated_ids) / elapsed if generated_ids else 0.0, - "stopped": bool(should_stop and should_stop()), - } - - def _render_prompt(self, prompt: str, system_prompt: str, reasoning_effort: str, thinking_enabled: bool) -> str: - """Render chat history into plain text for MicroGPT. - - Args: - prompt: Latest user message. - system_prompt: Optional system instruction. - reasoning_effort: Effort mode label. - thinking_enabled: Whether reasoning guidance is enabled. - - Returns: - Prompt text. - """ - - instruction = self._effort_instruction(reasoning_effort) if thinking_enabled else self._plain_instruction() - parts = [] - if system_prompt.strip() or instruction: - parts.append(f"System: {' '.join(part for part in (system_prompt.strip(), instruction) if part)}") - for message in self._messages[-12:]: - role = "User" if message["role"] == "user" else "Assistant" - parts.append(f"{role}: {message['content']}") - parts.append(f"User: {prompt}") - parts.append("Assistant:") - return "\n".join(parts) - - def _apply_repeat_penalty(self, logits: torch.Tensor, generated_ids: list[int], repeat_penalty: float) -> torch.Tensor: - """Apply a simple repeat penalty to recently generated tokens.""" - - if repeat_penalty <= 1.0 or not generated_ids: - return logits - for token in set(generated_ids[-128:]): - logits[:, token] = logits[:, token] / repeat_penalty - return logits - - def _sample_next_token(self, logits: torch.Tensor, temperature: float, top_p: float) -> int: - """Sample the next token from logits.""" - - temperature = max(float(temperature), 1e-5) - logits = logits / temperature - if top_p < 1.0: - sorted_logits, sorted_indices = torch.sort(logits, descending=True) - probs = F.softmax(sorted_logits, dim=-1) - cumulative = torch.cumsum(probs, dim=-1) - remove = cumulative > max(0.01, min(1.0, float(top_p))) - remove[..., 1:] = remove[..., :-1].clone() - remove[..., 0] = False - sorted_logits = sorted_logits.masked_fill(remove, -float("inf")) - filtered = torch.full_like(logits, -float("inf")) - filtered.scatter_(1, sorted_indices, sorted_logits) - logits = filtered - probs = F.softmax(logits, dim=-1) - return int(torch.multinomial(probs, num_samples=1).item()) - - @staticmethod - def _effort_instruction(reasoning_effort: str) -> str: - """Translate effort label into prompt guidance.""" - - if reasoning_effort == "Fast": - return "Answer concisely. Put code inside fenced Markdown code blocks with language labels." - if reasoning_effort == "Deep": - return "Think carefully and provide a detailed answer when useful. Put code inside fenced Markdown code blocks with language labels." - return "Use balanced reasoning and answer clearly. Put code inside fenced Markdown code blocks with language labels." - - @staticmethod - def _plain_instruction() -> str: - """Return direct answer guidance.""" - - return "Answer directly. Put code inside fenced Markdown code blocks with language labels." - - -def _resolve_model_checkpoint(path: Path) -> Path: - """Resolve a model folder or checkpoint path to a checkpoint file.""" - - path = Path(path) - if path.is_dir(): - final_model = path / "final_model.pt" - if final_model.exists(): - return final_model - checkpoints = sorted((path / "checkpoints").glob("checkpoint_*.pt"), key=lambda item: item.stat().st_mtime, reverse=True) - if checkpoints: - return checkpoints[0] - if path.exists() and path.suffix == ".pt": - return path - raise FileNotFoundError(f"MicroGPT checkpoint not found: {path}") - - -def load_microgpt_chat_session(model_path: Path, device: str = "auto") -> MicroGPTChatSession: - """Load a native MicroGPT chat session. - - Args: - model_path: Model folder or checkpoint path. - device: Device selector. - - Returns: - Loaded MicroGPT chat session. - """ - - return MicroGPTChatSession(model_path, device=device) - - -def stream_microgpt_chat_reply( - session: MicroGPTChatSession, - prompt: str, - system_prompt: str, - max_tokens: int, - temperature: float, - top_p: float, - repeat_penalty: float, - reasoning_effort: str, - thinking_enabled: bool = True, - progress: Optional[Callable[[Any], None]] = None, - should_stop: Optional[Callable[[], bool]] = None, -) -> dict[str, Any]: - """Stream a reply from a native MicroGPT chat session.""" - - return session.generate_stream( - prompt, - system_prompt=system_prompt, - max_tokens=max_tokens, - temperature=temperature, - top_p=top_p, - repeat_penalty=repeat_penalty, - reasoning_effort=reasoning_effort, - thinking_enabled=thinking_enabled, - progress=progress, - should_stop=should_stop, - ) \ No newline at end of file diff --git a/llm_trainer/model.py b/llm_trainer/model.py deleted file mode 100644 index a5521f7..0000000 --- a/llm_trainer/model.py +++ /dev/null @@ -1,732 +0,0 @@ -from __future__ import annotations - -import math -from typing import Optional - -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch.utils.checkpoint import checkpoint - -from .config import ModelConfig - - -class LayerNorm(nn.Module): - """Layer normalization with optional bias.""" - - def __init__(self, size: int, bias: bool) -> None: - """Create layer normalization. - - Args: - size: Feature dimension. - bias: Whether to include a bias vector. - """ - - super().__init__() - self.weight = nn.Parameter(torch.ones(size)) - self.bias = nn.Parameter(torch.zeros(size)) if bias else None - - def forward(self, value: torch.Tensor) -> torch.Tensor: - """Normalize an input tensor. - - Args: - value: Tensor to normalize. - - Returns: - Normalized tensor. - """ - - return F.layer_norm(value, self.weight.shape, self.weight, self.bias, 1e-5) - - -class RMSNorm(nn.Module): - """Root mean square normalization used by Llama-style models.""" - - def __init__(self, size: int, eps: float = 1e-6) -> None: - """Create RMSNorm. - - Args: - size: Feature dimension. - eps: Numerical stability value. - """ - - super().__init__() - self.weight = nn.Parameter(torch.ones(size)) - self.eps = eps - - def forward(self, value: torch.Tensor) -> torch.Tensor: - """Normalize by root mean square. - - Args: - value: Input tensor. - - Returns: - Normalized tensor. - """ - - return self.weight * value * torch.rsqrt(value.pow(2).mean(dim=-1, keepdim=True) + self.eps) - - -def make_norm(config: ModelConfig) -> nn.Module: - """Create the configured normalization layer. - - Args: - config: Model configuration. - - Returns: - Normalization module. - """ - - if config.norm_type == "rmsnorm": - return RMSNorm(config.embedding_size) - return LayerNorm(config.embedding_size, bias=config.bias) - - -class LoRALinear(nn.Module): - """Linear layer with trainable low-rank LoRA adapters.""" - - def __init__(self, base: nn.Linear, rank: int, alpha: float, dropout: float) -> None: - """Create a LoRA wrapper around an existing linear layer. - - Args: - base: Frozen base linear layer. - rank: Adapter rank. - alpha: LoRA scaling alpha. - dropout: Dropout probability before the adapter. - """ - - super().__init__() - self.base = base - self.rank = rank - self.alpha = alpha - self.scaling = alpha / rank - self.dropout = nn.Dropout(dropout) - self.lora_a = nn.Parameter(torch.zeros(rank, base.in_features, device=base.weight.device, dtype=base.weight.dtype)) - self.lora_b = nn.Parameter(torch.zeros(base.out_features, rank, device=base.weight.device, dtype=base.weight.dtype)) - nn.init.kaiming_uniform_(self.lora_a, a=math.sqrt(5)) - nn.init.zeros_(self.lora_b) - self.base.weight.requires_grad_(False) - if self.base.bias is not None: - self.base.bias.requires_grad_(False) - - def forward(self, value: torch.Tensor) -> torch.Tensor: - """Apply the base projection plus LoRA update. - - Args: - value: Input tensor. - - Returns: - Projected tensor. - """ - - update = F.linear(F.linear(self.dropout(value), self.lora_a), self.lora_b) * self.scaling - return self.base(value) + update - - def merged_linear(self) -> nn.Linear: - """Return a plain linear layer with LoRA weights merged. - - Returns: - Linear layer equivalent to base plus LoRA update. - """ - - merged = nn.Linear(self.base.in_features, self.base.out_features, bias=self.base.bias is not None) - merged.weight.data.copy_(self.base.weight.data + (self.lora_b @ self.lora_a) * self.scaling) - if self.base.bias is not None and merged.bias is not None: - merged.bias.data.copy_(self.base.bias.data) - return merged - - -def _set_nested_module(root: nn.Module, module_name: str, module: nn.Module) -> None: - """Replace a nested module by dotted name. - - Args: - root: Root module. - module_name: Dotted module name. - module: Replacement module. - """ - - parent_name, child_name = module_name.rsplit(".", 1) if "." in module_name else ("", module_name) - parent = root.get_submodule(parent_name) if parent_name else root - setattr(parent, child_name, module) - - -def _lora_target_names(model: nn.Module, target_modules: str) -> set[str]: - """Resolve LoRA target module names. - - Args: - model: Model to inspect. - target_modules: Comma-separated target groups. - - Returns: - Set of module names to wrap. - """ - - groups = {part.strip().lower() for part in target_modules.split(",") if part.strip()} - if "all" in groups: - groups.update({"attention", "mlp"}) - names: set[str] = set() - for name, module in model.named_modules(): - if not isinstance(module, nn.Linear): - continue - if name.endswith("lm_head"): - continue - is_attention = ".attn." in name - is_mlp = ".mlp." in name - if ("attention" in groups and is_attention) or ("mlp" in groups and is_mlp): - names.add(name) - return names - - -def apply_lora_adapters(model: nn.Module, rank: int, alpha: float, dropout: float, target_modules: str) -> int: - """Attach LoRA adapters to selected linear layers. - - Args: - model: Model to modify in place. - rank: LoRA rank. - alpha: LoRA alpha. - dropout: LoRA dropout. - target_modules: Comma-separated target groups. - - Returns: - Number of wrapped modules. - """ - - names = _lora_target_names(model, target_modules) - for name in sorted(names): - module = model.get_submodule(name) - if isinstance(module, nn.Linear): - _set_nested_module(model, name, LoRALinear(module, rank, alpha, dropout)) - return len(names) - - -def freeze_non_lora_parameters(model: nn.Module) -> None: - """Freeze all parameters except LoRA adapter parameters. - - Args: - model: Model to update. - """ - - for name, parameter in model.named_parameters(): - parameter.requires_grad_(("lora_a" in name) or ("lora_b" in name)) - - -def lora_state_dict(model: nn.Module) -> dict[str, torch.Tensor]: - """Return trainable LoRA adapter tensors. - - Args: - model: Model containing LoRA adapters. - - Returns: - LoRA-only state dictionary. - """ - - return { - name: tensor.detach().cpu() - for name, tensor in model.state_dict().items() - if ".lora_a" in name or ".lora_b" in name - } - - -def load_lora_state_dict(model: nn.Module, state: dict[str, torch.Tensor]) -> None: - """Load LoRA adapter tensors into a model. - - Args: - model: Model containing LoRA adapters. - state: LoRA-only state dictionary. - """ - - model.load_state_dict(state, strict=False) - - -def merge_lora_adapters(model: nn.Module) -> int: - """Merge LoRA adapters into plain linear layers. - - Args: - model: Model to modify in place. - - Returns: - Number of merged LoRA modules. - """ - - merged = 0 - for name, module in list(model.named_modules()): - if isinstance(module, LoRALinear): - # merged_linear() builds a fresh nn.Linear, which defaults to - # CPU regardless of what device the original layer was on. - # Moving it explicitly avoids ending up with a model whose - # merged layers are silently on a different device than - # everything else in it. - merged_linear = module.merged_linear().to(module.base.weight.device) - _set_nested_module(model, name, merged_linear) - merged += 1 - return merged - - -def lora_parameter_count(model: nn.Module) -> int: - """Count trainable LoRA parameters. - - Args: - model: Model containing LoRA adapters. - - Returns: - Number of trainable adapter parameters. - """ - - return sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad) - - -class RotaryEmbedding(nn.Module): - """Rotary positional embedding cache for attention heads.""" - - def __init__(self, head_size: int, context_length: int, theta: float) -> None: - """Create RoPE caches. - - Args: - head_size: Attention head dimension. - context_length: Maximum context length. - theta: Frequency base. - """ - - super().__init__() - inv_freq = 1.0 / (theta ** (torch.arange(0, head_size, 2).float() / head_size)) - positions = torch.arange(context_length, dtype=torch.float) - freqs = torch.einsum("i,j->ij", positions, inv_freq) - emb = torch.cat((freqs, freqs), dim=-1) - self.register_buffer("cos", emb.cos()[None, None, :, :], persistent=False) - self.register_buffer("sin", emb.sin()[None, None, :, :], persistent=False) - - def forward(self, query: torch.Tensor, key: torch.Tensor, start_pos: int = 0) -> tuple[torch.Tensor, torch.Tensor]: - """Apply RoPE to query and key tensors. - - Args: - query: Query tensor with shape ``[batch, heads, tokens, head_size]``. - key: Key tensor with shape ``[batch, heads, tokens, head_size]``. - start_pos: Absolute starting token position. - - Returns: - Rotated query and key tensors. - """ - - token_count = query.size(-2) - cos = self.cos[:, :, start_pos : start_pos + token_count, :] - sin = self.sin[:, :, start_pos : start_pos + token_count, :] - return (query * cos) + (_rotate_half(query) * sin), (key * cos) + (_rotate_half(key) * sin) - - -def _rotate_half(value: torch.Tensor) -> torch.Tensor: - """Rotate the last dimension in RoPE pairs. - - Args: - value: Tensor to rotate. - - Returns: - Rotated tensor. - """ - - first, second = value.chunk(2, dim=-1) - return torch.cat((-second, first), dim=-1) - - -class CausalSelfAttention(nn.Module): - """Causal multi-head self-attention block.""" - - def __init__(self, config: ModelConfig) -> None: - """Create attention module. - - Args: - config: Model architecture configuration. - """ - - super().__init__() - self.head_count = config.head_count - self.kv_head_count = config.resolved_kv_head_count() - self.embedding_size = config.embedding_size - self.position_encoding = config.position_encoding - self.attention_backend = config.attention_backend - self.attention_window = config.attention_window - self.head_size = config.embedding_size // config.head_count - self.kv_embedding_size = self.kv_head_count * self.head_size - self.c_attn = nn.Linear( - config.embedding_size, - config.embedding_size + (2 * self.kv_embedding_size), - bias=config.bias, - ) - self.c_proj = nn.Linear(config.embedding_size, config.embedding_size, bias=config.bias) - self.attn_dropout = nn.Dropout(config.dropout) - self.resid_dropout = nn.Dropout(config.dropout) - self.rotary = ( - RotaryEmbedding(self.head_size, config.context_length, config.rope_theta) - if config.position_encoding == "rope" - else None - ) - self.register_buffer( - "mask", - torch.tril(torch.ones(config.context_length, config.context_length, dtype=torch.bool)).view( - 1, 1, config.context_length, config.context_length - ), - ) - - def forward( - self, - value: torch.Tensor, - past_kv: Optional[tuple[torch.Tensor, torch.Tensor]] = None, - start_pos: int = 0, - use_cache: bool = False, - ) -> torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: - """Apply causal self-attention. - - Args: - value: Input hidden states. - past_kv: Optional cached key/value tensors. - start_pos: Absolute starting token position. - use_cache: Whether to return updated key/value cache. - - Returns: - Attention output tensor, plus cache when requested. - """ - - batch_size, token_count, channel_count = value.size() - qkv = self.c_attn(value) - query, key, val = qkv.split((self.embedding_size, self.kv_embedding_size, self.kv_embedding_size), dim=2) - - key = key.view(batch_size, token_count, self.kv_head_count, self.head_size).transpose(1, 2) - query = query.view(batch_size, token_count, self.head_count, self.head_size).transpose(1, 2) - val = val.view(batch_size, token_count, self.kv_head_count, self.head_size).transpose(1, 2) - if self.rotary is not None: - query, key = self.rotary(query, key, start_pos=start_pos) - - if past_kv is not None: - past_key, past_val = past_kv - key = torch.cat((past_key, key), dim=-2) - val = torch.cat((past_val, val), dim=-2) - if key.size(-2) > self.mask.size(-1): - key = key[:, :, -self.mask.size(-1) :, :] - val = val[:, :, -self.mask.size(-1) :, :] - present = (key, val) - expanded_key = self._expand_kv(key) - expanded_val = self._expand_kv(val) - - key_count = expanded_key.size(-2) - if past_kv is None: - mask = self.mask[:, :, :token_count, :key_count] - else: - start = max(0, key_count - token_count) - mask = self.mask[:, :, start : start + token_count, :key_count] - if self.attention_window > 0: - positions = torch.arange(key_count, device=value.device) - query_positions = torch.arange(key_count - token_count, key_count, device=value.device) - window_mask = positions[None, :] >= (query_positions[:, None] - self.attention_window + 1) - mask = mask & window_mask.view(1, 1, token_count, key_count) - - if self.attention_backend == "sdpa" and hasattr(F, "scaled_dot_product_attention"): - if self.attention_window <= 0 and past_kv is None: - # Plain full-sequence causal attention (the common training - # case): the mask built above is mathematically identical to - # is_causal=True. Passing an explicit attn_mask tensor here - # instead can prevent PyTorch from dispatching to the fused - # FlashAttention kernel on supported hardware, falling back - # to the slower/more memory-hungry "efficient" or "math" - # backends even when the SDPA/Flash backend is selected. - y = F.scaled_dot_product_attention( - query, - expanded_key, - expanded_val, - is_causal=True, - dropout_p=self.attn_dropout.p if self.training else 0.0, - ) - else: - attn_mask = mask[:, :, :, :].bool() - y = F.scaled_dot_product_attention( - query, - expanded_key, - expanded_val, - attn_mask=attn_mask, - dropout_p=self.attn_dropout.p if self.training else 0.0, - ) - else: - attention = (query @ expanded_key.transpose(-2, -1)) * (1.0 / math.sqrt(expanded_key.size(-1))) - attention = attention.masked_fill(mask == 0, float("-inf")) - attention = F.softmax(attention, dim=-1) - attention = self.attn_dropout(attention) - y = attention @ expanded_val - y = y.transpose(1, 2).contiguous().view(batch_size, token_count, channel_count) - output = self.resid_dropout(self.c_proj(y)) - if use_cache: - return output, present - return output - - def _expand_kv(self, value: torch.Tensor) -> torch.Tensor: - """Expand grouped key/value heads to query head count. - - Args: - value: Key or value tensor with key/value head count. - - Returns: - Tensor with one key/value head per query head. - """ - - if self.kv_head_count == self.head_count: - return value - repeat_count = self.head_count // self.kv_head_count - return value.repeat_interleave(repeat_count, dim=1) - - -class MLP(nn.Module): - """Feed-forward network inside a transformer block.""" - - def __init__(self, config: ModelConfig) -> None: - """Create feed-forward network. - - Args: - config: Model architecture configuration. - """ - - super().__init__() - self.mlp_type = config.mlp_type - hidden_size = 4 * config.embedding_size - if self.mlp_type == "swiglu": - self.w1 = nn.Linear(config.embedding_size, hidden_size, bias=config.bias) - self.w2 = nn.Linear(hidden_size, config.embedding_size, bias=config.bias) - self.w3 = nn.Linear(config.embedding_size, hidden_size, bias=config.bias) - self.dropout = nn.Dropout(config.dropout) - else: - self.net = nn.Sequential( - nn.Linear(config.embedding_size, hidden_size, bias=config.bias), - nn.GELU(), - nn.Linear(hidden_size, config.embedding_size, bias=config.bias), - nn.Dropout(config.dropout), - ) - - def forward(self, value: torch.Tensor) -> torch.Tensor: - """Apply feed-forward transformation. - - Args: - value: Input hidden states. - - Returns: - Transformed hidden states. - """ - - if self.mlp_type == "swiglu": - return self.dropout(self.w2(F.silu(self.w1(value)) * self.w3(value))) - return self.net(value) - - -class Block(nn.Module): - """Transformer block with attention and MLP.""" - - def __init__(self, config: ModelConfig) -> None: - """Create a transformer block. - - Args: - config: Model architecture configuration. - """ - - super().__init__() - self.ln_1 = make_norm(config) - self.attn = CausalSelfAttention(config) - self.ln_2 = make_norm(config) - self.mlp = MLP(config) - - def forward(self, value: torch.Tensor) -> torch.Tensor: - """Apply transformer block. - - Args: - value: Input hidden states. - - Returns: - Updated hidden states. - """ - - value = value + self.attn(self.ln_1(value)) - value = value + self.mlp(self.ln_2(value)) - return value - - def forward_with_cache( - self, - value: torch.Tensor, - past_kv: Optional[tuple[torch.Tensor, torch.Tensor]] = None, - start_pos: int = 0, - ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: - """Apply transformer block and return updated KV cache. - - Args: - value: Input hidden states. - past_kv: Optional cached key/value tensors. - start_pos: Absolute starting token position. - - Returns: - Updated hidden states and key/value cache. - """ - - attention_output, present = self.attn(self.ln_1(value), past_kv=past_kv, start_pos=start_pos, use_cache=True) - value = value + attention_output - value = value + self.mlp(self.ln_2(value)) - return value, present - - -class MicroGPT(nn.Module): - """Small GPT-style causal language model.""" - - def __init__(self, config: ModelConfig) -> None: - """Create the model. - - Args: - config: Model architecture configuration. - """ - - super().__init__() - config.validate() - self.config = config - self.token_embedding = nn.Embedding(config.vocab_size, config.embedding_size) - self.position_embedding = ( - nn.Embedding(config.context_length, config.embedding_size) - if config.position_encoding == "learned" - else None - ) - self.drop = nn.Dropout(config.dropout) - self.blocks = nn.Sequential(*[Block(config) for _ in range(config.layer_count)]) - self.ln_f = make_norm(config) - self.lm_head = nn.Linear(config.embedding_size, config.vocab_size, bias=False) - self.gradient_checkpointing = False - self.token_embedding.weight = self.lm_head.weight - self.apply(self._init_weights) - - def enable_gradient_checkpointing(self, enabled: bool = True) -> None: - """Trade extra compute for substantially lower activation memory.""" - - self.gradient_checkpointing = bool(enabled) - - def _init_weights(self, module: nn.Module) -> None: - """Initialize module weights. - - Args: - module: Module to initialize. - """ - - if isinstance(module, nn.Linear): - nn.init.normal_(module.weight, mean=0.0, std=0.02) - if module.bias is not None: - nn.init.zeros_(module.bias) - elif isinstance(module, nn.Embedding): - nn.init.normal_(module.weight, mean=0.0, std=0.02) - - def forward(self, idx: torch.Tensor) -> torch.Tensor: - """Run a forward pass. - - Args: - idx: Token IDs with shape ``[batch, tokens]``. - - Returns: - Logits with shape ``[batch, tokens, vocab]``. - - Raises: - ValueError: If the sequence is longer than context length. - """ - - _, token_count = idx.size() - if token_count > self.config.context_length: - raise ValueError("Input sequence is longer than context_length") - value = self.token_embedding(idx) - if self.position_embedding is not None: - positions = torch.arange(0, token_count, dtype=torch.long, device=idx.device) - value = value + self.position_embedding(positions) - value = self.drop(value) - for block in self.blocks: - if self.gradient_checkpointing and self.training: - value = checkpoint(block, value, use_reentrant=False) - else: - value = block(value) - value = self.ln_f(value) - return self.lm_head(value) - - def forward_with_cache( - self, - idx: torch.Tensor, - past_kv: Optional[list[tuple[torch.Tensor, torch.Tensor]]] = None, - start_pos: int = 0, - ) -> tuple[torch.Tensor, list[tuple[torch.Tensor, torch.Tensor]]]: - """Run a forward pass and return updated KV cache. - - Args: - idx: Token IDs with shape ``[batch, tokens]``. - past_kv: Optional per-layer key/value cache. - start_pos: Absolute starting token position. - - Returns: - Logits and updated per-layer KV cache. - """ - - _, token_count = idx.size() - if token_count > self.config.context_length: - raise ValueError("Input sequence is longer than context_length") - value = self.token_embedding(idx) - if self.position_embedding is not None: - positions = torch.arange(start_pos, start_pos + token_count, dtype=torch.long, device=idx.device) - positions = positions.clamp(max=self.config.context_length - 1) - value = value + self.position_embedding(positions) - value = self.drop(value) - next_cache: list[tuple[torch.Tensor, torch.Tensor]] = [] - for index, block in enumerate(self.blocks): - layer_cache = past_kv[index] if past_kv is not None and index < len(past_kv) else None - value, present = block.forward_with_cache(value, past_kv=layer_cache, start_pos=start_pos) - next_cache.append(present) - value = self.ln_f(value) - return self.lm_head(value), next_cache - - @torch.no_grad() - def generate( - self, - idx: torch.Tensor, - max_new_tokens: int, - temperature: float = 0.8, - top_k: Optional[int] = 50, - use_kv_cache: bool = True, - ) -> torch.Tensor: - """Autoregressively sample new tokens. - - Args: - idx: Starting token IDs. - max_new_tokens: Number of tokens to generate. - temperature: Sampling temperature. - top_k: Optional top-k cutoff. - use_kv_cache: Whether to reuse key/value tensors during generation. - - Returns: - Token IDs including the original context and generated tokens. - """ - - if max_new_tokens <= 0: - return idx - - past_kv: Optional[list[tuple[torch.Tensor, torch.Tensor]]] = None - cached_logits: Optional[torch.Tensor] = None - if use_kv_cache: - idx_cond = idx[:, -self.config.context_length :] - cached_logits, past_kv = self.forward_with_cache(idx_cond, start_pos=0) - - for step in range(max_new_tokens): - if use_kv_cache and cached_logits is not None: - logits = cached_logits[:, -1, :] / max(temperature, 1e-5) - else: - idx_cond = idx[:, -self.config.context_length :] - logits = self(idx_cond)[:, -1, :] / max(temperature, 1e-5) - if top_k is not None: - values, _ = torch.topk(logits, min(top_k, logits.size(-1))) - logits[logits < values[:, [-1]]] = -float("inf") - probs = F.softmax(logits, dim=-1) - idx_next = torch.multinomial(probs, num_samples=1) - idx = torch.cat((idx, idx_next), dim=1) - if use_kv_cache and step < max_new_tokens - 1: - if past_kv: - cached_length = int(past_kv[0][0].size(-2)) - rolling_start = min(cached_length, self.config.context_length - 1) - cached_logits, past_kv = self.forward_with_cache( - idx[:, -1:], - past_kv=past_kv, - start_pos=rolling_start, - ) - else: - idx_cond = idx[:, -self.config.context_length :] - cached_logits, past_kv = self.forward_with_cache(idx_cond, start_pos=0) - return idx \ No newline at end of file diff --git a/llm_trainer/notifier.py b/llm_trainer/notifier.py deleted file mode 100644 index 348668c..0000000 --- a/llm_trainer/notifier.py +++ /dev/null @@ -1,358 +0,0 @@ -from __future__ import annotations - -import json -import logging -import smtplib -import threading -import time -import urllib.parse -import urllib.request -from dataclasses import asdict, dataclass, field -from email.message import EmailMessage -from pathlib import Path -from typing import Any, Optional - - -LOGGER = logging.getLogger(__name__) -TELEGRAM_API = "https://api.telegram.org/bot{token}/{method}" -MAX_TELEGRAM_TEXT = 3900 - - -@dataclass -class TelegramNotifierConfig: - """Telegram Bot API settings for progress notifications.""" - - enabled: bool = False - bot_token: str = "" - chat_id: str = "" - - -@dataclass -class EmailNotifierConfig: - """SMTP settings for email progress notifications.""" - - enabled: bool = False - smtp_host: str = "" - smtp_port: int = 587 - username: str = "" - password: str = "" - from_email: str = "" - to_email: str = "" - use_tls: bool = True - - -@dataclass -class NotifierConfig: - """Notification manager settings stored in notifier_config.json.""" - - progress_interval_seconds: int = 60 - telegram: TelegramNotifierConfig = field(default_factory=TelegramNotifierConfig) - email: EmailNotifierConfig = field(default_factory=EmailNotifierConfig) - - -def default_notifier_config_path(project_dir: Optional[Path] = None) -> Path: - """Return the notifier config path for a project or user profile. - - Args: - project_dir: Optional project root directory. - - Returns: - Path where notifier_config.json should live. - """ - - if project_dir is not None: - return project_dir / "notifier_config.json" - return Path.home() / ".drunkenbot_ide" / "notifier_config.json" - - -def ensure_notifier_config(path: Path) -> Path: - """Create notifier_config.json with disabled defaults when missing. - - Args: - path: Desired configuration file path. - - Returns: - The same configuration path. - """ - - path.parent.mkdir(parents=True, exist_ok=True) - if not path.exists(): - path.write_text(json.dumps(asdict(NotifierConfig()), indent=2), encoding="utf-8") - LOGGER.info("Created notifier config: %s", path) - return path - - -def load_notifier_config(path: Path) -> NotifierConfig: - """Load notifier settings from JSON. - - Args: - path: Configuration file path. - - Returns: - Parsed notifier configuration. - """ - - ensure_notifier_config(path) - try: - raw = json.loads(path.read_text(encoding="utf-8-sig")) - except json.JSONDecodeError as exc: - LOGGER.error("Notifier config is invalid JSON: %s", exc) - return NotifierConfig() - telegram = TelegramNotifierConfig(**dict(raw.get("telegram", {}))) - email = EmailNotifierConfig(**dict(raw.get("email", {}))) - return NotifierConfig( - progress_interval_seconds=int(raw.get("progress_interval_seconds", 60) or 60), - telegram=telegram, - email=email, - ) - - -class NotificationManager: - """Send throttled progress and completion notifications. - - Telegram progress notifications are edited in-place when possible. Email - notifications are sent as periodic snapshots because email messages cannot - be edited after delivery. - """ - - def __init__(self, config_path: Path) -> None: - """Initialize the notification manager. - - Args: - config_path: Path to notifier_config.json. - """ - - self.config_path = ensure_notifier_config(config_path) - self.config = load_notifier_config(self.config_path) - self._last_progress_at: dict[str, float] = {} - self._telegram_message_ids: dict[str, int] = {} - self._lock = threading.RLock() - self._disabled_warning_logged = False - - @property - def enabled(self) -> bool: - """Whether any notification channel is configured and enabled.""" - - return self.telegram_enabled or self.email_enabled - - @property - def telegram_enabled(self) -> bool: - """Whether Telegram notifications can be sent.""" - - telegram = self.config.telegram - return telegram.enabled and bool(telegram.bot_token.strip()) and bool(telegram.chat_id.strip()) - - @property - def email_enabled(self) -> bool: - """Whether email notifications can be sent.""" - - email = self.config.email - return ( - email.enabled - and bool(email.smtp_host.strip()) - and bool(email.from_email.strip()) - and bool(email.to_email.strip()) - ) - - def reload(self) -> None: - """Reload notifier_config.json from disk.""" - - self.config = load_notifier_config(self.config_path) - - def notify_progress(self, stage_key: str, title: str, lines: list[str], percent: Optional[int] = None) -> None: - """Send or edit a throttled progress notification. - - Args: - stage_key: Stable task key such as dataset, training, or fine_tune. - title: User-facing notification title. - lines: Body lines. - percent: Optional progress percent. - """ - - self.reload() - if not self.enabled: - self._log_disabled() - return - now = time.time() - interval = max(10, int(self.config.progress_interval_seconds or 60)) - last = self._last_progress_at.get(stage_key, 0.0) - if now - last < interval: - return - self._last_progress_at[stage_key] = now - text = self._format_message(title, lines, percent) - self._submit(lambda: self._send_progress(stage_key, f"{title} progress", text)) - - def notify_complete(self, stage_key: str, title: str, lines: list[str]) -> None: - """Send a completion summary immediately. - - Args: - stage_key: Stable task key such as dataset, training, or fine_tune. - title: User-facing notification title. - lines: Body lines. - """ - - self.reload() - if not self.enabled: - self._log_disabled() - return - text = self._format_message(title, lines, 100) - self._submit(lambda: self._send_completion(stage_key, title, text)) - - def notify_failure(self, stage_key: str, title: str, message: str) -> None: - """Send a failure summary immediately. - - Args: - stage_key: Stable task key such as dataset, training, or fine_tune. - title: User-facing notification title. - message: Failure message. - """ - - self.reload() - if not self.enabled: - self._log_disabled() - return - text = self._format_message(title, [message], None) - self._submit(lambda: self._send_completion(stage_key, title, text)) - - def _send_progress(self, stage_key: str, subject: str, text: str) -> None: - """Dispatch a progress notification to enabled channels.""" - - if self.telegram_enabled: - self._send_or_edit_telegram(stage_key, text) - if self.email_enabled: - self._send_email(subject, text) - - def _send_completion(self, stage_key: str, subject: str, text: str) -> None: - """Dispatch a final notification to enabled channels.""" - - if self.telegram_enabled: - self._send_or_edit_telegram(stage_key, text) - if self.email_enabled: - self._send_email(subject, text) - - def _send_or_edit_telegram(self, stage_key: str, text: str) -> None: - """Send a new Telegram message or edit the existing stage message.""" - - with self._lock: - message_id = self._telegram_message_ids.get(stage_key) - if message_id is None: - response = self._telegram_post( - "sendMessage", - { - "chat_id": self.config.telegram.chat_id, - "text": self._truncate(text), - "disable_web_page_preview": "true", - }, - ) - result = response.get("result", {}) if isinstance(response, dict) else {} - new_id = result.get("message_id") - if isinstance(new_id, int): - self._telegram_message_ids[stage_key] = new_id - return - try: - self._telegram_post( - "editMessageText", - { - "chat_id": self.config.telegram.chat_id, - "message_id": str(message_id), - "text": self._truncate(text), - "disable_web_page_preview": "true", - }, - ) - except Exception as exc: - LOGGER.warning("Telegram edit failed, sending a new message: %s", exc) - self._telegram_message_ids.pop(stage_key, None) - self._send_or_edit_telegram(stage_key, text) - - def _telegram_post(self, method: str, payload: dict[str, str]) -> dict[str, Any]: - """Call the Telegram Bot API. - - Args: - method: Telegram method name. - payload: Form data. - - Returns: - Parsed JSON response. - """ - - token = self.config.telegram.bot_token.strip() - url = TELEGRAM_API.format(token=urllib.parse.quote(token), method=method) - data = urllib.parse.urlencode(payload).encode("utf-8") - request = urllib.request.Request(url, data=data, method="POST") - with urllib.request.urlopen(request, timeout=15) as response: - return json.loads(response.read().decode("utf-8")) - - def _send_email(self, subject: str, text: str) -> None: - """Send an email notification through SMTP.""" - - email = self.config.email - message = EmailMessage() - message["Subject"] = f"Micro LLM Creator - {subject}" - message["From"] = email.from_email - message["To"] = email.to_email - message.set_content(text) - if email.use_tls: - smtp = smtplib.SMTP(email.smtp_host, email.smtp_port, timeout=20) - try: - smtp.starttls() - if email.username: - smtp.login(email.username, email.password) - smtp.send_message(message) - finally: - smtp.quit() - else: - with smtplib.SMTP_SSL(email.smtp_host, email.smtp_port, timeout=20) as smtp_ssl: - if email.username: - smtp_ssl.login(email.username, email.password) - smtp_ssl.send_message(message) - - def _submit(self, fn) -> None: - """Run notification delivery without blocking the UI thread.""" - - def run_safely() -> None: - try: - fn() - except Exception: - LOGGER.exception("Notification delivery failed") - - threading.Thread(target=run_safely, daemon=True).start() - - def _log_disabled(self) -> None: - """Log a one-time hint when notifications are configured off.""" - - if self._disabled_warning_logged: - return - self._disabled_warning_logged = True - telegram = self.config.telegram - email = self.config.email - LOGGER.warning( - "Notifications skipped because no channel is enabled. " - "Config=%s telegram_enabled=%s telegram_token_set=%s telegram_chat_id_set=%s " - "email_enabled=%s email_host_set=%s email_to_set=%s", - self.config_path, - telegram.enabled, - bool(telegram.bot_token.strip()), - bool(telegram.chat_id.strip()), - email.enabled, - bool(email.smtp_host.strip()), - bool(email.to_email.strip()), - ) - - @staticmethod - def _format_message(title: str, lines: list[str], percent: Optional[int]) -> str: - """Build a compact plain-text notification body.""" - - now = time.strftime("%Y-%m-%d %H:%M:%S") - body = [title, f"Time: {now}"] - if percent is not None: - body.append(f"Progress: {max(0, min(100, int(percent)))}%") - body.extend(line for line in lines if line) - return "\n".join(body) - - @staticmethod - def _truncate(text: str) -> str: - """Keep Telegram messages within the API text limit.""" - - if len(text) <= MAX_TELEGRAM_TEXT: - return text - return text[: MAX_TELEGRAM_TEXT - 40] + "\n... truncated ..." diff --git a/llm_trainer/resume_checks.py b/llm_trainer/resume_checks.py deleted file mode 100644 index 3907d4b..0000000 --- a/llm_trainer/resume_checks.py +++ /dev/null @@ -1,151 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -from typing import Any, Optional - -import torch - -from .config import ModelConfig, TrainingConfig, dataclass_to_jsonable -from .data import file_sha256 -from .lineage import read_json -from .training import latest_checkpoint - - -def _resume_checkpoint_for(training_config: TrainingConfig) -> Optional[Path]: - """Return the checkpoint that will be used for resume, if any. - - Args: - training_config: Training configuration. - - Returns: - Resume checkpoint path or ``None``. - """ - - if not training_config.resume: - return None - if training_config.resume_from_checkpoint: - return Path(training_config.resume_from_checkpoint) - return latest_checkpoint(training_config.output_dir / "checkpoints") - - -def _compatible_model_config(model_config: ModelConfig) -> dict[str, Any]: - """Return checkpoint compatibility fields for a model config. - - Args: - model_config: Current model configuration. - - Returns: - Dictionary of architecture-shape fields. - """ - - data = dataclass_to_jsonable(model_config) - return { - key: data.get(key) - for key in ( - "vocab_size", - "context_length", - "embedding_size", - "head_count", - "layer_count", - "bias", - "norm_type", - "position_encoding", - "mlp_type", - "rope_theta", - "attention_type", - ) - } - - -def _tokenizer_files_match(left: Path, right: Path) -> bool: - """Return whether tokenizer files are byte-identical or JSON-equivalent. - - Args: - left: First tokenizer path. - right: Second tokenizer path. - - Returns: - True when tokenizers are compatible. - """ - - if file_sha256(left) == file_sha256(right): - return True - left_json = read_json(left, default=None) - right_json = read_json(right, default=None) - return left_json is not None and left_json == right_json - - -def _validate_resume_compatibility( - data_dir: Path, - tokenizer_path: Path, - model_config: ModelConfig, - training_config: TrainingConfig, -) -> Optional[Path]: - """Validate tokenizer and architecture before continuing training. - - Args: - data_dir: Prepared dataset folder. - tokenizer_path: Dataset tokenizer path. - model_config: Current model architecture. - training_config: Training configuration. - - Returns: - Resume checkpoint path when one exists. - - Raises: - ValueError: If tokenizer or architecture is incompatible. - """ - - resume_path = _resume_checkpoint_for(training_config) - if not resume_path or not resume_path.exists() or not training_config.require_compatible_resume: - return resume_path - - existing_tokenizer = training_config.output_dir / "tokenizer.json" - if existing_tokenizer.exists(): - if not _tokenizer_files_match(existing_tokenizer, tokenizer_path): - raise ValueError( - "Resume safety check failed: the selected dataset tokenizer does not match the tokenizer " - f"used by the existing model folder.\nExisting tokenizer: {existing_tokenizer}\n" - f"Dataset tokenizer: {tokenizer_path}\n\nUse the same tokenizer policy for continued training, " - "or choose a new model output folder to start a new model." - ) - - checkpoint = torch.load(resume_path, map_location="cpu") - checkpoint_config = checkpoint.get("model_config") - if not isinstance(checkpoint_config, dict): - raise ValueError(f"Resume safety check failed: checkpoint has no model_config: {resume_path}") - - current = _compatible_model_config(model_config) - legacy_defaults = { - "bias": True, - "norm_type": "layernorm", - "position_encoding": "learned", - "mlp_type": "gelu", - "rope_theta": 10000.0, - "attention_type": "mha", - } - previous = {key: checkpoint_config.get(key, legacy_defaults.get(key)) for key in current} - mismatches = { - key: {"checkpoint": previous.get(key), "current": current.get(key)} - for key in current - if previous.get(key) != current.get(key) - } - if mismatches: - mismatch_text = ", ".join( - f"{key} checkpoint={value['checkpoint']} current={value['current']}" - for key, value in mismatches.items() - ) - raise ValueError( - "Resume safety check failed: model architecture does not match the checkpoint. " - f"{mismatch_text}. Keep architecture settings identical for continued training, " - "or use a new model output folder." - ) - - return resume_path - -__all__ = [ - "_resume_checkpoint_for", - "_compatible_model_config", - "_tokenizer_files_match", - "_validate_resume_compatibility", -] diff --git a/llm_trainer/runpod_cloud.py b/llm_trainer/runpod_cloud.py deleted file mode 100644 index f5d4987..0000000 --- a/llm_trainer/runpod_cloud.py +++ /dev/null @@ -1,416 +0,0 @@ -from __future__ import annotations - -import json -import logging -import textwrap -import urllib.error -import urllib.parse -import urllib.request -import zipfile -from dataclasses import asdict, dataclass -from pathlib import Path -from typing import Any, Optional - - -LOGGER = logging.getLogger(__name__) -RUNPOD_REST_BASE = "https://rest.runpod.io/v1" - - -@dataclass -class RunPodConfig: - """Settings used to launch RunPod worker Pods.""" - - api_key: str = "" - image_name: str = "runpod/pytorch:2.4.0-py3.11-cuda12.4.1-devel-ubuntu22.04" - gpu_type_id: str = "NVIDIA GeForce RTX 4090" - gpu_count: int = 1 - cloud_type: str = "COMMUNITY" - interruptible: bool = True - container_disk_gb: int = 80 - volume_gb: int = 40 - min_vcpu_per_gpu: int = 4 - min_ram_per_gpu: int = 16 - auto_terminate: bool = True - worker_labels: str = "runpod,gpu" - - -@dataclass -class RunPodLaunchResult: - """Result returned after launching a RunPod worker Pod.""" - - pod_id: str - pod_name: str - cost_per_hour: str - gpu_name: str - worker_id: str - bootstrap_url: str - - -def default_runpod_config_path(project_dir: Optional[Path] = None) -> Path: - """Return the RunPod config path. - - Args: - project_dir: Optional project root folder. - - Returns: - Project-specific or user-profile RunPod config path. - """ - - if project_dir is not None: - return project_dir / "runpod_config.json" - return Path.home() / ".drunkenbot_ide" / "runpod_config.json" - - -def ensure_runpod_config(path: Path) -> Path: - """Create a disabled/default RunPod config when missing. - - Args: - path: Config path. - - Returns: - The same config path. - """ - - path.parent.mkdir(parents=True, exist_ok=True) - if not path.exists(): - path.write_text(json.dumps(asdict(RunPodConfig()), indent=2), encoding="utf-8") - LOGGER.info("Created RunPod config: %s", path) - return path - - -def load_runpod_config(path: Path) -> RunPodConfig: - """Load RunPod settings from JSON. - - Args: - path: Config path. - - Returns: - Parsed RunPod configuration. - """ - - ensure_runpod_config(path) - data = json.loads(path.read_text(encoding="utf-8-sig")) - return RunPodConfig(**{**asdict(RunPodConfig()), **data}) - - -def save_runpod_config(path: Path, config: RunPodConfig) -> None: - """Save RunPod settings to JSON. - - Args: - path: Config path. - config: RunPod configuration. - """ - - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps(asdict(config), indent=2), encoding="utf-8") - - -def create_runpod_worker_bundle(project_root: Path, artifact_root: Path, bundle_name: str = "runpod_worker_bootstrap.zip") -> Path: - """Create a worker source bundle served by the coordinator. - - Args: - project_root: Local micro_trainer project root containing llm_trainer/. - artifact_root: Coordinator artifact root. - bundle_name: Output zip file name. - - Returns: - Path to the created bootstrap bundle. - """ - - project_root = Path(project_root).resolve() - artifact_root = Path(artifact_root) - artifact_root.mkdir(parents=True, exist_ok=True) - bundle_path = artifact_root / bundle_name - with zipfile.ZipFile(bundle_path, "w", compression=zipfile.ZIP_DEFLATED) as archive: - _write_tree(archive, project_root / "llm_trainer", "worker_src/llm_trainer") - archive.writestr("worker_src/requirements.txt", _worker_requirements(project_root)) - run_app = project_root / "run_app.py" - if run_app.exists(): - archive.write(run_app, "worker_src/run_app.py") - init_file = project_root / "__init__.py" - if init_file.exists(): - archive.write(init_file, "worker_src/__init__.py") - archive.writestr("worker_src/run_worker.py", _worker_runner_script()) - return bundle_path - - -def build_runpod_start_command( - bootstrap_url: str, - coordinator_url: str, - worker_id: str, - labels: str, - claim_once: bool, -) -> list[str]: - """Build the Pod start command that installs and runs a worker. - - Args: - bootstrap_url: Public URL for the worker source bundle. - coordinator_url: Public coordinator URL. - worker_id: Stable worker identifier. - labels: Worker labels. - claim_once: Whether the worker exits after one job. - - Returns: - Docker start command array. - """ - - script = f""" -set -e -mkdir -p /workspace/micro_llm_worker -cd /workspace/micro_llm_worker -python - <<'PY' -import urllib.request -urllib.request.urlretrieve({bootstrap_url!r}, 'worker_bootstrap.zip') -PY -python - <<'PY' -import zipfile -with zipfile.ZipFile('worker_bootstrap.zip') as archive: - archive.extractall('.') -PY -cd worker_src -python -m pip install --upgrade pip -python -m pip install -r requirements.txt -python run_worker.py --coordinator-url {coordinator_url!r} --worker-id {worker_id!r} --labels {labels!r} {'--claim-once' if claim_once else ''} -""" - return ["bash", "-lc", textwrap.dedent(script).strip()] - - -class RunPodClient: - """Small REST client for RunPod Pods.""" - - def __init__(self, api_key: str, base_url: str = RUNPOD_REST_BASE) -> None: - """Create a RunPod API client. - - Args: - api_key: RunPod API key. - base_url: REST API base URL. - """ - - self.api_key = api_key.strip() - self.base_url = base_url.rstrip("/") - - def create_worker_pod( - self, - config: RunPodConfig, - pod_name: str, - worker_id: str, - coordinator_url: str, - bootstrap_url: str, - ) -> RunPodLaunchResult: - """Create and start a RunPod worker Pod. - - Args: - config: RunPod settings. - pod_name: Pod name. - worker_id: Worker identifier. - coordinator_url: Public coordinator URL. - bootstrap_url: Public worker bootstrap bundle URL. - - Returns: - Launch summary. - """ - - if not self.api_key: - raise ValueError("RunPod API key is missing. Edit runpod_config.json first.") - payload = { - "name": pod_name, - "cloudType": config.cloud_type, - "computeType": "GPU", - "imageName": config.image_name, - "gpuCount": max(1, int(config.gpu_count)), - "gpuTypeIds": [config.gpu_type_id] if config.gpu_type_id.strip() else [], - "gpuTypePriority": "availability", - "containerDiskInGb": int(config.container_disk_gb), - "volumeInGb": int(config.volume_gb), - "volumeMountPath": "/workspace", - "minVCPUPerGPU": int(config.min_vcpu_per_gpu), - "minRAMPerGPU": int(config.min_ram_per_gpu), - "interruptible": bool(config.interruptible), - "supportPublicIp": True, - "ports": ["22/tcp"], - "env": { - "MICRO_LLM_COORDINATOR_URL": coordinator_url, - "MICRO_LLM_WORKER_ID": worker_id, - "MICRO_LLM_BOOTSTRAP_URL": bootstrap_url, - "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", - }, - "dockerStartCmd": build_runpod_start_command( - bootstrap_url=bootstrap_url, - coordinator_url=coordinator_url, - worker_id=worker_id, - labels=config.worker_labels, - claim_once=config.auto_terminate, - ), - } - response = self._request("POST", "/pods", payload) - pod_id = str(response.get("id") or "") - if not pod_id: - raise RuntimeError(f"RunPod did not return a pod id: {response}") - gpu = response.get("gpu") or response.get("machine", {}) or {} - return RunPodLaunchResult( - pod_id=pod_id, - pod_name=str(response.get("name") or pod_name), - cost_per_hour=str(response.get("costPerHr") or response.get("adjustedCostPerHr") or "-"), - gpu_name=str(gpu.get("displayName") or gpu.get("gpuDisplayName") or config.gpu_type_id), - worker_id=worker_id, - bootstrap_url=bootstrap_url, - ) - - def stop_pod(self, pod_id: str) -> dict[str, Any]: - """Stop a RunPod Pod. - - Args: - pod_id: Pod identifier. - - Returns: - RunPod response. - """ - - return self._request("POST", f"/pods/{pod_id}/stop", {}) - - def delete_pod(self, pod_id: str) -> dict[str, Any]: - """Delete a RunPod Pod. - - Args: - pod_id: Pod identifier. - - Returns: - RunPod response. - """ - - return self._request("DELETE", f"/pods/{pod_id}", None) - - def list_pods(self) -> dict[str, Any]: - """List Pods in the RunPod account. - - Returns: - RunPod response. - """ - - return self._request("GET", "/pods", None) - - def _request(self, method: str, path: str, payload: Optional[dict[str, Any]]) -> dict[str, Any]: - """Send a JSON request to RunPod. - - Args: - method: HTTP method. - path: REST path. - payload: Optional JSON payload. - - Returns: - JSON response. - """ - - data = None if payload is None else json.dumps(payload).encode("utf-8") - request = urllib.request.Request( - f"{self.base_url}{path}", - data=data, - method=method, - headers={ - "Authorization": f"Bearer {self.api_key}", - "Content-Type": "application/json", - }, - ) - try: - with urllib.request.urlopen(request, timeout=60) as response: - raw = response.read().decode("utf-8") - except urllib.error.HTTPError as exc: - detail = exc.read().decode("utf-8", errors="replace") - raise RuntimeError(f"RunPod API error {exc.code}: {detail}") from exc - if not raw: - return {} - parsed = json.loads(raw) - if isinstance(parsed, dict): - return parsed - return {"data": parsed} - - -def public_url_is_cloud_reachable(url: str) -> bool: - """Return whether a URL looks reachable from RunPod. - - Args: - url: Coordinator URL. - - Returns: - True when the URL is not localhost/private loopback. - """ - - parsed = urllib.parse.urlparse(url) - host = (parsed.hostname or "").lower() - return bool(host and host not in {"127.0.0.1", "localhost", "::1", "0.0.0.0"}) - - -def _write_tree(archive: zipfile.ZipFile, source: Path, archive_root: str) -> None: - """Write a source tree into a zip archive.""" - - source = Path(source) - if not source.exists(): - raise FileNotFoundError(f"Worker source not found: {source}") - for path in source.rglob("*"): - if path.is_file() and "__pycache__" not in path.parts: - archive.write(path, Path(archive_root) / path.relative_to(source)) - - -def _worker_runner_script() -> str: - """Return the worker runner source used inside RunPod.""" - - return """from __future__ import annotations - -import argparse -from pathlib import Path - -from llm_trainer.worker import WorkerClientConfig, run_worker_client - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--coordinator-url", required=True) - parser.add_argument("--worker-id", required=True) - parser.add_argument("--labels", default="runpod,gpu") - parser.add_argument("--workspace-dir", default="/workspace/micro_llm_worker/jobs") - parser.add_argument("--claim-once", action="store_true") - args = parser.parse_args() - labels = [item.strip() for item in args.labels.split(",") if item.strip()] - config = WorkerClientConfig( - coordinator_url=args.coordinator_url, - worker_id=args.worker_id, - labels=labels, - workspace_dir=Path(args.workspace_dir), - execute_jobs=True, - claim_once=args.claim_once, - device="cuda", - ) - run_worker_client(config) - - -if __name__ == "__main__": - main() -""" - - -def _worker_requirements(project_root: Path) -> str: - """Return lean worker requirements. - - Args: - project_root: Local project root. - - Returns: - Requirements text for cloud workers. - """ - - requirements = project_root / "requirements.txt" - if not requirements.exists(): - return "\n".join(["numpy", "torch", "tokenizers", "psutil", "datasets", "PyPDF2", "nltk"]) + "\n" - excluded = {"pyside6", "llama-cpp-python", "pyqtgraph", "markdown", "pygments"} - lines: list[str] = [] - for raw in requirements.read_text(encoding="utf-8").splitlines(): - stripped = raw.strip() - if not stripped or stripped.startswith("#"): - continue - package = stripped.split("==", 1)[0].split(">=", 1)[0].split("<=", 1)[0].lower() - if package in excluded: - continue - lines.append(stripped) - if "torch" not in {line.split("==", 1)[0].split(">=", 1)[0].split("<=", 1)[0].lower() for line in lines}: - lines.append("torch") - return "\n".join(lines) + "\n" diff --git a/llm_trainer/services.py b/llm_trainer/services.py deleted file mode 100644 index 7a2166b..0000000 --- a/llm_trainer/services.py +++ /dev/null @@ -1,19 +0,0 @@ -from __future__ import annotations - -"""Compatibility facade for split service subsystems.""" - -from .dataset_build import DatasetBuildResult, build_dataset, content_warning, estimate_vocab_size -from .dataset_preview import DatasetPreviewResult, ProjectHealthResult, check_project_health, scan_dataset_preview -from .training_orchestrator import train_from_dataset - -__all__ = [ - "DatasetBuildResult", - "ProjectHealthResult", - "DatasetPreviewResult", - "build_dataset", - "train_from_dataset", - "check_project_health", - "scan_dataset_preview", - "estimate_vocab_size", - "content_warning", -] diff --git a/llm_trainer/telemetry_store.py b/llm_trainer/telemetry_store.py deleted file mode 100644 index 003aad2..0000000 --- a/llm_trainer/telemetry_store.py +++ /dev/null @@ -1,181 +0,0 @@ -from __future__ import annotations - -import sqlite3 -import time -from pathlib import Path -from typing import Any, Optional - - -METRIC_FIELDS = ( - "epoch", - "total_epochs", - "total_steps", - "train_loss", - "val_loss", - "learning_rate", - "grad_norm", - "weight_norm", - "update_ratio", - "tokens_per_second", - "samples_per_second", - "vram_allocated_gb", - "vram_reserved_gb", - "gpu_memory_percent", - "system_cpu_percent", - "system_ram_percent", - "data_loader_workers", - "sample_text", -) - - -def telemetry_db_path(model_dir: Path) -> Path: - """Return the telemetry SQLite path for a model directory. - - Args: - model_dir: Model output directory. - - Returns: - Path to the telemetry SQLite database. - """ - - return model_dir / "training_telemetry.sqlite" - - -def ensure_schema(connection: sqlite3.Connection) -> None: - """Create or migrate the live telemetry schema. - - Args: - connection: Open SQLite connection. - """ - - connection.execute( - """ - CREATE TABLE IF NOT EXISTS live_metrics ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - run_id TEXT NOT NULL, - recorded_at REAL NOT NULL, - step INTEGER NOT NULL, - epoch INTEGER, - total_epochs INTEGER, - total_steps INTEGER, - train_loss REAL, - val_loss REAL, - learning_rate REAL, - grad_norm REAL, - weight_norm REAL, - update_ratio REAL, - tokens_per_second REAL, - samples_per_second REAL, - vram_allocated_gb REAL, - vram_reserved_gb REAL, - gpu_memory_percent REAL, - system_cpu_percent REAL, - system_ram_percent REAL, - data_loader_workers INTEGER, - sample_text TEXT - ) - """ - ) - columns = {row[1] for row in connection.execute("PRAGMA table_info(live_metrics)")} - if "sample_text" not in columns: - connection.execute("ALTER TABLE live_metrics ADD COLUMN sample_text TEXT") - connection.execute("CREATE INDEX IF NOT EXISTS idx_live_metrics_run_id_id ON live_metrics(run_id, id)") - - -def initialize_store(model_dir: Path) -> Path: - """Create a telemetry database for a model directory. - - Args: - model_dir: Model output directory. - - Returns: - Telemetry database path. - """ - - model_dir.mkdir(parents=True, exist_ok=True) - db_path = telemetry_db_path(model_dir) - with sqlite3.connect(db_path) as connection: - ensure_schema(connection) - connection.commit() - return db_path - - -def insert_metric(db_path: Path, run_id: str, event: dict[str, Any]) -> int: - """Persist one training metric event. - - Args: - db_path: Telemetry database path. - run_id: Training run identifier. - event: Progress event emitted by training. - - Returns: - Inserted row id. - """ - - values = [event.get(field) for field in METRIC_FIELDS] - with sqlite3.connect(db_path) as connection: - ensure_schema(connection) - cursor = connection.execute( - """ - INSERT INTO live_metrics ( - run_id, recorded_at, step, epoch, total_epochs, total_steps, - train_loss, val_loss, learning_rate, grad_norm, weight_norm, update_ratio, - tokens_per_second, samples_per_second, vram_allocated_gb, vram_reserved_gb, - gpu_memory_percent, system_cpu_percent, system_ram_percent, data_loader_workers, - sample_text - ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """, - [run_id, time.time(), int(event["step"]), *values], - ) - connection.commit() - return int(cursor.lastrowid) - - -def latest_run(db_path: Path) -> Optional[sqlite3.Row]: - """Return latest telemetry run metadata. - - Args: - db_path: Telemetry database path. - - Returns: - Row with run id, sample count, and latest row id, or None. - """ - - with sqlite3.connect(db_path) as connection: - ensure_schema(connection) - connection.row_factory = sqlite3.Row - return connection.execute( - """ - SELECT run_id, COUNT(*) AS sample_count, MAX(id) AS latest_id - FROM live_metrics - GROUP BY run_id - ORDER BY MAX(id) DESC - LIMIT 1 - """ - ).fetchone() - - -def rows_until(db_path: Path, run_id: str, sample_index: int) -> list[sqlite3.Row]: - """Load telemetry rows up to a selected sample. - - Args: - db_path: Telemetry database path. - run_id: Training run identifier. - sample_index: Maximum row count to load. - - Returns: - Ordered telemetry rows. - """ - - if sample_index <= 0: - return [] - with sqlite3.connect(db_path) as connection: - ensure_schema(connection) - connection.row_factory = sqlite3.Row - return list( - connection.execute( - "SELECT * FROM live_metrics WHERE run_id = ? ORDER BY id LIMIT ?", - (run_id, int(sample_index)), - ) - ) diff --git a/llm_trainer/tokenizer.py b/llm_trainer/tokenizer.py deleted file mode 100644 index 1997bea..0000000 --- a/llm_trainer/tokenizer.py +++ /dev/null @@ -1,481 +0,0 @@ -from __future__ import annotations - -import json -from pathlib import Path -from random import Random -from typing import Callable, Iterator, Optional - -import numpy as np -import numpy.lib.format as npy_format -from tokenizers import Tokenizer -from tokenizers.decoders import ByteLevel as ByteLevelDecoder -from tokenizers.models import BPE -from tokenizers.pre_tokenizers import ByteLevel -from tokenizers.processors import TemplateProcessing -from tokenizers.trainers import BpeTrainer - - -PAD_TOKEN = "" -UNK_TOKEN = "" -BOS_TOKEN = "" -EOS_TOKEN = "" -SPECIAL_TOKENS = [PAD_TOKEN, UNK_TOKEN, BOS_TOKEN, EOS_TOKEN] -DEFAULT_CHAT_TEMPLATE = """{% for message in messages %}{{ '' if loop.first else '' }}{{ message['role'] | capitalize }}: {{ message['content'] }}{{ '' if loop.last else '\\n' }}{% endfor %}{% if add_generation_prompt %}{{ '\\nAssistant:' }}{% endif %}""" -MAX_TOKENIZER_LINE_CHARS = 8_192 -# The Rust BPE trainer builds an in-memory pretoken frequency table sized to -# whatever corpus it is shown. Vocabulary quality saturates well before a -# multi-gigabyte corpus is fully consumed, so by default only a bounded, -# evenly-spread sample of the corpus is shown to the trainer. The full -# corpus is still encoded with the resulting tokenizer afterward -- only -# *training* the merges is sampled. -DEFAULT_TOKENIZER_TRAINING_MAX_BYTES = 2 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024 # 12 GiB -TOKENIZER_SAMPLE_SEED = 1337 -# Tokens are streamed to disk in fixed-size batches rather than accumulated -# into one giant Python list, so peak RAM during encoding stays roughly -# constant regardless of corpus size. -ENCODE_FLUSH_TOKEN_COUNT = 200_000 -# Number of line-chunks encoded per tokenizer.encode_batch() call. The Rust -# tokenizers library parallelizes encode_batch() internally across CPU -# cores (via Rayon); calling encode() one string at a time from a Python -# loop -- the previous approach -- pays Python/Rust boundary overhead per -# call and never engages that internal parallelism, which dominates total -# time at billions-of-tokens scale. -ENCODE_BATCH_SIZE = 1_000 -# Chunk size (in tokens) used when streaming raw token bytes into the final -# .npy file. Keeps the .bin -> .npy conversion step's RAM use flat too. -NPY_CONVERT_CHUNK_TOKENS = 1_000_000 - - -def train_tokenizer( - corpus_path: Path, - output_path: Path, - vocab_size: int = 8000, - min_frequency: int = 2, - should_stop: Optional[Callable[[], bool]] = None, - max_training_bytes: Optional[int] = DEFAULT_TOKENIZER_TRAINING_MAX_BYTES, -) -> Tokenizer: - """Train a byte-level BPE tokenizer. - - Args: - corpus_path: Text corpus used for tokenizer training. - output_path: Destination tokenizer JSON path. - vocab_size: Target vocabulary size. - min_frequency: Minimum token frequency for BPE merges. - should_stop: Optional callback returning true when training should stop. - max_training_bytes: Maximum corpus bytes shown to the BPE trainer. - The trainer keeps a frequency table in memory sized to whatever it - is shown, so on very large corpora only an evenly-spread sample - up to this many bytes is used to fit merges. Pass ``None`` to - disable sampling and train on the entire corpus. The full corpus - is always encoded with the resulting tokenizer regardless of this - setting. - - Returns: - Trained tokenizer instance. - """ - - tokenizer = Tokenizer(BPE(unk_token=UNK_TOKEN)) - tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False) - tokenizer.decoder = ByteLevelDecoder() - - trainer = BpeTrainer( - vocab_size=vocab_size, - min_frequency=min_frequency, - special_tokens=SPECIAL_TOKENS, - initial_alphabet=ByteLevel.alphabet(), - show_progress=True, - ) - corpus_size = corpus_path.stat().st_size - sample_ratio = 1.0 - if max_training_bytes is not None and corpus_size > max_training_bytes: - sample_ratio = max_training_bytes / corpus_size - tokenizer.train_from_iterator( - _iter_corpus_lines(corpus_path, should_stop, sample_ratio=sample_ratio), - trainer=trainer, - ) - tokenizer.post_processor = TemplateProcessing( - single=f"{BOS_TOKEN} $A {EOS_TOKEN}", - special_tokens=[ - (BOS_TOKEN, tokenizer.token_to_id(BOS_TOKEN)), - (EOS_TOKEN, tokenizer.token_to_id(EOS_TOKEN)), - ], - ) - - output_path.parent.mkdir(parents=True, exist_ok=True) - tokenizer.save(str(output_path)) - save_tokenizer_package(tokenizer, output_path) - return tokenizer - - -def save_tokenizer_package( - tokenizer: Tokenizer, - tokenizer_path: Path, - model_max_length: Optional[int] = None, -) -> None: - """Write standard tokenizer metadata beside the native tokenizer JSON.""" - - tokenizer_path = Path(tokenizer_path) - tokens = { - "bos_token": BOS_TOKEN, - "eos_token": EOS_TOKEN, - "unk_token": UNK_TOKEN, - "pad_token": PAD_TOKEN, - } - (tokenizer_path.parent / "special_tokens_map.json").write_text( - json.dumps(tokens, indent=2) + "\n", encoding="utf-8" - ) - config = { - "tokenizer_class": "PreTrainedTokenizerFast", - "tokenizer_file": tokenizer_path.name, - "model_max_length": int(model_max_length) if model_max_length else 1_000_000_000_000_000_000_000_000_000_000, - "clean_up_tokenization_spaces": False, - "add_bos_token": True, - "add_eos_token": True, - "chat_template": DEFAULT_CHAT_TEMPLATE, - **tokens, - } - (tokenizer_path.parent / "tokenizer_config.json").write_text( - json.dumps(config, indent=2) + "\n", encoding="utf-8" - ) - - -def _iter_corpus_lines( - corpus_path: Path, - should_stop: Optional[Callable[[], bool]], - sample_ratio: float = 1.0, -) -> Iterator[str]: - """Yield corpus lines and check for cancellation between chunks. - - Args: - corpus_path: Text corpus path to read line by line. - should_stop: Optional callback returning true when reading should stop. - sample_ratio: Fraction of lines to keep, in ``(0.0, 1.0]``. Lines are - kept via an independent per-line random draw (not a positional - head-of-file cut), so the kept sample is spread evenly across the - whole file rather than biased toward whatever content appears - first. - - Raises: - RuntimeError: If cancellation is requested. - """ - - sample_ratio = max(0.0, min(1.0, sample_ratio)) - sampling = sample_ratio < 1.0 - rng = Random(TOKENIZER_SAMPLE_SEED) if sampling else None - - with corpus_path.open("r", encoding="utf-8") as handle: - for line in handle: - if should_stop and should_stop(): - raise RuntimeError("Dataset preparation stopped by user.") - if sampling and rng.random() > sample_ratio: - continue - - for start in range(0, len(line), MAX_TOKENIZER_LINE_CHARS): - chunk = line[start : start + MAX_TOKENIZER_LINE_CHARS] - if chunk: - yield chunk - - -def load_tokenizer(path: Path) -> Tokenizer: - """Load a tokenizer from disk. - - Args: - path: Tokenizer JSON path. - - Returns: - Loaded tokenizer. - """ - - return Tokenizer.from_file(str(path)) - - -def token_id(tokenizer: Tokenizer, token: str) -> int: - """Return the integer ID for a required special token. - - Args: - tokenizer: Tokenizer to query. - token: Token string to find. - - Returns: - Token ID. - - Raises: - ValueError: If the tokenizer does not contain the token. - """ - - value = tokenizer.token_to_id(token) - if value is None: - raise ValueError(f"Tokenizer is missing required token: {token}") - return value - - -def missing_training_special_tokens(tokenizer: Tokenizer) -> list[str]: - """Return special tokens missing from a tokenizer. - - Args: - tokenizer: Tokenizer to inspect. - - Returns: - Missing special token strings. - """ - - return [token for token in SPECIAL_TOKENS if tokenizer.token_to_id(token) is None] - - -def validate_training_tokenizer(tokenizer: Tokenizer) -> None: - """Validate that a tokenizer can be used by the MicroGPT trainer. - - Args: - tokenizer: Tokenizer to validate. - - Raises: - ValueError: If required special tokens are missing. - """ - - missing = missing_training_special_tokens(tokenizer) - if missing: - raise ValueError( - "Tokenizer is not compatible with Micro LLM Creator training. " - f"Missing required special token(s): {', '.join(missing)}. " - "Use a tokenizer created by this app, or import a tokenizer containing " - ", , , and ." - ) - - -def encode_text(tokenizer: Tokenizer, text: str) -> list[int]: - """Encode text into token IDs. - - Args: - tokenizer: Tokenizer used for encoding. - text: Text to encode. - - Returns: - List of token IDs. - """ - - token_ids: list[int] = [] - for start in range(0, len(text), MAX_TOKENIZER_LINE_CHARS): - chunk = text[start : start + MAX_TOKENIZER_LINE_CHARS] - if chunk: - token_ids.extend(tokenizer.encode(chunk).ids) - return token_ids - - -def encode_file( - tokenizer: Tokenizer, - corpus_path: Path, - should_stop: Optional[Callable[[], bool]] = None, -) -> list[int]: - """Encode a corpus file into token IDs without loading it all at once. - - Args: - tokenizer: Tokenizer used for encoding. - corpus_path: Text corpus path. - should_stop: Optional callback returning true when encoding should stop. - - Returns: - Token IDs for the corpus. - - Raises: - RuntimeError: If cancellation is requested. - """ - - token_ids: list[int] = [] - with corpus_path.open("r", encoding="utf-8") as handle: - for line in handle: - if should_stop and should_stop(): - raise RuntimeError("Dataset preparation stopped by user.") - for start in range(0, len(line), MAX_TOKENIZER_LINE_CHARS): - chunk = line[start : start + MAX_TOKENIZER_LINE_CHARS] - if chunk: - token_ids.extend(tokenizer.encode(chunk).ids) - return token_ids - - -def token_dtype_for_vocab(vocab_size: int) -> np.dtype: - """Pick the smallest unsigned integer dtype that can hold every token ID. - - Args: - vocab_size: Tokenizer vocabulary size. - - Returns: - ``uint16`` for the (overwhelmingly common) case of a vocab under - 65,536, otherwise ``uint32``. - """ - - return np.dtype(np.uint16) if vocab_size <= 65_535 else np.dtype(np.uint32) - - -def encode_file_to_bin( - tokenizer: Tokenizer, - corpus_path: Path, - output_path: Path, - dtype: np.dtype, - should_stop: Optional[Callable[[], bool]] = None, -) -> int: - """Encode a corpus file straight to a flat, headerless binary token file. - - Unlike ``encode_file``, this never accumulates the whole corpus as a - Python list of ints in memory. Token IDs are buffered in small batches - and flushed to disk as raw fixed-width integers (``dtype``), so peak RAM - stays roughly constant no matter how large the corpus is. This is an - intermediate format (no shape/dtype header) -- see ``encode_file_to_npy`` - for the version that produces a directly loadable ``.npy`` file. - - Line-chunks are also batched into groups of ``ENCODE_BATCH_SIZE`` and - encoded with a single ``tokenizer.encode_batch()`` call per group rather - than one ``tokenizer.encode()`` call per chunk -- the Rust tokenizer - parallelizes ``encode_batch()`` across CPU cores internally, which - matters enormously at multi-billion-token scale. - - Args: - tokenizer: Tokenizer used for encoding. - corpus_path: Text corpus path. - output_path: Destination raw ``.bin`` path for the encoded token stream. - dtype: Integer dtype to store each token ID as (see - ``token_dtype_for_vocab``). - should_stop: Optional callback returning true when encoding should stop. - - Returns: - Total number of tokens written. - - Raises: - RuntimeError: If cancellation is requested. - """ - - output_path.parent.mkdir(parents=True, exist_ok=True) - buffer: list[int] = [] - pending_chunks: list[str] = [] - total_tokens = 0 - - def flush_pending_chunks(sink) -> None: - """Encode any buffered chunks as one batch and append their IDs.""" - - nonlocal total_tokens - if pending_chunks: - for encoding in tokenizer.encode_batch(pending_chunks): - buffer.extend(encoding.ids) - pending_chunks.clear() - if len(buffer) >= ENCODE_FLUSH_TOKEN_COUNT: - np.asarray(buffer, dtype=dtype).tofile(sink) - total_tokens += len(buffer) - buffer.clear() - - with corpus_path.open("r", encoding="utf-8") as source, output_path.open("wb") as sink: - for line in source: - if should_stop and should_stop(): - raise RuntimeError("Dataset preparation stopped by user.") - for start in range(0, len(line), MAX_TOKENIZER_LINE_CHARS): - chunk = line[start : start + MAX_TOKENIZER_LINE_CHARS] - if not chunk: - continue - pending_chunks.append(chunk) - if len(pending_chunks) >= ENCODE_BATCH_SIZE: - flush_pending_chunks(sink) - flush_pending_chunks(sink) - if buffer: - np.asarray(buffer, dtype=dtype).tofile(sink) - total_tokens += len(buffer) - return total_tokens - - -def token_count_in_bin(path: Path, dtype: np.dtype) -> int: - """Return how many tokens a raw ``.bin`` file holds, without loading it. - - Args: - path: Raw token ``.bin`` file path. - dtype: Integer dtype the tokens were stored as. - - Returns: - Number of tokens in the file. - """ - - itemsize = np.dtype(dtype).itemsize - return path.stat().st_size // itemsize - - -def load_token_memmap(path: Path, dtype: np.dtype) -> np.memmap: - """Open a raw token ``.bin`` file as a read-only memory-mapped array. - - Args: - path: Raw token ``.bin`` file path. - dtype: Integer dtype the tokens were stored as. - - Returns: - Read-only memory-mapped array of token IDs. - """ - - return np.memmap(path, dtype=dtype, mode="r") - - -def convert_bin_to_npy(bin_path: Path, npy_path: Path, dtype: np.dtype, token_count: int) -> None: - """Convert a flat raw token ``.bin`` file into a self-describing ``.npy`` file. - - Only a small, fixed-size ``.npy`` header (built from ``dtype`` and - ``token_count``, which ``encode_file_to_bin`` already gave us for free) - needs to be known upfront. The token data itself is streamed across in - fixed-size chunks via plain file reads/writes -- the full token array is - never materialized in memory during conversion, regardless of how large - the file is. The result is byte-for-byte a normal ``.npy`` file, loadable - with ``numpy.load(path, mmap_mode="r")`` like any other. - - Args: - bin_path: Source raw ``.bin`` file (as written by ``encode_file_to_bin``). - npy_path: Destination ``.npy`` file path. - dtype: Integer dtype the tokens were stored as. - token_count: Number of tokens in ``bin_path``. - """ - - npy_path.parent.mkdir(parents=True, exist_ok=True) - header = { - "descr": npy_format.dtype_to_descr(np.dtype(dtype)), - "fortran_order": False, - "shape": (token_count,), - } - chunk_bytes = NPY_CONVERT_CHUNK_TOKENS * np.dtype(dtype).itemsize - with bin_path.open("rb") as source, npy_path.open("wb") as sink: - npy_format.write_array_header_1_0(sink, header) - while True: - block = source.read(chunk_bytes) - if not block: - break - sink.write(block) - - -def encode_file_to_npy( - tokenizer: Tokenizer, - corpus_path: Path, - output_path: Path, - dtype: np.dtype, - should_stop: Optional[Callable[[], bool]] = None, -) -> int: - """Encode a corpus file straight to a memory-map-friendly ``.npy`` file. - - Combines ``encode_file_to_bin`` (single-pass streaming encode, constant - RAM) with ``convert_bin_to_npy`` (header + streamed byte copy, also - constant RAM) so the whole corpus is never held in memory as a Python - list or a full array at any point, while still producing a standard - ``.npy`` file that ``numpy.load(path, mmap_mode="r")`` opens directly. - - Args: - tokenizer: Tokenizer used for encoding. - corpus_path: Text corpus path. - output_path: Destination ``.npy`` path for the encoded token stream. - dtype: Integer dtype to store each token ID as (see - ``token_dtype_for_vocab``). - should_stop: Optional callback returning true when encoding should stop. - - Returns: - Total number of tokens written. - - Raises: - RuntimeError: If cancellation is requested. - """ - - temp_bin_path = output_path.with_suffix(output_path.suffix + ".raw_tmp") - try: - token_count = encode_file_to_bin(tokenizer, corpus_path, temp_bin_path, dtype, should_stop=should_stop) - convert_bin_to_npy(temp_bin_path, output_path, dtype, token_count) - finally: - temp_bin_path.unlink(missing_ok=True) - return token_count \ No newline at end of file diff --git a/llm_trainer/training.py b/llm_trainer/training.py deleted file mode 100644 index cea2920..0000000 --- a/llm_trainer/training.py +++ /dev/null @@ -1,1497 +0,0 @@ -from __future__ import annotations - -import json -import math -import os -import random -from dataclasses import dataclass -from pathlib import Path -from time import perf_counter -from typing import Any, Callable, Optional, Union - -import numpy as np -import numpy.lib.format as npy_format -import torch -import torch.nn.functional as F -from torch.amp import GradScaler, autocast -from torch.utils.data import DataLoader, Dataset - -from .config import ModelConfig, TrainingConfig, dataclass_to_jsonable -from .model import ( - MicroGPT, - apply_lora_adapters, - freeze_non_lora_parameters, - load_lora_state_dict, - lora_parameter_count, - lora_state_dict, - merge_lora_adapters, -) - -try: - import psutil -except ImportError: - psutil = None - - -class Lion(torch.optim.Optimizer): - """Lion optimizer with decoupled weight decay. - - The implementation follows the common Lion update rule and keeps the - optimizer self-contained so the app does not require an extra dependency. - """ - - def __init__( - self, - params, - lr: float = 1e-4, - betas: tuple[float, float] = (0.9, 0.99), - weight_decay: float = 0.0, - ) -> None: - """Create a Lion optimizer. - - Args: - params: Iterable of parameters to optimize. - lr: Learning rate. - betas: Momentum coefficients. - weight_decay: Decoupled weight decay. - """ - - if lr <= 0.0: - raise ValueError("lr must be greater than 0") - if not 0.0 <= betas[0] < 1.0 or not 0.0 <= betas[1] < 1.0: - raise ValueError("betas must be in [0, 1)") - defaults = {"lr": lr, "betas": betas, "weight_decay": weight_decay} - super().__init__(params, defaults) - - @torch.no_grad() - def step(self, closure=None): - """Perform one optimization step. - - Args: - closure: Optional closure that reevaluates the model. - - Returns: - Closure loss when a closure is provided. - """ - - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - for group in self.param_groups: - lr = group["lr"] - beta1, beta2 = group["betas"] - weight_decay = group["weight_decay"] - for parameter in group["params"]: - if parameter.grad is None: - continue - grad = parameter.grad - if weight_decay: - parameter.mul_(1.0 - lr * weight_decay) - state = self.state[parameter] - if len(state) == 0: - state["exp_avg"] = torch.zeros_like(parameter) - exp_avg = state["exp_avg"] - update = exp_avg.mul(beta1).add(grad, alpha=1.0 - beta1) - parameter.add_(update.sign(), alpha=-lr) - exp_avg.mul_(beta2).add_(grad, alpha=1.0 - beta2) - return loss - - -class TokenDataset(Dataset): - """Sliding-window token dataset for next-token prediction. - - Accepts a list of ints, a numpy array, or a numpy memmap. When backed by a - memmap the full token stream lives on disk — only individual windows are - loaded into RAM on each ``__getitem__`` call, so datasets of any size can be - used without exhausting system memory. - """ - - def __init__(self, tokens: Union[list[int], np.ndarray], context_length: int, stride: int = 1) -> None: - """Create a token dataset. - - Args: - tokens: Complete token stream (list, ndarray, or memmap). - context_length: Number of input tokens per sample. - stride: Token offset step between consecutive windows. - - Raises: - ValueError: If there are not enough tokens. - """ - - if len(tokens) <= context_length: - raise ValueError("Not enough tokens for the selected context length") - if stride <= 0: - raise ValueError("stride must be greater than 0") - # Keep the backing store as-is (memmap stays on disk). - if isinstance(tokens, np.ndarray): - self._tokens_np: Optional[np.ndarray] = tokens - self._tokens_tensor: Optional[torch.Tensor] = None - else: - self._tokens_np = None - self._tokens_tensor = torch.tensor(tokens, dtype=torch.long) - self.context_length = context_length - self.stride = stride - available_windows = len(tokens) - self.context_length - self.sample_count = (available_windows + self.stride - 1) // self.stride - - def __len__(self) -> int: - """Return the number of sliding windows available. - - Returns: - Dataset length. - """ - - return self.sample_count - - def __getitem__(self, index: int) -> tuple[torch.Tensor, torch.Tensor]: - """Return one input/target token window. - - Args: - index: Starting token index. - - Returns: - Pair of input tokens and next-token targets. - """ - - start = index * self.stride - end = start + self.context_length + 1 - if self._tokens_np is not None: - # Read the slice from the numpy array / memmap and convert to tensor. - chunk = torch.from_numpy(np.array(self._tokens_np[start:end], dtype=np.int64)) - else: - chunk = self._tokens_tensor[start:end] # type: ignore[index] - return chunk[:-1], chunk[1:] - - -@dataclass -class TrainingResult: - """Result returned after training. - - Attributes: - checkpoint_path: Final model checkpoint path. - summary_path: Training summary JSON path. - final_train_loss: Final epoch training loss. - final_val_loss: Final validation loss when available. - stopped: Whether training was stopped by the user. - """ - - checkpoint_path: Path - summary_path: Path - final_train_loss: float - final_val_loss: Optional[float] - stopped: bool = False - - -@dataclass -class ResumeCompatibilityReport: - """Compatibility result for a checkpoint resume attempt. - - Attributes: - checkpoint_path: Checkpoint path that was inspected. - errors: Blocking compatibility problems. - warnings: Non-blocking but important differences. - info: Informational compatibility details. - can_load_optimizer_state: Whether optimizer state can be safely loaded. - can_load_scheduler_state: Whether scheduler state can be safely loaded. - can_load_scaler_state: Whether AMP scaler state can be safely loaded. - """ - - checkpoint_path: Path - errors: list[str] - warnings: list[str] - info: list[str] - can_load_optimizer_state: bool = True - can_load_scheduler_state: bool = True - can_load_scaler_state: bool = True - - -def emit_progress( - progress: Optional[Callable[[Any], None]], - message: str, - percent: Optional[int] = None, - **metrics: Any, -) -> None: - """Emit training progress if a callback is available. - - Args: - progress: Optional callback for progress dictionaries. - message: Human-readable status message. - percent: Optional progress percentage. - **metrics: Optional structured metrics for UI dashboards. - """ - - if progress: - progress({"message": message, "percent": percent, **metrics}) - - -def set_seed(seed: int) -> None: - """Set random seeds for repeatable training. - - Args: - seed: Integer seed value. - """ - - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(seed) - - -def split_tokens( - tokens: list[int], - validation_split: float, - chunk_size: int = 2048, - seed: int = 1337, -) -> tuple[list[int], list[int]]: - """Split tokens into train and validation streams. - - The corpus is written to disk as one big concatenation of source - documents, then tokenized into a single flat stream. A plain positional - split would make validation depend on whichever source happened to be at - the tail of the corpus. This chunks and deterministically shuffles the - stream first, so validation samples are drawn from across the corpus. - - Args: - tokens: Full token stream. - validation_split: Fraction reserved for validation. - chunk_size: Number of tokens per shuffle unit. - seed: Fixed seed for reproducible train/validation assignment. - - Returns: - Pair of training tokens and validation tokens. - """ - - total = len(tokens) - if total <= 1 or validation_split <= 0: - return list(tokens), [] - if validation_split >= 1: - return [], list(tokens) - - chunk_size = max(1, chunk_size) - chunk_ranges = [(start, min(start + chunk_size, total)) for start in range(0, total, chunk_size)] - if len(chunk_ranges) <= 1: - split_at = int(total * (1.0 - validation_split)) - split_at = max(1, min(split_at, total - 1)) - return tokens[:split_at], tokens[split_at:] - - shuffled_indices = list(range(len(chunk_ranges))) - random.Random(seed).shuffle(shuffled_indices) - val_chunk_count = max(1, round(len(chunk_ranges) * validation_split)) - val_chunk_count = min(val_chunk_count, len(chunk_ranges) - 1) - val_chunk_indices = set(shuffled_indices[:val_chunk_count]) - - train_tokens: list[int] = [] - val_tokens: list[int] = [] - for chunk_index, (start, end) in enumerate(chunk_ranges): - piece = tokens[start:end] - if chunk_index in val_chunk_indices: - val_tokens.extend(piece) - else: - train_tokens.extend(piece) - return train_tokens, val_tokens - - -def split_tokens_to_files( - tokens: np.memmap, - train_path: Path, - val_path: Path, - validation_split: float, - dtype: np.dtype, - chunk_size: int = 2048, - seed: int = 1337, - should_stop: Optional[Callable[[], bool]] = None, -) -> tuple[int, int]: - """Split a token stream into train/validation ``.npy`` files on disk. - - Behaves like :func:`split_tokens` (same chunked, seeded shuffle so - validation samples are drawn from across the corpus, not just the tail), - but never materializes the full train or validation token stream in - memory. ``tokens`` is expected to be a read-only memmap (or any - ``__len__``/slice-able array) backed by disk; each chunk is read, cast to - ``dtype``, and written straight to the appropriate output file. Peak - memory use is therefore bounded by ``chunk_size`` regardless of corpus - size. - - Args: - tokens: Full token stream, typically a memory-mapped ``.npy`` array. - train_path: Destination path for the training token ``.npy`` file. - val_path: Destination path for the validation token ``.npy`` file. - validation_split: Fraction of chunks reserved for validation. - dtype: Integer dtype to store each token ID as (see - ``tokenizer.token_dtype_for_vocab``). - chunk_size: Number of tokens per shuffle unit. - seed: Fixed seed for reproducible train/validation assignment. - should_stop: Optional callback returning true when the split should - stop early. - - Returns: - Pair of ``(train_token_count, val_token_count)``. - - Raises: - RuntimeError: If cancellation is requested. - """ - - total = len(tokens) - chunk_size = max(1, chunk_size) - validation_split = max(0.0, min(1.0, validation_split)) - - if total <= 1 or validation_split <= 0: - chunk_ranges: list[tuple[int, int]] = [(0, total)] - val_chunk_indices: set[int] = set() - elif validation_split >= 1: - chunk_ranges = [(0, total)] - val_chunk_indices = {0} - else: - chunk_ranges = [ - (start, min(start + chunk_size, total)) for start in range(0, total, chunk_size) - ] - if len(chunk_ranges) <= 1: - split_at = int(total * (1.0 - validation_split)) - split_at = max(1, min(split_at, total - 1)) - chunk_ranges = [(0, split_at), (split_at, total)] - val_chunk_indices = {1} - else: - shuffled_indices = list(range(len(chunk_ranges))) - random.Random(seed).shuffle(shuffled_indices) - val_chunk_count = max(1, round(len(chunk_ranges) * validation_split)) - val_chunk_count = min(val_chunk_count, len(chunk_ranges) - 1) - val_chunk_indices = set(shuffled_indices[:val_chunk_count]) - - train_token_count = sum( - end - start for index, (start, end) in enumerate(chunk_ranges) if index not in val_chunk_indices - ) - val_token_count = sum( - end - start for index, (start, end) in enumerate(chunk_ranges) if index in val_chunk_indices - ) - - train_path.parent.mkdir(parents=True, exist_ok=True) - val_path.parent.mkdir(parents=True, exist_ok=True) - train_header = { - "descr": npy_format.dtype_to_descr(np.dtype(dtype)), - "fortran_order": False, - "shape": (train_token_count,), - } - val_header = { - "descr": npy_format.dtype_to_descr(np.dtype(dtype)), - "fortran_order": False, - "shape": (val_token_count,), - } - with train_path.open("wb") as train_file, val_path.open("wb") as val_file: - npy_format.write_array_header_1_0(train_file, train_header) - npy_format.write_array_header_1_0(val_file, val_header) - for chunk_index, (start, end) in enumerate(chunk_ranges): - if should_stop and should_stop(): - raise RuntimeError("Dataset preparation stopped by user.") - piece = np.asarray(tokens[start:end], dtype=dtype) - if chunk_index in val_chunk_indices: - piece.tofile(val_file) - else: - piece.tofile(train_file) - - return train_token_count, val_token_count - - -def make_optimizer(model: MicroGPT, training_config: TrainingConfig) -> torch.optim.Optimizer: - """Create the configured optimizer. - - Args: - model: Model whose parameters will be optimized. - training_config: Training configuration. - - Returns: - Configured optimizer. - - Raises: - ValueError: If the optimizer is unsupported by the installed PyTorch. - """ - - name = training_config.optimizer_name - common = { - "lr": training_config.learning_rate, - "weight_decay": training_config.weight_decay, - } - parameters = [parameter for parameter in model.parameters() if parameter.requires_grad] - if not parameters: - raise ValueError("No trainable parameters are available for optimization") - if name == "adamw": - return torch.optim.AdamW(parameters, betas=(0.9, 0.95), **common) - if name == "adam": - return torch.optim.Adam(parameters, betas=(0.9, 0.95), **common) - if name == "lion": - return Lion(parameters, betas=(0.9, 0.99), **common) - if name == "adafactor": - adafactor = getattr(torch.optim, "Adafactor", None) - if adafactor is None: - raise ValueError("Adafactor requires a newer PyTorch build that includes torch.optim.Adafactor") - return adafactor(parameters, **common) - raise ValueError(f"Unsupported optimizer: {name}") - - -def make_scheduler( - optimizer: torch.optim.Optimizer, - total_steps: int, - training_config: TrainingConfig, -) -> torch.optim.lr_scheduler.LambdaLR: - """Create the configured learning-rate scheduler. - - Args: - optimizer: Optimizer to schedule. - total_steps: Total optimizer steps. - training_config: Training configuration. - - Returns: - Lambda learning-rate scheduler. - """ - - warmup_steps = training_config.warmup_steps - warmup_steps = min(warmup_steps, max(total_steps - 1, 1)) - min_ratio = training_config.scheduler_min_lr_ratio - schedule = training_config.scheduler_name - - def lr_lambda(step: int) -> float: - if step < warmup_steps: - return max(step, 1) / max(warmup_steps, 1) - if schedule == "constant": - return 1.0 - progress = (step - warmup_steps) / max(total_steps - warmup_steps, 1) - progress = max(0.0, min(progress, 1.0)) - if schedule == "cosine": - value = 0.5 * (1.0 + math.cos(math.pi * progress)) - return min_ratio + (1.0 - min_ratio) * value - if schedule == "polynomial": - value = (1.0 - progress) ** training_config.polynomial_power - return min_ratio + (1.0 - min_ratio) * value - if schedule == "one_cycle": - if progress < 0.3: - return min_ratio + (1.0 - min_ratio) * (progress / 0.3) - decay_progress = (progress - 0.3) / 0.7 - value = 0.5 * (1.0 + math.cos(math.pi * decay_progress)) - return min_ratio + (1.0 - min_ratio) * value - return max(min_ratio, 1.0 - progress) - - return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) - - -def amp_settings(training_config: TrainingConfig) -> tuple[bool, bool, torch.dtype]: - """Return autocast and scaler settings for the selected precision. - - Args: - training_config: Training configuration. - - Returns: - Tuple of ``use_autocast``, ``use_scaler``, and autocast dtype. - """ - - use_cuda_amp = training_config.use_amp and training_config.device == "cuda" - if not use_cuda_amp or training_config.precision == "fp32": - return False, False, torch.float32 - if training_config.precision == "bf16": - return True, False, torch.bfloat16 - return True, True, torch.float16 - - -def system_ram_percent() -> Optional[float]: - """Return system RAM utilization when psutil is available. - - Returns: - RAM utilization percentage, or None when unavailable. - """ - - if psutil is None: - return None - return float(psutil.virtual_memory().percent) - - -def system_cpu_percent() -> Optional[float]: - """Return system CPU utilization when psutil is available. - - Returns: - CPU utilization percentage, or None when unavailable. - """ - - if psutil is None: - return None - return float(psutil.cpu_percent(interval=None)) - - -def evaluate( - model: MicroGPT, - loader: DataLoader, - device: str, - pad_token_id: int, - max_batches: int = 50, - progress: Optional[Callable[[Any], None]] = None, - should_stop: Optional[Callable[[], bool]] = None, - step: Optional[int] = None, - total_steps: Optional[int] = None, - percent: Optional[int] = None, -) -> float: - """Evaluate validation loss. - - Args: - model: Model to evaluate. - loader: Validation data loader. - device: Device used for evaluation. - pad_token_id: Token ID ignored in loss. - max_batches: Maximum validation batches to evaluate. Zero evaluates the full loader. - progress: Optional progress callback. - should_stop: Optional cancellation callback. - step: Current optimizer step for progress metrics. - total_steps: Total planned optimizer steps for progress metrics. - percent: Current outer training progress percentage. - - Returns: - Mean validation loss. - """ - - model.eval() - losses: list[float] = [] - batch_limit = len(loader) if max_batches <= 0 else min(len(loader), max_batches) - with torch.no_grad(): - for batch_index, (x, y) in enumerate(loader, start=1): - if should_stop and should_stop(): - model.train() - raise RuntimeError("Training stopped by user during validation.") - if batch_index > batch_limit: - break - x = x.to(device) - y = y.to(device) - logits = model(x) - loss = F.cross_entropy( - logits.reshape(-1, logits.size(-1)), - y.reshape(-1), - ignore_index=pad_token_id, - ) - losses.append(float(loss.item())) - if progress and (batch_index == 1 or batch_index == batch_limit or batch_index % 10 == 0): - emit_progress( - progress, - f"Validation running: batch {batch_index}/{batch_limit}.", - percent, - step=step, - total_steps=total_steps, - system_cpu_percent=system_cpu_percent(), - system_ram_percent=system_ram_percent(), - validation_batch=batch_index, - validation_batches=batch_limit, - ) - model.train() - _release_cuda_cache() - return sum(losses) / max(len(losses), 1) - - -def latest_checkpoint(checkpoints_dir: Path) -> Optional[Path]: - """Find the newest checkpoint in a folder. - - Args: - checkpoints_dir: Directory containing checkpoint files. - - Returns: - Newest checkpoint path, or ``None``. - """ - - checkpoints = sorted( - checkpoints_dir.glob("checkpoint_*.pt"), - key=lambda path: path.stat().st_mtime, - reverse=True, - ) - return checkpoints[0] if checkpoints else None - - -def _config_value(data: dict[str, Any], key: str, default: Any) -> Any: - """Return a saved config value with a default for old checkpoints. - - Args: - data: Saved configuration dictionary. - key: Configuration key. - default: Default value when the key is missing. - - Returns: - Saved or default value. - """ - - return data[key] if key in data else default - - -def _same_config_value(left: Any, right: Any) -> bool: - """Compare config values with tolerance for numeric fields. - - Args: - left: First value. - right: Second value. - - Returns: - True when values are effectively equal. - """ - - if isinstance(left, float) or isinstance(right, float): - try: - return abs(float(left) - float(right)) <= 1e-9 - except (TypeError, ValueError): - return False - return left == right - - -def _saved_model_default(key: str) -> Any: - """Return ModelConfig defaults for legacy checkpoints. - - Args: - key: ModelConfig field name. - - Returns: - Default value used by current ModelConfig. - """ - - defaults = { - "context_length": 128, - "embedding_size": 256, - "head_count": 4, - "layer_count": 4, - "dropout": 0.1, - "bias": True, - "norm_type": "layernorm", - "position_encoding": "learned", - "mlp_type": "gelu", - "rope_theta": 10000.0, - "attention_type": "mha", - "kv_head_count": 0, - "attention_backend": "sdpa", - "attention_window": 0, - } - return defaults.get(key) - - -def _saved_training_default(key: str) -> Any: - """Return TrainingConfig defaults for legacy checkpoints. - - Args: - key: TrainingConfig field name. - - Returns: - Default value used by current TrainingConfig. - """ - - defaults = { - "optimizer_name": "adamw", - "scheduler_name": "warmup_linear", - "scheduler_min_lr_ratio": 0.1, - "polynomial_power": 1.0, - "learning_rate": 3e-4, - "weight_decay": 0.1, - "max_grad_norm": 1.0, - "precision": "fp16", - "use_amp": True, - "training_mode": "pretrain", - "fine_tune_from_checkpoint": None, - } - return defaults.get(key) - - -def check_resume_compatibility( - checkpoint_path: Path, - model_config: ModelConfig, - training_config: TrainingConfig, -) -> ResumeCompatibilityReport: - """Check whether a checkpoint can be safely resumed. - - Args: - checkpoint_path: Checkpoint to inspect. - model_config: Current model configuration. - training_config: Current training configuration. - - Returns: - Resume compatibility report. - """ - - checkpoint = torch.load(checkpoint_path, map_location="cpu") - saved_model = checkpoint.get("model_config", {}) - saved_training = checkpoint.get("training_config", {}) - if not isinstance(saved_model, dict): - saved_model = {} - if not isinstance(saved_training, dict): - saved_training = {} - - errors: list[str] = [] - warnings: list[str] = [] - info: list[str] = [f"Resume checkpoint: {checkpoint_path.name}."] - - critical_model_fields = ( - ("vocab_size", model_config.vocab_size, "Tokenizer vocabulary"), - ("context_length", model_config.context_length, "Context length"), - ("embedding_size", model_config.embedding_size, "n_embd"), - ("head_count", model_config.head_count, "n_head"), - ("layer_count", model_config.layer_count, "n_layer"), - ("bias", model_config.bias, "Bias layout"), - ("norm_type", model_config.norm_type, "Normalization"), - ("position_encoding", model_config.position_encoding, "Position encoding"), - ("mlp_type", model_config.mlp_type, "MLP type"), - ("rope_theta", model_config.rope_theta, "RoPE theta"), - ("attention_type", model_config.attention_type, "Attention type"), - ) - for key, current_value, label in critical_model_fields: - saved_value = _config_value(saved_model, key, _saved_model_default(key)) - if not _same_config_value(saved_value, current_value): - errors.append(f"{label} changed: checkpoint={saved_value}, current={current_value}.") - - saved_attention_type = _config_value(saved_model, "attention_type", "mha") - saved_kv_heads = _config_value(saved_model, "kv_head_count", 0) - try: - saved_kv_effective = ModelConfig( - vocab_size=int(_config_value(saved_model, "vocab_size", model_config.vocab_size)), - context_length=int(_config_value(saved_model, "context_length", _saved_model_default("context_length"))), - embedding_size=int(_config_value(saved_model, "embedding_size", _saved_model_default("embedding_size"))), - head_count=int(_config_value(saved_model, "head_count", _saved_model_default("head_count"))), - layer_count=int(_config_value(saved_model, "layer_count", _saved_model_default("layer_count"))), - attention_type=str(saved_attention_type), - kv_head_count=int(saved_kv_heads), - ).resolved_kv_head_count() - except Exception: - saved_kv_effective = saved_kv_heads - current_kv_effective = model_config.resolved_kv_head_count() - if saved_kv_effective != current_kv_effective: - errors.append(f"Effective KV heads changed: checkpoint={saved_kv_effective}, current={current_kv_effective}.") - - warning_model_fields = ( - ("dropout", model_config.dropout, "Dropout"), - ("attention_backend", model_config.attention_backend, "Attention backend"), - ("attention_window", model_config.attention_window, "Sliding attention window"), - ) - for key, current_value, label in warning_model_fields: - saved_value = _config_value(saved_model, key, _saved_model_default(key)) - if not _same_config_value(saved_value, current_value): - warnings.append(f"{label} changed: checkpoint={saved_value}, current={current_value}.") - - can_load_optimizer_state = True - can_load_scheduler_state = True - can_load_scaler_state = True - if "optimizer_state_dict" in checkpoint: - saved_optimizer = _config_value(saved_training, "optimizer_name", _saved_training_default("optimizer_name")) - if saved_optimizer != training_config.optimizer_name: - warnings.append( - f"Optimizer changed: checkpoint={saved_optimizer}, current={training_config.optimizer_name}. " - "Optimizer state will not be loaded." - ) - can_load_optimizer_state = False - if "scheduler_state_dict" in checkpoint: - saved_scheduler = _config_value(saved_training, "scheduler_name", _saved_training_default("scheduler_name")) - if saved_scheduler != training_config.scheduler_name: - warnings.append( - f"LR scheduler changed: checkpoint={saved_scheduler}, current={training_config.scheduler_name}. " - "Scheduler state will not be loaded." - ) - can_load_scheduler_state = False - for key, current_value, label in ( - ("scheduler_min_lr_ratio", training_config.scheduler_min_lr_ratio, "Scheduler min LR ratio"), - ("polynomial_power", training_config.polynomial_power, "Polynomial power"), - ): - saved_value = _config_value(saved_training, key, _saved_training_default(key)) - if not _same_config_value(saved_value, current_value): - warnings.append(f"{label} changed: checkpoint={saved_value}, current={current_value}.") - if "scaler_state_dict" in checkpoint: - saved_precision = _config_value(saved_training, "precision", _saved_training_default("precision")) - if saved_precision != training_config.precision: - warnings.append(f"Precision changed: checkpoint={saved_precision}, current={training_config.precision}.") - can_load_scaler_state = saved_precision == "fp16" and training_config.precision == "fp16" - - for key, current_value, label in ( - ("learning_rate", training_config.learning_rate, "Learning rate"), - ("weight_decay", training_config.weight_decay, "Weight decay"), - ("max_grad_norm", training_config.max_grad_norm, "Gradient clipping"), - ): - saved_value = _config_value(saved_training, key, _saved_training_default(key)) - if not _same_config_value(saved_value, current_value): - warnings.append(f"{label} changed: checkpoint={saved_value}, current={current_value}.") - - if not errors: - info.append("Checkpoint architecture and tokenizer are compatible.") - return ResumeCompatibilityReport( - checkpoint_path=checkpoint_path, - errors=errors, - warnings=warnings, - info=info, - can_load_optimizer_state=can_load_optimizer_state, - can_load_scheduler_state=can_load_scheduler_state, - can_load_scaler_state=can_load_scaler_state, - ) - - -def _configure_cuda_allocator() -> None: - """Configure the CUDA caching allocator to reduce VRAM over-reservation. - - Sets ``expandable_segments:True`` so PyTorch grows GPU memory in smaller - increments rather than grabbing large contiguous blocks up front. - """ - - key = "PYTORCH_CUDA_ALLOC_CONF" - current = os.environ.get(key, "") - if "expandable_segments" not in current: - new_value = "expandable_segments:True" - if current: - new_value = current + "," + new_value - os.environ[key] = new_value - - -def _release_cuda_cache() -> None: - """Release unused cached VRAM back to the OS.""" - - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - -def _estimated_training_vram_bytes(model: MicroGPT, model_config: ModelConfig, training_config: TrainingConfig) -> int: - """Return a conservative, explainable training-memory estimate. - - Parameters, gradients, and Adam-style optimizer states are generally - fp32. Activations vary by kernel, therefore the estimate intentionally - includes headroom rather than pretending to be exact. - """ - - parameter_count = sum(parameter.numel() for parameter in model.parameters()) - parameter_and_optimizer = parameter_count * 16 # fp32 weights, grads, m/v - activation_bytes_per_value = 2 if training_config.use_amp and training_config.precision != "fp32" else 4 - activation_multiplier = 3 if training_config.activation_checkpointing else 10 - activations = ( - training_config.batch_size - * model_config.context_length - * model_config.embedding_size - * model_config.layer_count - * activation_bytes_per_value - * activation_multiplier - ) - logits = training_config.batch_size * model_config.context_length * model_config.vocab_size * activation_bytes_per_value - return int((parameter_and_optimizer + activations + logits) * 1.15) - - -def train_model( - model_config: ModelConfig, - training_config: TrainingConfig, - train_tokens: Union[list[int], np.ndarray], - val_tokens: Union[list[int], np.ndarray], - pad_token_id: int, - progress: Optional[Callable[[Any], None]] = None, - should_stop: Optional[Callable[[], bool]] = None, - decode_preview: Optional[Callable[[list[int]], str]] = None, -) -> TrainingResult: - """Train a MicroGPT model. - - Args: - model_config: Architecture settings. - training_config: Optimizer, device, and checkpoint settings. - train_tokens: Training token stream. - val_tokens: Validation token stream. - pad_token_id: Token ID ignored by cross-entropy loss. - progress: Optional callback receiving progress dictionaries. - should_stop: Optional callback returning true when training should stop. - decode_preview: Optional callback that decodes token IDs into a short text preview. - - Returns: - Training result with checkpoint and summary paths. - """ - - model_config.validate() - training_config.validate() - set_seed(training_config.seed) - - # Reduce VRAM over-reservation by the CUDA caching allocator. - if training_config.device.startswith("cuda") and torch.cuda.is_available(): - _configure_cuda_allocator() - - training_config.output_dir.mkdir(parents=True, exist_ok=True) - checkpoints_dir = training_config.output_dir / "checkpoints" - checkpoints_dir.mkdir(parents=True, exist_ok=True) - - emit_progress(progress, "Building model...", 2) - model = MicroGPT(model_config).to(training_config.device) - model.enable_gradient_checkpointing(training_config.activation_checkpointing) - if training_config.device.startswith("cuda") and torch.cuda.is_available(): - free_vram, total_vram = torch.cuda.mem_get_info() - estimate = _estimated_training_vram_bytes(model, model_config, training_config) - emit_progress( - progress, - ( - f"VRAM preflight: estimated {estimate / (1024 ** 3):.2f} GB; " - f"currently free {free_vram / (1024 ** 3):.2f} GB of {total_vram / (1024 ** 3):.2f} GB." - ), - 3, - estimated_vram_gb=estimate / (1024 ** 3), - free_vram_gb=free_vram / (1024 ** 3), - ) - if estimate > free_vram * 0.85: - emit_progress( - progress, - "[WARN] Estimated training memory is close to available VRAM. Reduce micro-batch size, enable activation checkpointing, or use gradient accumulation.", - 3, - ) - _release_cuda_cache() - emit_progress(progress, "Preparing token batches...", 4) - loader_workers = max(0, int(training_config.data_loader_workers)) - pin_memory = training_config.device.startswith("cuda") and torch.cuda.is_available() - loader_kwargs = { - "num_workers": loader_workers, - "pin_memory": pin_memory, - "persistent_workers": loader_workers > 0, - } - train_loader = DataLoader( - TokenDataset(train_tokens, model_config.context_length, stride=training_config.sample_stride), - batch_size=training_config.batch_size, - shuffle=True, - drop_last=True, - **loader_kwargs, - ) - val_loader = None - if len(val_tokens) > model_config.context_length: - val_loader = DataLoader( - TokenDataset(val_tokens, model_config.context_length), - batch_size=training_config.batch_size, - shuffle=False, - drop_last=False, - **loader_kwargs, - ) - - global_step = 0 - start_epoch = 0 - final_train_loss = 0.0 - final_val_loss: Optional[float] = None - best_val_loss: Optional[float] = None - best_checkpoint_path: Optional[Path] = None - early_stop_counter = 0 - early_stopped = False - - resume_path = training_config.resume_from_checkpoint if training_config.resume else None - if resume_path is None and training_config.resume: - resume_path = latest_checkpoint(checkpoints_dir) - resume_checkpoint: Optional[dict[str, Any]] = None - resume_compatibility: Optional[ResumeCompatibilityReport] = None - if training_config.peft_method == "lora": - base_path = training_config.fine_tune_from_checkpoint - if resume_path and Path(resume_path).exists(): - resume_checkpoint = torch.load(resume_path, map_location="cpu") - checkpoint_base = resume_checkpoint.get("fine_tune_base_checkpoint") - if checkpoint_base: - base_path = Path(checkpoint_base) - if base_path is None: - raise ValueError("LoRA fine-tuning requires a base checkpoint.") - base_path = Path(base_path) - if not base_path.exists(): - raise FileNotFoundError(f"LoRA base checkpoint not found: {base_path}") - emit_progress(progress, f"Loading LoRA base checkpoint: {base_path}", 5) - base_checkpoint = torch.load(base_path, map_location="cpu") - model.load_state_dict(base_checkpoint["model_state_dict"]) - wrapped = apply_lora_adapters( - model, - training_config.lora_rank, - training_config.lora_alpha, - training_config.lora_dropout, - training_config.lora_target_modules, - ) - freeze_non_lora_parameters(model) - emit_progress( - progress, - f"LoRA enabled: {wrapped} module(s), {lora_parameter_count(model):,} trainable adapter parameter(s).", - 6, - ) - - optimizer = make_optimizer(model, training_config) - steps_per_epoch = max(math.ceil(len(train_loader) / training_config.gradient_accumulation), 1) - total_steps = max(steps_per_epoch * training_config.epochs, 1) - scheduler = make_scheduler(optimizer, total_steps, training_config) - use_autocast, use_scaler, autocast_dtype = amp_settings(training_config) - scaler = GradScaler("cuda", enabled=use_scaler) - emit_progress( - progress, - "Optimizer: " - f"{training_config.optimizer_name}, schedule: {training_config.scheduler_name}, " - f"precision: {training_config.precision}.", - 5, - ) - if resume_path and Path(resume_path).exists(): - emit_progress(progress, f"Resuming from checkpoint: {resume_path}", 6) - compatibility = resume_compatibility or check_resume_compatibility(Path(resume_path), model_config, training_config) - for line in compatibility.info: - emit_progress(progress, line, 6) - for line in compatibility.warnings: - emit_progress(progress, f"[WARN] {line}", 6) - strict_resume_errors = list(compatibility.errors) - if training_config.require_compatible_resume: - if not compatibility.can_load_optimizer_state: - strict_resume_errors.append("Safe resume requires matching optimizer state.") - if not compatibility.can_load_scheduler_state: - strict_resume_errors.append("Safe resume requires matching scheduler state.") - if not compatibility.can_load_scaler_state: - strict_resume_errors.append("Safe resume requires matching AMP scaler state.") - if strict_resume_errors: - message = "Checkpoint is not compatible with the current training settings:\n" + "\n".join( - f"- {line}" for line in strict_resume_errors - ) - raise ValueError(message) - checkpoint = resume_checkpoint or torch.load(resume_path, map_location="cpu") - if training_config.peft_method == "lora" and "adapter_state_dict" in checkpoint: - load_lora_state_dict(model, checkpoint["adapter_state_dict"]) - else: - model.load_state_dict(checkpoint["model_state_dict"]) - if "optimizer_state_dict" in checkpoint and compatibility.can_load_optimizer_state: - optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) - if "scaler_state_dict" in checkpoint and use_scaler and compatibility.can_load_scaler_state: - scaler.load_state_dict(checkpoint["scaler_state_dict"]) - global_step = int(checkpoint.get("global_step", 0)) - start_epoch = min(int(checkpoint.get("epoch", 0)), training_config.epochs) - final_train_loss = float(checkpoint.get("train_loss", 0.0)) - final_val_loss = checkpoint.get("val_loss") - # Recompute total_steps to account for the resumed global_step so that - # progress tracking, ETA, and the LR scheduler use a consistent target. - remaining_epochs = max(training_config.epochs - start_epoch, 1) - total_steps = max(total_steps, global_step + steps_per_epoch * remaining_epochs, global_step + 1,) - scheduler = make_scheduler(optimizer, total_steps, training_config) - - if "scheduler_state_dict" in checkpoint and compatibility.can_load_scheduler_state: - scheduler.load_state_dict(checkpoint["scheduler_state_dict"]) - emit_progress(progress, f"Checkpoint loaded at step {global_step}.", 8) - _release_cuda_cache() - else: - if ( - training_config.training_mode == "fine_tune" - and training_config.fine_tune_from_checkpoint is not None - and training_config.peft_method != "lora" - ): - base_path = Path(training_config.fine_tune_from_checkpoint) - if not base_path.exists(): - raise FileNotFoundError(f"Fine-tune base checkpoint not found: {base_path}") - emit_progress(progress, f"Fine-tuning from base checkpoint: {base_path}", 6) - compatibility = check_resume_compatibility(base_path, model_config, training_config) - for line in compatibility.info: - emit_progress(progress, line, 6) - for line in compatibility.warnings: - emit_progress(progress, f"[WARN] {line}", 6) - if compatibility.errors: - message = "Fine-tune base checkpoint is not compatible with the current model settings:\n" + "\n".join( - f"- {line}" for line in compatibility.errors - ) - raise ValueError(message) - checkpoint = torch.load(base_path, map_location="cpu") - model.load_state_dict(checkpoint["model_state_dict"]) - emit_progress(progress, "Base model weights loaded. Starting fresh fine-tune optimizer state.", 8) - else: - emit_progress(progress, "Starting new training run.", 6) - - model.train() - optimizer.zero_grad(set_to_none=True) - last_metric_time = perf_counter() - step_time_window: list[float] = [] - for epoch in range(start_epoch, training_config.epochs): - epoch_losses: list[float] = [] - epoch_batch_count = len(train_loader) - for batch_index, (x, y) in enumerate(train_loader): - if should_stop and should_stop(): - final_train_loss = sum(epoch_losses) / max(len(epoch_losses), 1) if epoch_losses else final_train_loss - stopped_path = checkpoints_dir / f"checkpoint_stopped_step_{global_step}.pt" - save_checkpoint( - stopped_path, - model, - optimizer, - scheduler, - scaler, - model_config, - training_config, - global_step, - epoch, - final_train_loss, - final_val_loss, - ) - emit_progress(progress, f"Training stopped. Resume checkpoint saved: {stopped_path}", 100) - summary_path = training_config.output_dir / "training_summary.json" - summary = { - "model_config": dataclass_to_jsonable(model_config), - "training_config": dataclass_to_jsonable(training_config), - "final_train_loss": final_train_loss, - "final_val_loss": final_val_loss, - "total_steps": global_step, - "stopped": True, - "resume_checkpoint": str(stopped_path), - "parameters": sum(p.numel() for p in model.parameters()), - } - summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8") - return TrainingResult(stopped_path, summary_path, final_train_loss, final_val_loss, stopped=True) - x = x.to(training_config.device, non_blocking=pin_memory) - y = y.to(training_config.device, non_blocking=pin_memory) - with autocast("cuda", enabled=use_autocast, dtype=autocast_dtype): - logits = model(x) - loss = F.cross_entropy( - logits.reshape(-1, logits.size(-1)), - y.reshape(-1), - ignore_index=pad_token_id, - ) - loss = loss / training_config.gradient_accumulation - - scaler.scale(loss).backward() - should_step = ( - (batch_index + 1) % training_config.gradient_accumulation == 0 - or (batch_index + 1) == epoch_batch_count - ) - if should_step: - scaler.unscale_(optimizer) - grad_norm_tensor = torch.nn.utils.clip_grad_norm_(model.parameters(), training_config.max_grad_norm) - grad_norm = float(grad_norm_tensor.item() if hasattr(grad_norm_tensor, "item") else grad_norm_tensor) - weight_norm = math.sqrt( - sum(float(parameter.detach().float().norm(2).item()) ** 2 for parameter in model.parameters()) - ) - learning_rate = float(scheduler.get_last_lr()[0]) - update_ratio = learning_rate * grad_norm / max(weight_norm, 1e-12) - scaler.step(optimizer) - scaler.update() - optimizer.zero_grad(set_to_none=True) - scheduler.step() - global_step += 1 - now = perf_counter() - step_seconds = max(now - last_metric_time, 1e-9) - last_metric_time = now - step_time_window.append(step_seconds) - step_time_window = step_time_window[-50:] - average_step_seconds = sum(step_time_window) / max(len(step_time_window), 1) - remaining_steps = max(total_steps - global_step, 0) - eta_seconds = remaining_steps * average_step_seconds - samples_seen = training_config.batch_size * training_config.gradient_accumulation - tokens_seen = samples_seen * model_config.context_length - vram_allocated_gb = None - vram_reserved_gb = None - gpu_memory_percent = None - if training_config.device.startswith("cuda") and torch.cuda.is_available(): - device_index = torch.cuda.current_device() - vram_allocated_gb = torch.cuda.memory_allocated(device_index) / (1024 ** 3) - vram_reserved_gb = torch.cuda.memory_reserved(device_index) / (1024 ** 3) - free_vram, total_vram = torch.cuda.mem_get_info(device_index) - gpu_memory_percent = 100.0 * (1.0 - (free_vram / max(total_vram, 1))) - sample_text = None - if decode_preview is not None: - try: - sample_text = decode_preview(x[0].detach().cpu().tolist()) - except Exception: - sample_text = None - current_progress = 8 + int(86 * min(global_step, total_steps) / max(total_steps, 1)) - emit_progress( - progress, - f"Epoch {epoch + 1}/{training_config.epochs}, step {global_step}/{total_steps}, loss {float(loss.item() * training_config.gradient_accumulation):.4f}", - current_progress, - epoch=epoch + 1, - total_epochs=training_config.epochs, - step=global_step, - total_steps=total_steps, - train_loss=float(loss.item() * training_config.gradient_accumulation), - val_loss=final_val_loss, - learning_rate=learning_rate, - grad_norm=grad_norm, - weight_norm=weight_norm, - update_ratio=update_ratio, - tokens_per_second=tokens_seen / step_seconds, - samples_per_second=samples_seen / step_seconds, - step_seconds=step_seconds, - average_step_seconds=average_step_seconds, - eta_seconds=eta_seconds, - remaining_steps=remaining_steps, - vram_allocated_gb=vram_allocated_gb, - vram_reserved_gb=vram_reserved_gb, - gpu_memory_percent=gpu_memory_percent, - system_cpu_percent=system_cpu_percent(), - system_ram_percent=system_ram_percent(), - data_loader_workers=loader_workers, - sample_text=sample_text, - ) - - if ( - val_loader is not None - and training_config.eval_interval > 0 - and global_step % training_config.eval_interval == 0 - ): - emit_progress( - progress, - f"Running validation at step {global_step}...", - current_progress, - epoch=epoch + 1, - total_epochs=training_config.epochs, - step=global_step, - total_steps=total_steps, - train_loss=float(loss.item() * training_config.gradient_accumulation), - val_loss=final_val_loss, - system_cpu_percent=system_cpu_percent(), - system_ram_percent=system_ram_percent(), - ) - final_val_loss = evaluate( - model, - val_loader, - training_config.device, - pad_token_id, - training_config.max_eval_batches, - progress, - should_stop, - global_step, - total_steps, - current_progress, - ) - emit_progress( - progress, - f"Validation loss at step {global_step}: {final_val_loss:.4f}", - current_progress, - epoch=epoch + 1, - total_epochs=training_config.epochs, - step=global_step, - total_steps=total_steps, - train_loss=epoch_losses[-1] if epoch_losses else None, - val_loss=final_val_loss, - system_cpu_percent=system_cpu_percent(), - system_ram_percent=system_ram_percent(), - ) - if best_val_loss is None or final_val_loss < best_val_loss: - best_val_loss = final_val_loss - best_checkpoint_path = checkpoints_dir / "checkpoint_best_val.pt" - save_checkpoint( - best_checkpoint_path, - model, - optimizer, - scheduler, - scaler, - model_config, - training_config, - global_step, - epoch + 1, - epoch_losses[-1] if epoch_losses else final_train_loss, - final_val_loss, - ) - emit_progress( - progress, - f"New best validation checkpoint: {best_checkpoint_path.name} ({best_val_loss:.4f}).", - current_progress, - checkpoint_quality="best_validation", - best_val_loss=best_val_loss, - best_checkpoint_path=str(best_checkpoint_path), - ) - early_stop_counter = 0 - elif training_config.early_stopping and best_val_loss is not None: - early_stop_counter += 1 - if early_stop_counter >= training_config.early_stopping_patience: - reason = ( - f"Early stopping: validation loss has not improved for " - f"{early_stop_counter} consecutive evaluation(s). " - f"Best val loss: {best_val_loss:.4f}, current: {final_val_loss:.4f}. " - f"Best checkpoint: {best_checkpoint_path}." - ) - emit_progress(progress, reason, current_progress) - early_stopped = True - break - - if training_config.save_interval > 0 and global_step % training_config.save_interval == 0: - save_checkpoint( - checkpoints_dir / f"checkpoint_{global_step}.pt", - model, - optimizer, - scheduler, - scaler, - model_config, - training_config, - global_step, - epoch + 1, - final_train_loss, - final_val_loss, - ) - emit_progress(progress, f"Saved checkpoint at step {global_step}.", current_progress) - - epoch_losses.append(float(loss.item() * training_config.gradient_accumulation)) - - if early_stopped: - break - final_train_loss = sum(epoch_losses) / max(len(epoch_losses), 1) - if val_loader is not None: - final_val_loss = evaluate( - model, - val_loader, - training_config.device, - pad_token_id, - training_config.max_eval_batches, - progress, - should_stop, - global_step, - total_steps, - 8 + int(86 * (epoch + 1) / max(training_config.epochs, 1)), - ) - if best_val_loss is None or final_val_loss < best_val_loss: - best_val_loss = final_val_loss - best_checkpoint_path = checkpoints_dir / "checkpoint_best_val.pt" - save_checkpoint( - best_checkpoint_path, - model, - optimizer, - scheduler, - scaler, - model_config, - training_config, - global_step, - epoch + 1, - final_train_loss, - final_val_loss, - ) - emit_progress( - progress, - f"New best validation checkpoint: {best_checkpoint_path.name} ({best_val_loss:.4f}).", - 8 + int(86 * (epoch + 1) / max(training_config.epochs, 1)), - checkpoint_quality="best_validation", - best_val_loss=best_val_loss, - best_checkpoint_path=str(best_checkpoint_path), - ) - early_stop_counter = 0 - elif training_config.early_stopping and best_val_loss is not None: - early_stop_counter += 1 - if early_stop_counter >= training_config.early_stopping_patience: - reason = ( - f"Early stopping at epoch {epoch + 1}: validation loss has not improved for " - f"{early_stop_counter} consecutive evaluation(s). " - f"Best val loss: {best_val_loss:.4f}, current: {final_val_loss:.4f}. " - f"Best checkpoint: {best_checkpoint_path}." - ) - emit_progress(progress, reason, 8 + int(86 * (epoch + 1) / max(training_config.epochs, 1))) - early_stopped = True - print(f"epoch {epoch + 1}/{training_config.epochs}: train_loss={final_train_loss:.4f}") - save_checkpoint( - checkpoints_dir / f"checkpoint_epoch_{epoch + 1}.pt", - model, - optimizer, - scheduler, - scaler, - model_config, - training_config, - global_step, - epoch + 1, - final_train_loss, - final_val_loss, - ) - emit_progress( - progress, - f"Epoch {epoch + 1} complete. Checkpoint saved.", - 8 + int(86 * (epoch + 1) / max(training_config.epochs, 1)), - epoch=epoch + 1, - total_epochs=training_config.epochs, - step=global_step, - total_steps=total_steps, - train_loss=final_train_loss, - val_loss=final_val_loss, - system_cpu_percent=system_cpu_percent(), - system_ram_percent=system_ram_percent(), - ) - if early_stopped: - break - - if training_config.peft_method == "lora": - adapter_path = training_config.output_dir / "final_adapter.pt" - save_checkpoint( - adapter_path, - model, - optimizer, - scheduler, - scaler, - model_config, - training_config, - global_step, - training_config.epochs, - final_train_loss, - final_val_loss, - ) - merged_count = merge_lora_adapters(model) - emit_progress(progress, f"Merged {merged_count} LoRA adapter module(s) into final model weights.", 96) - checkpoint_path = training_config.output_dir / "final_model.pt" - save_checkpoint( - checkpoint_path, - model, - optimizer, - scheduler, - scaler, - model_config, - training_config, - global_step, - training_config.epochs, - final_train_loss, - final_val_loss, - ) - summary_path = training_config.output_dir / "training_summary.json" - summary = { - "model_config": dataclass_to_jsonable(model_config), - "training_config": dataclass_to_jsonable(training_config), - "final_train_loss": final_train_loss, - "final_val_loss": final_val_loss, - "best_val_loss": best_val_loss, - "best_checkpoint_path": str(best_checkpoint_path) if best_checkpoint_path else None, - "recommended_checkpoint_path": str(best_checkpoint_path or checkpoint_path), - "total_steps": global_step, - "parameters": sum(p.numel() for p in model.parameters()), - "adapter_checkpoint": str(training_config.output_dir / "final_adapter.pt") if training_config.peft_method == "lora" else None, - "early_stopped": early_stopped, - } - summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8") - emit_progress( - progress, - "Training stopped early — validation loss converged." if early_stopped else "Training complete.", - 100, - epoch=training_config.epochs, - total_epochs=training_config.epochs, - step=global_step, - total_steps=total_steps, - train_loss=final_train_loss, - val_loss=final_val_loss, - ) - return TrainingResult(checkpoint_path, summary_path, final_train_loss, final_val_loss) - - -def save_checkpoint( - path: Path, - model: MicroGPT, - optimizer: torch.optim.Optimizer, - scheduler, - scaler: GradScaler, - model_config: ModelConfig, - training_config: TrainingConfig, - global_step: int, - epoch: int, - train_loss: float, - val_loss: Optional[float], -) -> None: - """Save a resumable training checkpoint. - - Args: - path: Destination checkpoint path. - model: Model being trained. - optimizer: Optimizer state to save. - scheduler: Learning-rate scheduler state to save. - scaler: AMP scaler state to save. - model_config: Model configuration. - training_config: Training configuration. - global_step: Current optimizer step. - epoch: Current epoch number. - train_loss: Most recent training loss. - val_loss: Most recent validation loss. - """ - - payload = { - "optimizer_state_dict": optimizer.state_dict(), - "scheduler_state_dict": scheduler.state_dict(), - "scaler_state_dict": scaler.state_dict(), - "model_config": dataclass_to_jsonable(model_config), - "training_config": dataclass_to_jsonable(training_config), - "global_step": global_step, - "epoch": epoch, - "train_loss": train_loss, - "val_loss": val_loss, - } - if training_config.peft_method == "lora" and path.name != "final_model.pt": - payload["adapter_state_dict"] = lora_state_dict(model) - payload["fine_tune_base_checkpoint"] = ( - str(training_config.fine_tune_from_checkpoint) - if training_config.fine_tune_from_checkpoint - else None - ) - payload["lora_config"] = { - "rank": training_config.lora_rank, - "alpha": training_config.lora_alpha, - "dropout": training_config.lora_dropout, - "target_modules": training_config.lora_target_modules, - } - else: - payload["model_state_dict"] = model.state_dict() - torch.save(payload, path) \ No newline at end of file diff --git a/llm_trainer/training_orchestrator.py b/llm_trainer/training_orchestrator.py deleted file mode 100644 index 46f7c36..0000000 --- a/llm_trainer/training_orchestrator.py +++ /dev/null @@ -1,129 +0,0 @@ -from __future__ import annotations - -import json -import shutil -from pathlib import Path -from typing import Any, Callable, Optional - -import numpy as np - -from .config import ModelConfig, TrainingConfig -from .data import file_sha256 -from .lineage import read_json, stable_json_hash, utc_timestamp, write_json -from .resume_checks import _validate_resume_compatibility -from .tokenizer import PAD_TOKEN, load_tokenizer, token_id, validate_training_tokenizer -from .training import TrainingResult, train_model - - -def _load_tokens_for_training(data_dir: Path) -> tuple[Any, Any]: - train_npy = data_dir / "train_tokens.npy" - val_npy = data_dir / "val_tokens.npy" - if train_npy.exists() and val_npy.exists(): - train_tokens = np.load(train_npy, mmap_mode="r", allow_pickle=False) - val_tokens = np.load(val_npy, mmap_mode="r", allow_pickle=False) - return train_tokens, val_tokens - train_json = data_dir / "train_tokens.json" - val_json = data_dir / "val_tokens.json" - if train_json.exists() and val_json.exists(): - train_tokens = json.loads(train_json.read_text(encoding="utf-8")) - val_tokens = json.loads(val_json.read_text(encoding="utf-8")) - return train_tokens, val_tokens - raise FileNotFoundError("Prepared dataset is missing token files (expected .npy or .json train/val tokens).") - - -def train_from_dataset( - data_dir: Path, - model_config: ModelConfig, - training_config: TrainingConfig, - progress: Optional[Callable[[Any], None]] = None, - should_stop: Optional[Callable[[], bool]] = None, -) -> TrainingResult: - """Train a model using a prepared dataset folder. - - Args: - data_dir: Prepared dataset folder. - model_config: Model architecture settings. - training_config: Optimizer and checkpoint settings. - progress: Optional callback receiving progress event dictionaries. - should_stop: Optional callback returning true when the user requested stop. - - Returns: - Training result with final checkpoint and summary paths. - - Raises: - FileNotFoundError: If the tokenizer is missing. - """ - - tokenizer_path = data_dir / "tokenizer.json" - if not tokenizer_path.exists(): - raise FileNotFoundError(f"Tokenizer not found: {tokenizer_path}") - - dataset_summary = read_json(data_dir / "dataset_summary.json", default={}) or {} - dataset_lineage = read_json(data_dir / "dataset_lineage.json", default={}) or {} - tokenizer = load_tokenizer(tokenizer_path) - validate_training_tokenizer(tokenizer) - train_tokens, val_tokens = _load_tokens_for_training(data_dir) - - if model_config.vocab_size != tokenizer.get_vocab_size(): - model_config.vocab_size = tokenizer.get_vocab_size() - - training_config.output_dir.mkdir(parents=True, exist_ok=True) - resume_path = _validate_resume_compatibility(data_dir, tokenizer_path, model_config, training_config) - if resume_path and progress: - progress({"message": f"Resume safety check passed: {resume_path}", "percent": 3}) - shutil.copy2(tokenizer_path, training_config.output_dir / "tokenizer.json") - for metadata_name in ("tokenizer_config.json", "special_tokens_map.json"): - metadata_path = data_dir / metadata_name - if metadata_path.exists(): - shutil.copy2(metadata_path, training_config.output_dir / metadata_name) - result = train_model( - model_config, - training_config, - train_tokens, - val_tokens, - pad_token_id=token_id(tokenizer, PAD_TOKEN), - progress=progress, - should_stop=should_stop, - decode_preview=lambda ids: tokenizer.decode(ids, skip_special_tokens=True), - ) - training_summary = read_json(result.summary_path, default={}) or {} - run_id = ( - f"run_{utc_timestamp()}_" - f"{stable_json_hash({'dataset': dataset_summary.get('dataset_version'), 'model': training_summary.get('model_config'), 'training': training_summary.get('training_config')})}" - ) - lineage = { - "schema": "micro_llm_model_lineage", - "version": 1, - "training_run_id": run_id, - "created_at": utc_timestamp(), - "dataset_dir": str(data_dir), - "dataset_id": dataset_summary.get("dataset_id") or dataset_lineage.get("dataset_id"), - "dataset_version": dataset_summary.get("dataset_version"), - "dataset_fingerprint": (dataset_summary.get("dataset_version") or {}).get("source_fingerprint"), - "tokenizer_path": str(tokenizer_path), - "tokenizer_vocab_size": tokenizer.get_vocab_size(), - "tokenizer_sha256": file_sha256(tokenizer_path), - "training_mode": training_config.training_mode, - "fine_tune_from_checkpoint": ( - str(training_config.fine_tune_from_checkpoint) - if training_config.fine_tune_from_checkpoint - else None - ), - "peft_method": training_config.peft_method, - "lora_rank": training_config.lora_rank if training_config.peft_method == "lora" else None, - "lora_alpha": training_config.lora_alpha if training_config.peft_method == "lora" else None, - "lora_target_modules": training_config.lora_target_modules if training_config.peft_method == "lora" else None, - "resume_checkpoint": str(resume_path) if resume_path else None, - "resume_safety_required": training_config.require_compatible_resume, - "checkpoint_path": str(result.checkpoint_path), - "summary_path": str(result.summary_path), - "stopped": result.stopped, - } - training_summary["training_run_id"] = run_id - training_summary["model_lineage"] = lineage - write_json(result.summary_path, training_summary) - write_json(training_config.output_dir / "model_lineage.json", lineage) - write_json(training_config.output_dir / "dataset_summary.json", dataset_summary) - return result - -__all__ = ["train_from_dataset"] diff --git a/llm_trainer/training_planning.py b/llm_trainer/training_planning.py deleted file mode 100644 index ac3d35a..0000000 --- a/llm_trainer/training_planning.py +++ /dev/null @@ -1,141 +0,0 @@ -from __future__ import annotations - -from .config import ModelConfig, TrainingConfig - - -def estimate_model_parameters(model_config: ModelConfig) -> int: - """Estimate trainable parameters for a MicroGPT architecture. - - Args: - model_config: Model architecture configuration. - - Returns: - Approximate trainable parameter count. - """ - - vocab = model_config.vocab_size - emb = model_config.embedding_size - layers = model_config.layer_count - breakdown = estimate_parameter_breakdown(model_config) - return int(sum(breakdown.values())) - - -def estimate_parameter_breakdown(model_config: ModelConfig) -> dict[str, int]: - """Estimate parameter groups for a MicroGPT architecture. - - Args: - model_config: Model architecture configuration. - - Returns: - Dictionary with parameter counts by major model component. - """ - - vocab = model_config.vocab_size - emb = model_config.embedding_size - layers = model_config.layer_count - token_embedding = vocab * emb - position_embedding = model_config.context_length * emb if model_config.position_encoding == "learned" else 0 - head_size = emb // max(model_config.head_count, 1) - kv_emb = model_config.resolved_kv_head_count() * head_size - attention = (emb * (emb + (2 * kv_emb))) + (emb * emb) - if model_config.bias: - attention += emb + (2 * kv_emb) + emb - if model_config.mlp_type == "swiglu": - mlp = emb * 4 * emb * 3 - if model_config.bias: - mlp += 9 * emb - else: - mlp = (emb * 4 * emb) + (4 * emb * emb) - if model_config.bias: - mlp += 5 * emb - norms = 4 * emb - return { - "token_embedding": int(token_embedding), - "position_embedding": int(position_embedding), - "attention": int(layers * attention), - "mlp": int(layers * mlp), - "norms": int(layers * norms + (2 * emb)), - } - - -def estimate_training_resources( - model_config: ModelConfig, - training_config: TrainingConfig, - train_tokens: int, -) -> dict[str, int]: - """Estimate model size, VRAM, steps, and storage footprint. - - Args: - model_config: Selected model architecture. - training_config: Selected training settings. - train_tokens: Number of training tokens. - - Returns: - Estimate dictionary. - """ - - parameter_breakdown = estimate_parameter_breakdown(model_config) - params = int(sum(parameter_breakdown.values())) - mixed_precision = training_config.use_amp and training_config.device == "cuda" and training_config.precision in {"fp16", "bf16"} - param_bytes = params * (2 if mixed_precision else 4) - optimizer_bytes = params * 8 - activation_bytes = ( - training_config.batch_size - * model_config.context_length - * model_config.embedding_size - * model_config.layer_count - * 8 - ) - vram_bytes = param_bytes + optimizer_bytes + activation_bytes - kv_cache_bytes = ( - training_config.batch_size - * model_config.context_length - * model_config.layer_count - * model_config.resolved_kv_head_count() - * (model_config.embedding_size // max(model_config.head_count, 1)) - * 2 - * (2 if mixed_precision else 4) - ) - checkpoint_bytes = params * 16 - steps_per_epoch = max( - (train_tokens - model_config.context_length) - // max(model_config.context_length * training_config.batch_size, 1), - 1, - ) - total_steps = max(steps_per_epoch * training_config.epochs, 1) - checkpoint_count = max(total_steps // max(training_config.save_interval, 1), 1) + training_config.epochs + 2 - estimated_storage = checkpoint_bytes * checkpoint_count - return { - "parameters": params, - "parameter_breakdown": parameter_breakdown, - "checkpoint_bytes": checkpoint_bytes, - "vram_bytes": vram_bytes, - "memory_breakdown": { - "weights": int(param_bytes), - "optimizer": int(optimizer_bytes), - "activations": int(activation_bytes), - "kv_cache": int(kv_cache_bytes), - }, - "steps_per_epoch": steps_per_epoch, - "total_steps": total_steps, - "checkpoint_count": checkpoint_count, - "estimated_storage": estimated_storage, - } - - -def format_bytes(byte_count: float) -> str: - """Format a byte count for compact display. - - Args: - byte_count: Number of bytes. - - Returns: - Human-readable storage size. - """ - - value = float(byte_count) - for unit in ("B", "KB", "MB", "GB", "TB"): - if value < 1024 or unit == "TB": - return f"{value:.1f} {unit}" if unit != "B" else f"{value:.0f} B" - value /= 1024 - return f"{value:.1f} TB" diff --git a/llm_trainer/training_service.py b/llm_trainer/training_service.py deleted file mode 100644 index 36166b4..0000000 --- a/llm_trainer/training_service.py +++ /dev/null @@ -1,136 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from pathlib import Path -from typing import Any, Callable, Optional - -from .backends import LocalTrainerBackend, TrainerBackend -from .backends.registry import DEFAULT_BACKEND_REGISTRY, BackendRegistry -from .config import ModelConfig, TrainingConfig -from .contracts import BackendKind, TrainingJobSpec -from .coordinator import JobManager -from .training import TrainingResult - - -ProgressCallback = Callable[[Any], None] -StopCallback = Callable[[], bool] - - -@dataclass -class TrainingJobRequest: - """Request payload for a training service job. - - Args: - dataset_dir: Prepared dataset directory. - model_config: Model architecture settings. - training_config: Training runtime and optimizer settings. - """ - - dataset_dir: Path - model_config: ModelConfig - training_config: TrainingConfig - backend: BackendKind = BackendKind.LOCAL - metadata: Optional[dict[str, Any]] = None - - def to_spec(self) -> TrainingJobSpec: - """Convert this request to a backend-neutral job spec. - - Returns: - Training job specification. - """ - - if self.backend != BackendKind.LOCAL: - raise ValueError(f"Unsupported backend for local service: {self.backend.value}") - return TrainingJobSpec.local( - self.dataset_dir, - self.model_config, - self.training_config, - metadata=self.metadata, - ) - - -class TrainingService: - """Training service that dispatches jobs to a backend.""" - - def __init__( - self, - backend: Optional[TrainerBackend] = None, - registry: Optional[BackendRegistry] = None, - ) -> None: - """Create a training service. - - Args: - backend: Backend used to execute jobs. - registry: Backend registry used when backend is not provided. - """ - - self.backend = backend - self.registry = registry or DEFAULT_BACKEND_REGISTRY - - def run( - self, - request: TrainingJobRequest, - progress: Optional[ProgressCallback] = None, - should_stop: Optional[StopCallback] = None, - ) -> TrainingResult: - """Run a training request through the configured backend. - - Args: - request: Training job request. - progress: Optional progress callback. - should_stop: Optional cooperative cancellation callback. - - Returns: - Training result. - """ - - job = request.to_spec() - registry = self.registry - if self.backend is not None: - registry = BackendRegistry() - registry.register(job.runtime.backend, self.backend) - manager = JobManager(registry=registry) - manager.submit(job) - return manager.run_job(job.job_id, progress=progress, should_stop=should_stop) - - -class LocalTrainerService(TrainingService): - """Local trainer service used by the desktop coordinator. - - This service is intentionally small: the GUI depends on this API boundary - instead of calling the trainer directly, which leaves room for a future - cloud or multi-machine implementation behind the same contract. - """ - - def __init__(self) -> None: - """Create a local trainer service.""" - - super().__init__(LocalTrainerBackend()) - - -def run_training_job( - dataset_dir: Path, - model_config: ModelConfig, - training_config: TrainingConfig, - progress: Optional[ProgressCallback] = None, - should_stop: Optional[StopCallback] = None, -) -> TrainingResult: - """Run a training job through the configured training service. - - Args: - dataset_dir: Prepared dataset directory. - model_config: Model architecture settings. - training_config: Training runtime and optimizer settings. - progress: Optional progress callback. - should_stop: Optional cooperative cancellation callback. - - Returns: - Training result from the active trainer service. - """ - - service = TrainingService(LocalTrainerBackend()) - return service.run( - TrainingJobRequest(dataset_dir, model_config, training_config), - progress=progress, - should_stop=should_stop, - ) diff --git a/llm_trainer/ui/app.py b/llm_trainer/ui/app.py deleted file mode 100644 index 4c036aa..0000000 --- a/llm_trainer/ui/app.py +++ /dev/null @@ -1,7581 +0,0 @@ -from __future__ import annotations - -import ctypes -from datetime import datetime -import html -import importlib -import json -from functools import partial -import logging -import math -import os -from queue import Empty, Queue -import re -import shutil -import signal -import sqlite3 -import subprocess -import sys -from pathlib import Path -from threading import Event, Thread -from typing import Any, Optional, Union - -import torch -from PySide6.QtCore import QObject, QEvent, QPoint, Qt, QThread, QTimer, Slot, qInstallMessageHandler -from PySide6.QtGui import QBrush, QColor, QFont, QIcon, QPainter, QPen, QPixmap, QPolygon -from PySide6.QtGui import QFontDatabase -from PySide6.QtWidgets import ( - QApplication, - QAbstractButton, - QComboBox, - QDoubleSpinBox, - QDialog, - QFileDialog, - QFormLayout, - QGridLayout, - QHBoxLayout, - QInputDialog, - QLabel, - QLineEdit, - QListWidget, - QListWidgetItem, - QTreeWidget, - QTreeWidgetItem, - QMainWindow, - QMessageBox, - QProgressBar, - QPushButton, - QSizePolicy, - QStackedWidget, - QSpinBox, - QTextBrowser, - QTextEdit, - QVBoxLayout, - QWidget, -) - -from llm_trainer.app_logging import qt_message_handler, setup_logging -from llm_trainer.app_logging import DEFAULT_LOG_DIR -from llm_trainer.config import DatasetConfig, ModelConfig, TrainingConfig -from llm_trainer.conversation_datasets import CONVERSATION_DATASET_PRESETS, dataset_ids_for_stage, dataset_stage_label -from llm_trainer.contracts import BackendKind -from llm_trainer.contracts.jobs import RuntimeSpec, TrainingJobSpec -from llm_trainer.coordinator import CoordinatorApiServer, JobManager, create_job_artifact_bundle -from llm_trainer.evaluation import DEFAULT_BENCHMARK_PROMPTS, evaluate_checkpoint, normalize_prompts -from llm_trainer.export import export_gguf_with_llama_cpp, export_hf_microgpt_package, export_llama_adapter_package, export_project_bundle, quantize_checkpoint -from llm_trainer.fine_tuning_service import run_fine_tuning_job -from llm_trainer.llama_chat import LlamaChatSession, load_llama_chat_session, stream_chat_reply -from llm_trainer.lineage import read_json -from llm_trainer.microgpt_chat import load_microgpt_chat_session, stream_microgpt_chat_reply -from llm_trainer.notifier import NotificationManager, default_notifier_config_path, ensure_notifier_config -from llm_trainer.runpod_cloud import ( - RunPodClient, - RunPodConfig, - create_runpod_worker_bundle, - default_runpod_config_path, - ensure_runpod_config, - load_runpod_config, - public_url_is_cloud_reachable, - save_runpod_config, -) -from llm_trainer.dataset_build import build_dataset -from llm_trainer.dataset_preview import check_project_health, scan_dataset_preview -from llm_trainer.telemetry_store import initialize_store, insert_metric, latest_run, rows_until, telemetry_db_path -from llm_trainer.training import check_resume_compatibility, latest_checkpoint -from llm_trainer.training_planning import estimate_training_resources, format_bytes -from llm_trainer.training_service import run_training_job -from llm_trainer.external_dataset import ( - DEFAULT_MANIFEST_URL, - download_latest_dataset, - is_newer_version, - load_manifest, -) -from llm_trainer.ui.chat_widgets import ChatMessageWidget -from llm_trainer.ui.markdown_renderer import markdown_to_html -from llm_trainer.ui.workers import ProcessTaskWorker, TaskWorker -from llm_trainer.ui.startup_splash import StartupSplash -from llm_trainer.ui.tabs.benchmark_tab import build_benchmark_tab -from llm_trainer.ui.tabs.chat_tab import build_chat_tab -from llm_trainer.ui.tabs.dataset_tab import build_dataset_tab -from llm_trainer.ui.tabs.dataset_plan_tab import ( - build_dataset_plan_tab, - default_data_root, - default_data_stage, - dataset_plan_defaults, - iter_default_data_files, - populate_default_data_tree, -) -from llm_trainer.ui.tabs.live_tab import build_live_training_tab -from llm_trainer.ui.tabs.training_tab import build_training_tab -from llm_trainer.ui.tabs.export_tab import build_export_tab -from llm_trainer.ui.tabs.fine_tuning_tab import build_fine_tuning_tab -from llm_trainer.ui.tabs.job_manager_tab import build_job_manager_tab, set_table_rows -from llm_trainer.license_client import load_stored_license_key -from llm_trainer.ui.license_activation_dialog import LicenseActivationDialog, run_license_check_responsively - -try: - import psutil -except ImportError: - psutil = None - - -APP_NAME = "DrunkenBot LLM-IDE" -# Bump on every release that should require a version-ceiling check against -# licenses -- this is what license_client.check_license_at_launch compares -# against a license's version_ceiling/grace_period_until. -APP_VERSION = "1.0.0" -# TODO: point at the real deployed cloud-service URL once it has one. -# Overridable via env var so ops can point a build at a different -# deployment (dev/staging/prod) without a code change or rebuild. -# LICENSE_SERVER_URL = os.environ.get("DRUNKENBOT_LICENSE_SERVER_URL", "https://license.drunkenbot.ai") -# LICENSE_SERVER_URL = "http://127.0.0.1:8000/" -LICENSE_SERVER_URL = "https://drunkenbot.store" -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 _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 = MainWindow._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 = MainWindow._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 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"
    • ✓ {escaped}
      • ") - elif state == "running": - rows.append(f"
    • ● {escaped}
      • ") - elif state == "failed": - rows.append(f"
    • ✗ {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 = MainWindow._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) - - -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 _run_startup_validations(splash: StartupValidationSplash) -> None: - """Run startup checks shown on the splash screen.""" - - 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(module_name) for module_name in required_modules], - ), - ] - if tests_root.is_dir(): - steps.append(("Running test suite", lambda: _run_startup_tests(repo_root, tests_root, splash.append_log))) - else: - splash.append_log("Repository tests are not included in this packaged installation; skipping test suite.") - - splash.set_checks([label for label, _ in steps]) - splash.append_log(f"Workspace: {repo_root}") - for index, (label, action) in enumerate(steps, start=1): - splash.update_step(f"[{index}/{len(steps)}] {label}...", index - 1, len(steps)) - splash.mark_check_running(label) - try: - action() - except Exception: - splash.mark_check_failed(label) - raise - splash.mark_check_done(label) - splash.append_log(f"Completed: {label}") - splash.update_step("Startup checks complete", len(steps), len(steps)) - splash.append_log("All startup validations passed.") - - -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: - output = "\n".join(output_lines).strip() - tail = "\n".join(output.splitlines()[-25:]) - raise RuntimeError(f"Startup tests failed.\n{tail}") - - -class MainWindow(QMainWindow): - """Main PySide6 window for DrunkenBot LLM-IDE.""" - - def __init__(self) -> None: - """Create the main application window.""" - - super().__init__() - self.log_file_path = setup_logging() - LOGGER.info("Creating %s main window", APP_NAME) - if QApplication.instance(): - QApplication.instance().setFont(QFont("Arial", 10)) - self.setWindowTitle(APP_NAME) - self.setWindowIcon(self._app_icon()) - self._windows_icon_handles: list[int] = [] - self.resize(1240, 820) - self.thread: Optional[QThread] = None - self.worker: Optional[TaskWorker] = None - self.stop_event: Optional[Event] = None - self.progress_queue: Optional[Queue] = None - self.active_log: Optional[QTextEdit] = None - self.active_progress_bar: Optional[QProgressBar] = None - self.active_button: Optional[QPushButton] = None - self.active_stop_button: Optional[QPushButton] = None - self.active_button_text = "" - self.active_button_restore_text = "" - self.active_task_kind = "" - self.notification_manager: Optional[NotificationManager] = None - self.current_project_file: Optional[Path] = None - self.telemetry_db_path: Optional[Path] = None - self.telemetry_run_id = "" - self.telemetry_latest_id = 0 - self.telemetry_latest_index = 0 - self.live_scrub_active = False - self.hardware_meter_labels: dict[int, QLabel] = {} - self.training_cards: list[QWidget] = [] - self.training_controls_grid: Optional[QGridLayout] = None - self.training_controls_columns = 3 - self.training_health_points: list[tuple[int, Optional[float], Optional[float]]] = [] - self.active_training_log: Optional[QTextEdit] = None - self.active_training_progress: Optional[QProgressBar] = None - self.active_training_final_button_text = "Start Training" - self.active_training_output_dir: Optional[Path] = None - self.interrupt_count = 0 - self.chat_session: Optional[LlamaChatSession] = None - self.chat_markdown = "" - self.chat_stream_prefix = "" - self.chat_stream_reply = "" - self.current_assistant_browser: Optional[QTextBrowser] = None - self.current_assistant_meta: Optional[QLabel] = None - self.current_assistant_message: Optional[ChatMessageWidget] = None - self.pending_user_message = "" - self.spinner_index = 0 - self.spinner_timer = QTimer(self) - self.spinner_timer.timeout.connect(self._tick_spinner) - self.progress_timer = QTimer(self) - self.progress_timer.timeout.connect(self._drain_progress_queue) - self.job_manager = JobManager() - self.coordinator_server: Optional[CoordinatorApiServer] = None - self.coordinator_thread: Optional[Thread] = None - self.job_manager_timer = QTimer(self) - self.job_manager_timer.setInterval(2500) - self.job_manager_timer.timeout.connect(self.refresh_job_manager_tab) - - self._apply_style() - - shell = self._build_shell() - self.setCentralWidget(shell) - self._install_ui_event_logging(shell) - self._install_wheel_guard(shell) - self._refresh_notification_manager() - self.job_manager_timer.start() - - def eventFilter(self, watched: QObject, event: QEvent) -> bool: - """Prevent accidental wheel changes on compact option widgets. - - Args: - watched: Widget receiving the event. - event: Qt event. - - Returns: - True when the event is handled by the filter. - """ - - guarded_types = (QSpinBox, QDoubleSpinBox, QComboBox) - if isinstance(watched, guarded_types): - if event.type() == QEvent.Type.MouseButtonPress: - watched.setProperty("_wheel_enabled_after_click", True) - elif event.type() == QEvent.Type.FocusOut: - watched.setProperty("_wheel_enabled_after_click", False) - elif event.type() == QEvent.Type.Wheel and not watched.property("_wheel_enabled_after_click"): - return True - return super().eventFilter(watched, event) - - def _install_wheel_guard(self, root: QWidget) -> None: - """Require a click before spin boxes and combos react to mouse wheel. - - Args: - root: Root widget to scan for child controls. - """ - - for widget in root.findChildren(QWidget): - if not isinstance(widget, (QSpinBox, QDoubleSpinBox, QComboBox)): - continue - widget.setFocusPolicy(Qt.FocusPolicy.ClickFocus) - widget.setProperty("_wheel_enabled_after_click", False) - widget.installEventFilter(self) - - def _install_ui_event_logging(self, root: QWidget) -> None: - """Log user-facing widget actions and parameter changes. - - Args: - root: Root widget to scan for child controls. - """ - - for widget in root.findChildren(QWidget): - if isinstance(widget, QAbstractButton): - if widget.isCheckable(): - widget.toggled.connect( - lambda checked, item=widget: self._log_ui_event("toggled", item, checked) - ) - else: - widget.clicked.connect( - lambda checked=False, item=widget: self._log_ui_event("clicked", item, checked) - ) - elif isinstance(widget, QComboBox): - widget.currentTextChanged.connect( - lambda value, item=widget: self._log_ui_event("changed", item, value) - ) - elif isinstance(widget, QSpinBox): - widget.valueChanged.connect( - lambda value, item=widget: self._log_ui_event("changed", item, value) - ) - elif isinstance(widget, QDoubleSpinBox): - widget.valueChanged.connect( - lambda value, item=widget: self._log_ui_event("changed", item, value) - ) - elif isinstance(widget, QLineEdit): - widget.editingFinished.connect( - lambda item=widget: self._log_ui_event("edited", item, item.text()) - ) - - def _log_ui_event(self, action: str, widget: QWidget, value: Any) -> None: - """Log a UI action or parameter value. - - Args: - action: Event label. - widget: Widget that emitted the event. - value: Current value. - """ - - if action == "clicked" and isinstance(widget, QAbstractButton) and not widget.isCheckable(): - LOGGER.info("UI clicked: %s", self._widget_log_name(widget)) - return - LOGGER.info("UI %s: %s = %s", action, self._widget_log_name(widget), value) - - @staticmethod - def _widget_log_name(widget: QWidget) -> str: - """Return a useful log label for a widget. - - Args: - widget: Widget to describe. - - Returns: - Human-readable widget label. - """ - - if isinstance(widget, QAbstractButton) and widget.text(): - return widget.text().replace("\n", " ") - if isinstance(widget, QLineEdit) and widget.placeholderText(): - return widget.placeholderText() - if widget.objectName(): - return widget.objectName() - return widget.__class__.__name__ - - def _apply_style(self) -> None: - """Load the application stylesheet from the QSS module file.""" - - qss_path = Path(__file__).with_name("styles.qss") - self.setStyleSheet(qss_path.read_text(encoding="utf-8")) - - def _build_shell(self) -> QWidget: - """Build the top-level dashboard shell. - - Returns: - Root shell widget. - """ - - shell = QWidget() - shell.setObjectName("AppShell") - root = QVBoxLayout(shell) - root.setContentsMargins(8, 8, 8, 8) - root.setSpacing(0) - - top = QWidget() - top.setObjectName("TopBar") - self.top_bar = top - top_layout = QHBoxLayout(top) - top_layout.setContentsMargins(16, 8, 16, 8) - top_layout.setSpacing(8) - logo = QLabel() - logo.setObjectName("Logo") - logo_pixmap = self._app_logo_pixmap(36) - if logo_pixmap.isNull(): - logo.setText("DB") - else: - logo.setPixmap(logo_pixmap) - logo.setFixedSize(42, 42) - logo.setScaledContents(False) - self.search_box = QLineEdit() - self.search_box.setPlaceholderText("Project name...") - self.search_box.setMaximumWidth(260) - self._tip(self.search_box, f"Project name used when saving or reopening a {APP_NAME} project.") - self.new_project_button = QPushButton("New Project") - self.new_project_button.setMaximumWidth(130) - self.new_project_button.clicked.connect(self.new_project) - self._tip(self.new_project_button, f"Start a fresh {APP_NAME} project with default paths and settings.") - self.save_project_button = QPushButton("Save Project") - self.save_project_button.setMaximumWidth(130) - self.save_project_button.clicked.connect(self.save_project) - self._tip(self.save_project_button, "Save all current paths and settings into a project.json file.") - self.open_project_button = QPushButton("Open Project") - self.open_project_button.setMaximumWidth(130) - self.open_project_button.clicked.connect(self.open_project) - self._tip(self.open_project_button, "Open a saved project.json file and restore the UI settings.") - self.dataset_status = QLabel("Dataset: not prepared") - self.train_status = QLabel("Training: idle") - self.export_status = QLabel("Export: waiting") - self.chat_status = QLabel("Chat: no model loaded") - for label in (self.dataset_status, self.train_status, self.export_status, self.chat_status): - label.setObjectName("TopStatus") - label.setMinimumWidth(0) - label.setMaximumWidth(180) - label.setSizePolicy(QSizePolicy.Ignored, QSizePolicy.Fixed) - label.setWordWrap(False) - self.project_state = QLabel("Ready") - self.project_state.setObjectName("Metric") - self.project_state.setMinimumWidth(0) - self.project_state.setSizePolicy(QSizePolicy.Ignored, QSizePolicy.Fixed) - top_layout.addWidget(logo) - top_layout.addSpacing(12) - top_layout.addWidget(self.search_box) - top_layout.addWidget(self.new_project_button) - top_layout.addWidget(self.save_project_button) - top_layout.addWidget(self.open_project_button) - top_layout.addSpacing(10) - top_layout.addWidget(self.dataset_status) - top_layout.addWidget(self.train_status) - top_layout.addWidget(self.export_status) - top_layout.addWidget(self.chat_status) - top_layout.addStretch(1) - top_layout.addWidget(self.project_state) - root.addWidget(top) - - body = QHBoxLayout() - body.setContentsMargins(0, 0, 0, 0) - body.setSpacing(0) - rail = QWidget() - rail.setObjectName("SideRail") - self.side_rail = rail - rail.setFixedWidth(82) - rail_layout = QVBoxLayout(rail) - rail_layout.setContentsMargins(12, 18, 12, 18) - rail_layout.setSpacing(12) - self.dataset_plan_nav = self._nav_button("PLAN") - self.dataset_nav = self._nav_button("IN") - self.training_nav = self._nav_button("✦\nAI") - self.training_nav.setText("AI") - self.fine_tune_nav = self._nav_button("FT") - self.live_nav = self._nav_button("LIVE") - self.jobs_nav = self._nav_button("JOB") - self.benchmark_nav = self._nav_button("◷\nBench") - self.export_nav = self._nav_button("⇧\nX") - self.chat_nav = self._nav_button("◌\nChat") - self._tip(self.dataset_plan_nav, "Open Dataset Blueprint: plan the target data mix before ingestion.") - self._tip(self.dataset_nav, "Open dataset preparation: load text/PDF files and build tokenizer data.") - self._tip(self.training_nav, "Open model training: configure architecture and optimization settings.") - self._tip(self.fine_tune_nav, "Open fine-tuning: adapt checkpoints with instruction, conversation, or LoRA settings.") - self._tip(self.live_nav, "Open the live training tracker with model flow, charts, metrics, and telemetry.") - self._tip(self.jobs_nav, "Open Job Manager: monitor workers, remote connections, assignments, and job controls.") - self._tip(self.benchmark_nav, "Open benchmark prompts: test checkpoint quality with repeatable prompts.") - self._tip(self.export_nav, "Open export tools: bundle or quantize the trained model artifacts.") - self._tip(self.chat_nav, "Open Chat: load a GGUF or native MicroGPT model once and send prompts.") - self.dataset_plan_nav.setChecked(True) - self.dataset_plan_nav.clicked.connect(lambda: self._switch_page(0)) - self.dataset_nav.clicked.connect(lambda: self._switch_page(1)) - self.training_nav.clicked.connect(lambda: self._switch_page(2)) - self.fine_tune_nav.clicked.connect(lambda: self._switch_page(3)) - self.live_nav.clicked.connect(lambda: self._switch_page(4)) - self.jobs_nav.clicked.connect(lambda: self._switch_page(5)) - self.benchmark_nav.clicked.connect(lambda: self._switch_page(6)) - self.export_nav.clicked.connect(lambda: self._switch_page(7)) - self.chat_nav.clicked.connect(lambda: self._switch_page(8)) - rail_layout.addWidget(self.dataset_plan_nav) - rail_layout.addWidget(self.dataset_nav) - rail_layout.addWidget(self.training_nav) - rail_layout.addWidget(self.fine_tune_nav) - rail_layout.addWidget(self.live_nav) - rail_layout.addWidget(self.jobs_nav) - rail_layout.addWidget(self.benchmark_nav) - rail_layout.addWidget(self.export_nav) - rail_layout.addWidget(self.chat_nav) - rail_layout.addStretch(1) - - self.pages = QStackedWidget() - self.pages.addWidget(self._build_dataset_plan_tab()) - self.pages.addWidget(self._build_dataset_tab()) - self.pages.addWidget(self._build_training_tab()) - self.pages.addWidget(self._build_fine_tuning_tab()) - self.pages.addWidget(self._build_live_training_tab()) - self.pages.addWidget(self._build_job_manager_tab()) - self.pages.addWidget(self._build_benchmark_tab()) - self.pages.addWidget(self._build_export_tab()) - self.pages.addWidget(self._build_chat_tab()) - - body.addWidget(rail) - body.addWidget(self.pages, 1) - root.addLayout(body, 1) - return shell - - def _nav_button(self, text: str) -> QPushButton: - """Create a left-rail navigation button. - - Args: - text: Button label. - - Returns: - Configured navigation button. - """ - - button = QPushButton(text) - button.setObjectName("NavButton") - button.setCheckable(True) - return button - - def _switch_page(self, index: int) -> None: - """Switch the visible page. - - Args: - index: Page index in the stacked widget. - """ - - self.pages.setCurrentIndex(index) - buttons = [ - self.dataset_plan_nav, - self.dataset_nav, - self.training_nav, - self.fine_tune_nav, - self.live_nav, - self.jobs_nav, - self.benchmark_nav, - self.export_nav, - self.chat_nav, - ] - for button_index, button in enumerate(buttons): - button.setChecked(button_index == index) - self._refresh_training_layout() - if index == 5: - QTimer.singleShot(20, self.refresh_job_manager_tab) - - def show_chat_only_mode(self) -> None: - """Collapse the UI to chat-only view for quick local LLM testing.""" - - if hasattr(self, "top_bar"): - self.top_bar.hide() - if hasattr(self, "side_rail"): - self.side_rail.hide() - self._switch_page(8) - self.setWindowTitle("DrunkenBot - Chat") - self.resize(980, 760) - - def resizeEvent(self, event: Any) -> None: - """Refresh responsive layouts when the main window changes size. - - Args: - event: Qt resize event. - """ - - super().resizeEvent(event) - self._refresh_training_layout() - - def _refresh_training_layout(self) -> None: - """Apply responsive card columns on the training page.""" - - if not self.training_cards or self.training_controls_grid is None: - return - width = self.pages.width() if hasattr(self, "pages") else self.width() - if width >= 900: - columns = 2 - else: - columns = 1 - if columns == self.training_controls_columns: - return - self._set_training_card_columns(columns) - - def _set_training_card_columns(self, columns: int) -> None: - """Reflow the training cards into the requested column count. - - Args: - columns: Number of columns to use. - """ - - if self.training_controls_grid is None: - return - while self.training_controls_grid.count(): - self.training_controls_grid.takeAt(0) - for index, card in enumerate(self.training_cards): - row = index // columns - column = index % columns - self.training_controls_grid.addWidget(card, row, column) - for column in range(2): - self.training_controls_grid.setColumnStretch(column, 1 if column < columns else 0) - self.training_controls_columns = columns - - def _build_dataset_plan_tab(self) -> QWidget: - """Build the dataset blueprint page. - - Returns: - Dataset blueprint page widget. - """ - - return build_dataset_plan_tab(self) - - def _build_dataset_tab(self) -> QWidget: - """Build the dataset preparation page. - - Returns: - Dataset page widget. - """ - - return build_dataset_tab(self) - - def _build_training_tab(self) -> QWidget: - """Build the training configuration page. - - Returns: - Training page widget. - """ - - return build_training_tab(self) - - def _build_fine_tuning_tab(self) -> QWidget: - """Build the fine-tuning page. - - Returns: - Fine-tuning page widget. - """ - - return build_fine_tuning_tab(self) - - def _build_live_training_tab(self) -> QWidget: - """Build the live training tracker page. - - Returns: - Live training tracker page widget. - """ - - return build_live_training_tab(self) - - def _build_job_manager_tab(self) -> QWidget: - """Build the distributed job manager page. - - Returns: - Job manager page widget. - """ - - return build_job_manager_tab(self) - - def refresh_job_manager_tab(self) -> None: - """Refresh the job manager dashboard tables.""" - - if not hasattr(self, "job_worker_table"): - return - workers = self.job_manager.list_workers() - jobs = self.job_manager.list_jobs() - heartbeats = self.job_manager.state_store.latest_heartbeats() - worker_rows = [] - for worker in workers: - heartbeat = heartbeats.get(worker.worker_id, {}) - metrics = heartbeat.get("metrics") or {} - active_job = heartbeat.get("active_job_id") or self._active_job_for_worker(worker.worker_id) - capabilities = worker.capabilities or {} - cpu_ram_gpu = ( - f"CPU {capabilities.get('cpu_count', '-')}, " - f"RAM {capabilities.get('system_ram_gb', '-')} GB, " - f"VRAM {capabilities.get('total_vram_gb', '-')} GB" - ) - if metrics: - cpu_ram_gpu = f"{cpu_ram_gpu}, util {metrics.get('gpu_util', metrics.get('gpu_memory_percent', '-'))}" - worker_rows.append( - [ - worker.worker_id, - worker.status.value, - worker.backend.value, - worker.device, - worker.last_heartbeat_at or "-", - active_job or "-", - cpu_ram_gpu, - ", ".join(capabilities.get("labels") or []) or "-", - ] - ) - set_table_rows(self.job_worker_table, worker_rows) - - job_rows = [] - for managed in jobs: - job = managed.spec - metrics = managed.latest_metrics - stage_label = str(job.metadata.get("training_stage") or job.metadata.get("training_mode") or job.training.training_mode) - job_rows.append( - [ - job.job_id, - stage_label, - job.status.value, - managed.assigned_worker_id or "-", - job.runtime.backend.value, - self._metric_pair(metrics.epoch if metrics else None, metrics.total_epochs if metrics else None), - self._metric_pair(metrics.step if metrics else None, metrics.total_steps if metrics else None), - str(job.training.batch_size), - str(job.model.config.layer_count), - self._metric_float(metrics.train_loss if metrics else None), - self._metric_float(metrics.tokens_per_second if metrics else None, suffix=" tok/s"), - managed.updated_at, - ] - ) - set_table_rows(self.job_table, job_rows) - active_count = sum(1 for item in jobs if item.spec.status.value in {"assigned", "running", "paused", "stopping"}) - queued_count = sum(1 for item in jobs if item.spec.status.value == "queued") - self.job_worker_count_label.setText(f"Workers: {len(workers)}") - self.job_active_count_label.setText(f"Active jobs: {active_count}") - self.job_queue_count_label.setText(f"Queued jobs: {queued_count}") - self.job_db_label.setText(f"State DB: {self.job_manager.state_store.db_path}") - self.job_manager_progress.setValue(100) - - def pause_all_managed_jobs(self) -> None: - """Pause all managed jobs.""" - - count = self.job_manager.pause_all_jobs() - self.job_manager_log.append(f"Pause requested for {count} job(s).") - self.refresh_job_manager_tab() - - def resume_all_managed_jobs(self) -> None: - """Resume all paused managed jobs.""" - - count = self.job_manager.resume_all_jobs() - self.job_manager_log.append(f"Resumed {count} job(s).") - self.refresh_job_manager_tab() - - def stop_all_managed_jobs(self) -> None: - """Stop all managed jobs.""" - - count = self.job_manager.stop_all_jobs() - self.job_manager_log.append(f"Stop requested for {count} job(s).") - self.refresh_job_manager_tab() - - def mark_stale_workers_offline(self) -> None: - """Mark stale remote workers offline.""" - - workers = self.job_manager.mark_stale_workers_offline() - if workers: - self.job_manager_log.append(f"Marked offline: {', '.join(workers)}") - else: - self.job_manager_log.append("No stale remote workers found.") - self.refresh_job_manager_tab() - - def start_coordinator_server(self) -> None: - """Start the coordinator API used by remote workers.""" - - if self.coordinator_server is not None: - self.job_manager_log.append("Coordinator API is already running.") - return - host = self.coordinator_host.text().strip() or "0.0.0.0" - port = self.coordinator_port.value() - artifact_root = Path(self.coordinator_artifact_root.text().strip()).expanduser() - artifact_root.mkdir(parents=True, exist_ok=True) - try: - self.coordinator_server = CoordinatorApiServer( - manager=self.job_manager, - host=host, - port=port, - artifact_root=artifact_root, - ) - self.coordinator_thread = Thread(target=self.coordinator_server.serve_forever, daemon=True) - self.coordinator_thread.start() - except Exception as exc: - self.coordinator_server = None - self.coordinator_thread = None - QMessageBox.warning(self, "Coordinator failed", f"Could not start coordinator API:\n{exc}") - return - public_url = self.coordinator_public_url.text().strip() or f"http://127.0.0.1:{port}" - self.coordinator_public_url.setText(public_url.rstrip("/")) - self.coordinator_status_label.setText(f"Coordinator: running at {public_url.rstrip('/')}") - self.coordinator_start_button.setEnabled(False) - self.coordinator_stop_button.setEnabled(True) - self.project_state.setText("Coordinator running") - self.job_manager_log.append(f"Coordinator API started on {host}:{port}.") - self.job_manager_log.append(f"Artifact sync root: {artifact_root}") - - def stop_coordinator_server(self) -> None: - """Stop the coordinator API.""" - - if self.coordinator_server is None: - return - self.coordinator_server.shutdown() - if self.coordinator_thread is not None: - self.coordinator_thread.join(timeout=3) - self.coordinator_server = None - self.coordinator_thread = None - self.coordinator_status_label.setText("Coordinator: stopped") - self.coordinator_start_button.setEnabled(True) - self.coordinator_stop_button.setEnabled(False) - self.project_state.setText("Coordinator stopped") - self.job_manager_log.append("Coordinator API stopped.") - - def _runpod_config_path(self) -> Path: - """Return the active RunPod config path. - - Returns: - Project-local RunPod config path when a project is open. - """ - - project_dir = self.current_project_file.parent if self.current_project_file is not None else None - return default_runpod_config_path(project_dir) - - def load_runpod_settings(self) -> None: - """Load RunPod settings into the Job Manager UI.""" - - if not hasattr(self, "runpod_api_key"): - return - config_path = self._runpod_config_path() - try: - config = load_runpod_config(config_path) - except Exception as exc: - LOGGER.error("Could not load RunPod config: %s", exc) - self.runpod_status_label.setText(f"RunPod config error: {exc}") - return - self.runpod_api_key.setText(config.api_key) - self._set_combo_text(self.runpod_gpu_type, config.gpu_type_id) - self._set_combo_text(self.runpod_cloud_type, config.cloud_type) - self.runpod_image.setText(config.image_name) - self.runpod_container_disk.setValue(config.container_disk_gb) - self.runpod_volume_gb.setValue(config.volume_gb) - self.runpod_min_ram.setValue(config.min_ram_per_gpu) - self.runpod_min_vcpu.setValue(config.min_vcpu_per_gpu) - self.runpod_spot.setChecked(config.interruptible) - self.runpod_auto_terminate.setChecked(config.auto_terminate) - status = "configured" if config.api_key.strip() else "API key needed" - self.runpod_status_label.setText(f"RunPod: {status} ({config_path})") - - def save_runpod_settings(self) -> None: - """Save RunPod settings from the Job Manager UI.""" - - config = self._runpod_config_from_ui() - config_path = self._runpod_config_path() - save_runpod_config(config_path, config) - self.runpod_status_label.setText(f"RunPod settings saved: {config_path}") - self.job_manager_log.append(f"RunPod settings saved: {config_path}") - LOGGER.info("RunPod settings saved: %s", config_path) - - def _runpod_config_from_ui(self) -> RunPodConfig: - """Collect RunPod settings from the UI. - - Returns: - RunPod configuration. - """ - - return RunPodConfig( - api_key=self.runpod_api_key.text().strip(), - image_name=self.runpod_image.text().strip(), - gpu_type_id=self.runpod_gpu_type.currentText().strip(), - gpu_count=1, - cloud_type=self.runpod_cloud_type.currentText().strip(), - interruptible=self.runpod_spot.isChecked(), - container_disk_gb=self.runpod_container_disk.value(), - volume_gb=self.runpod_volume_gb.value(), - min_vcpu_per_gpu=self.runpod_min_vcpu.value(), - min_ram_per_gpu=self.runpod_min_ram.value(), - auto_terminate=self.runpod_auto_terminate.isChecked(), - worker_labels="runpod,gpu", - ) - - def launch_runpod_worker_for_current_training(self, training_mode: str = "pretrain", stage: str = "base") -> None: - """Publish the current training job and launch a RunPod worker Pod. - - Args: - training_mode: Training mode for the queued job. - stage: Dataset/training stage label. - """ - - if isinstance(training_mode, bool): - training_mode = "pretrain" - stage = "base" - try: - config = self._runpod_config_from_ui() - save_runpod_config(self._runpod_config_path(), config) - coordinator_url = self.coordinator_public_url.text().strip().rstrip("/") - if not public_url_is_cloud_reachable(coordinator_url): - raise ValueError( - "RunPod needs a public Worker URL. Start a tunnel or set Worker URL to a public address, " - "not localhost/127.0.0.1." - ) - if self.coordinator_server is None: - self.start_coordinator_server() - if self.coordinator_server is None: - return - job, bundle_path = self._publish_remote_training_job_spec( - training_mode=training_mode, - stage=stage, - backend_label="runpod", - ) - artifact_root = Path(self.coordinator_artifact_root.text().strip()).expanduser() - bootstrap_path = create_runpod_worker_bundle(Path(__file__).resolve().parents[2], artifact_root) - bootstrap_url = f"{coordinator_url}/artifacts/{bootstrap_path.name}" - worker_id = f"runpod-{job.job_id}" - pod_name = f"micro-llm-{self._safe_project_name(self.search_box.text().strip() or 'project')}-{job.job_id[-8:]}" - result = RunPodClient(config.api_key).create_worker_pod( - config=config, - pod_name=pod_name, - worker_id=worker_id, - coordinator_url=coordinator_url, - bootstrap_url=bootstrap_url, - ) - managed = self.job_manager.get_job(job.job_id) - managed.spec.metadata["runpod_pod_id"] = result.pod_id - managed.spec.metadata["runpod_worker_id"] = result.worker_id - managed.spec.metadata["runpod_cost_per_hour"] = result.cost_per_hour - self.job_manager._persist_job(job.job_id) - except Exception as exc: - LOGGER.exception("RunPod launch failed") - QMessageBox.warning(self, "RunPod launch failed", str(exc)) - if hasattr(self, "runpod_status_label"): - self.runpod_status_label.setText(f"RunPod launch failed: {exc}") - return - self.runpod_status_label.setText( - f"RunPod pod {result.pod_id} launched for {job.job_id} ({result.gpu_name}, {result.cost_per_hour}/hr)" - ) - self.job_manager_log.append(f"RunPod pod launched: {result.pod_id}") - self.job_manager_log.append(f"RunPod worker: {result.worker_id}") - self.job_manager_log.append(f"RunPod GPU: {result.gpu_name}, cost/hr: {result.cost_per_hour}") - self.job_manager_log.append(f"Worker bootstrap: {result.bootstrap_url}") - self.project_state.setText("RunPod worker launched") - self.refresh_job_manager_tab() - - def publish_remote_training_job(self, training_mode: str = "pretrain", stage: str = "base") -> None: - """Bundle the current training setup and queue it for remote workers. - - Args: - training_mode: Trainer mode to publish, either ``pretrain`` or ``fine_tune``. - stage: Higher-level stage label for job manager display. - """ - - if isinstance(training_mode, bool): - training_mode = "pretrain" - stage = "base" - if self.coordinator_server is None: - self.start_coordinator_server() - if self.coordinator_server is None: - return - try: - job, bundle_path = self._publish_remote_training_job_spec(training_mode=training_mode, stage=stage) - except Exception as exc: - QMessageBox.warning(self, "Publish failed", f"Could not publish remote job:\n{exc}") - return - self.job_manager_log.append(f"Published remote job: {job.job_id}") - self.job_manager_log.append(f"Input bundle: {bundle_path}") - self.job_manager_log.append(f"Worker download URL: {job.metadata.get('artifact_bundle_url')}") - self.project_state.setText("Remote job queued") - self.refresh_job_manager_tab() - - def _publish_remote_training_job_spec( - self, - training_mode: str = "pretrain", - stage: str = "base", - backend_label: str = "remote", - ) -> tuple[TrainingJobSpec, Path]: - """Bundle and queue the current remote training job. - - Args: - training_mode: Trainer mode to publish. - stage: Higher-level stage label. - backend_label: Human-readable backend label stored in metadata. - - Returns: - Queued job and bundle path. - """ - - job = self._current_remote_training_job(training_mode=training_mode, stage=stage) - job.metadata["launch_backend"] = backend_label - artifact_root = Path(self.coordinator_artifact_root.text().strip()).expanduser() - base_url = f"{self.coordinator_public_url.text().strip().rstrip('/')}/artifacts" - bundle_path = create_job_artifact_bundle(job, artifact_root=artifact_root, base_url=base_url) - self.job_manager.submit(job) - return job, bundle_path - - def _current_remote_training_job(self, training_mode: str = "pretrain", stage: str = "base") -> TrainingJobSpec: - """Build a remote-worker job from current training controls. - - Args: - training_mode: Trainer mode to publish. - stage: Higher-level stage label for job manager display. - - Returns: - Complete training job spec ready to bundle and queue. - - Raises: - FileNotFoundError: If the prepared dataset is missing. - ValueError: If model or training options are invalid. - """ - - dataset_dir = Path(self.train_data_dir.text().strip()) - if not dataset_dir.exists(): - raise FileNotFoundError(f"Prepared dataset folder does not exist: {dataset_dir}") - if not self._dataset_artifacts_exist(dataset_dir): - raise FileNotFoundError( - "Prepared dataset is missing tokenizer or token files. " - "Expected tokenizer.json plus train/val tokens in .npy or .json." - ) - vocab_size = self._current_training_vocab_size(dataset_dir) - if vocab_size <= 0: - raise ValueError("Could not determine tokenizer vocabulary size from the prepared dataset.") - resume_path = Path(self.resume_checkpoint.text()) if self.resume_checkpoint.text().strip() else None - if resume_path is None and self.resume_training.isChecked(): - resume_path = latest_checkpoint(self._training_output_dir_for_mode(training_mode) / "checkpoints") - model_config = self._current_model_config(vocab_size=vocab_size) - training_config = self._current_training_config(resume_path, training_mode=training_mode) - model_config.validate() - training_config.validate() - job = TrainingJobSpec.local( - dataset_dir, - model_config, - training_config, - metadata={ - "project_name": self.search_box.text().strip(), - "submitted_from": "desktop_ui", - "coordinator_url": self.coordinator_public_url.text().strip().rstrip("/"), - "training_mode": training_mode, - "training_stage": stage, - }, - ) - job.runtime = RuntimeSpec( - backend=BackendKind.REMOTE_CLIENT, - device=training_config.device, - tags=[training_config.device, "remote"], - ) - return job - - def _active_job_for_worker(self, worker_id: str) -> str: - """Return the active job ID for a worker. - - Args: - worker_id: Worker identifier. - - Returns: - Active job ID or empty string. - """ - - for managed in self.job_manager.list_jobs(): - if managed.assigned_worker_id == worker_id and managed.spec.status.value in {"assigned", "running", "paused", "stopping"}: - return managed.spec.job_id - return "" - - @staticmethod - def _metric_pair(value: Optional[int], total: Optional[int]) -> str: - """Format a metric pair. - - Args: - value: Current value. - total: Total value. - - Returns: - Display text. - """ - - if value is None: - return "-" - if total is None: - return str(value) - return f"{value}/{total}" - - @staticmethod - def _metric_float(value: Optional[float], suffix: str = "") -> str: - """Format a floating-point metric. - - Args: - value: Metric value. - suffix: Optional suffix. - - Returns: - Display text. - """ - - if value is None: - return "-" - return f"{value:.4g}{suffix}" - - def _init_telemetry_store(self, model_dir: Path) -> None: - """Create or reset the SQLite telemetry store for a training run. - - Args: - model_dir: Model output directory. - """ - - self.telemetry_db_path = initialize_store(model_dir) - self.telemetry_run_id = datetime.now().strftime("%Y%m%d_%H%M%S") - self.telemetry_latest_id = 0 - self.telemetry_latest_index = 0 - self.live_time_slider.setRange(0, 0) - self.live_time_slider.setValue(0) - self.live_timeline_label.setText("Timeline: live") - self.live_scrub_active = False - - def _record_live_metric(self, event: dict[str, Any]) -> None: - """Persist one live training metric event to SQLite. - - Args: - event: Training progress event. - """ - - if self.telemetry_db_path is None or not self.telemetry_run_id or event.get("step") is None: - return - self.telemetry_latest_id = insert_metric(self.telemetry_db_path, self.telemetry_run_id, event) - self.telemetry_latest_index += 1 - self.live_time_slider.blockSignals(True) - self.live_time_slider.setRange(0, self.telemetry_latest_index) - if not self.live_scrub_active: - self.live_time_slider.setValue(self.telemetry_latest_index) - self.live_timeline_label.setText("Timeline: live") - self.live_time_slider.blockSignals(False) - - def _load_existing_telemetry(self, model_dir: Path) -> None: - """Load the latest saved telemetry run for an opened project. - - Args: - model_dir: Model output directory that may contain ``training_telemetry.sqlite``. - """ - - db_path = telemetry_db_path(model_dir) - self.telemetry_db_path = db_path if db_path.exists() else None - self.telemetry_run_id = "" - self.telemetry_latest_id = 0 - self.telemetry_latest_index = 0 - self.live_scrub_active = False - self.live_time_slider.blockSignals(True) - self.live_time_slider.setRange(0, 0) - self.live_time_slider.setValue(0) - self.live_time_slider.blockSignals(False) - self.live_timeline_label.setText("Timeline: no saved telemetry") - self.live_sample_text.setText("Training text: -") - if self.telemetry_db_path is None: - return - try: - run_row = latest_run(self.telemetry_db_path) - if run_row is None: - self.live_timeline_label.setText("Timeline: no samples") - return - self.telemetry_run_id = str(run_row["run_id"]) - self.telemetry_latest_index = int(run_row["sample_count"] or 0) - self.telemetry_latest_id = int(run_row["latest_id"] or 0) - except sqlite3.Error as exc: - self.live_timeline_label.setText("Timeline: could not load") - self.training_log.append(f"Telemetry load warning: {exc}") - return - self.live_time_slider.blockSignals(True) - self.live_time_slider.setRange(0, self.telemetry_latest_index) - self.live_time_slider.setValue(self.telemetry_latest_index) - self.live_time_slider.blockSignals(False) - if self.telemetry_latest_index: - rows = self._timeline_rows_until(self.telemetry_latest_index) - if rows: - self._apply_timeline_rows(rows) - - def _timeline_rows_until(self, sample_index: int) -> list[sqlite3.Row]: - """Load telemetry rows up to a selected sample index. - - Args: - sample_index: Maximum number of samples to load for the active run. - - Returns: - Ordered telemetry rows for the active run. - """ - - if self.telemetry_db_path is None or not self.telemetry_run_id or sample_index <= 0: - return [] - return rows_until(self.telemetry_db_path, self.telemetry_run_id, sample_index) - - def _begin_live_scrub(self) -> None: - """Pause live auto-follow while the timeline slider is being dragged.""" - - self.live_scrub_active = True - - def _end_live_scrub(self) -> None: - """Apply the selected timeline snapshot after slider drag.""" - - self._scrub_live_timeline(self.live_time_slider.value()) - - def _jump_live_timeline_to_latest(self) -> None: - """Return timeline display to the latest live point.""" - - self.live_scrub_active = False - self.live_time_slider.setValue(self.telemetry_latest_index) - self._scrub_live_timeline(self.telemetry_latest_index) - self.live_timeline_label.setText("Timeline: live") - - def _scrub_live_timeline(self, sample_index: int) -> None: - """Replay charts and live visual widgets to a selected telemetry point. - - Args: - sample_index: Timeline sample selected by the slider. - """ - - rows = self._timeline_rows_until(sample_index) - if not rows: - return - self._apply_timeline_rows(rows) - - def _apply_timeline_rows(self, rows: list[sqlite3.Row]) -> None: - """Apply historical telemetry rows to charts and live widgets. - - Args: - rows: Ordered SQLite telemetry rows. - """ - - def series(name: str) -> list[tuple[int, float]]: - return [(int(row["step"]), float(row[name])) for row in rows if row[name] is not None] - - latest = rows[-1] - self.loss_chart.set_points(series("train_loss"), series("val_loss")) - self.optimization_chart.set_points(series("learning_rate"), series("grad_norm")) - self.stability_chart.set_points(series("weight_norm"), series("update_ratio")) - self.throughput_chart.set_points(series("tokens_per_second"), series("samples_per_second")) - self.memory_chart.set_points(series("vram_allocated_gb"), series("vram_reserved_gb")) - snapshot = {key: latest[key] for key in latest.keys()} - sample_text = str(snapshot.get("sample_text") or "").strip() - if sample_text: - self.live_sample_text.setText(f"Training text: {self._compact_preview_text(sample_text, 220)}") - else: - self.live_sample_text.setText("Training text: -") - self._update_live_training_metrics( - int(latest["step"]), - snapshot, - snapshot.get("train_loss"), - snapshot.get("learning_rate"), - snapshot.get("grad_norm"), - snapshot.get("update_ratio"), - snapshot.get("tokens_per_second"), - snapshot.get("samples_per_second"), - snapshot.get("vram_allocated_gb"), - snapshot.get("vram_reserved_gb"), - snapshot.get("gpu_memory_percent"), - snapshot.get("system_cpu_percent"), - snapshot.get("system_ram_percent"), - snapshot.get("data_loader_workers"), - ) - timestamp = datetime.fromtimestamp(float(latest["recorded_at"])).strftime("%H:%M:%S") - self.live_timeline_label.setText(f"Timeline: step {int(latest['step']):,} @ {timestamp}") - - @staticmethod - def _compact_preview_text(text: str, limit: int = 220) -> str: - """Normalize a training preview into a compact single line. - - Args: - text: Raw decoded preview text. - limit: Maximum number of displayed characters. - - Returns: - Single-line text preview. - """ - - compact = re.sub(r"\s+", " ", text).strip() - if len(compact) <= limit: - return compact - return compact[: max(0, limit - 3)].rstrip() + "..." - - def _build_export_tab(self) -> QWidget: - """Build the export page. - - Returns: - Export page widget. - """ - - return build_export_tab(self) - - def _build_benchmark_tab(self) -> QWidget: - """Build the benchmark prompt page. - - Returns: - Benchmark page widget. - """ - - return build_benchmark_tab(self) - - def _build_chat_tab(self) -> QWidget: - """Build the model test chat page. - - Returns: - Chat page widget. - """ - - return build_chat_tab(self) - - def _panel(self) -> QWidget: - """Create a base page panel. - - Returns: - Panel widget. - """ - - page = QWidget() - page.setObjectName("Panel") - return page - - def _page_title(self, text: str) -> QLabel: - """Create a page title label. - - Args: - text: Title text. - - Returns: - Label configured as a page title. - """ - - label = QLabel(text) - label.setObjectName("PageTitle") - return label - - def _metric_chip(self, text: str, tooltip: str) -> QLabel: - """Create a compact metric display label. - - Args: - text: Initial metric text. - tooltip: User-facing explanation. - - Returns: - Configured metric label. - """ - - label = QLabel(text) - label.setObjectName("MetricChip") - label.setMinimumWidth(150) - label.setMinimumHeight(28) - label.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed) - self._tip(label, tooltip) - return label - - def _hardware_meter(self, name: str) -> QProgressBar: - """Create a slider-like hardware utilization meter. - - Args: - name: Display name for the meter. - - Returns: - Configured progress bar. - """ - - meter = QProgressBar() - meter.setObjectName("HardwareMeter") - meter.setRange(0, 100) - meter.setValue(0) - meter.setTextVisible(False) - meter.setFixedHeight(8) - meter.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed) - self._tip(meter, f"Live {name} utilization.") - return meter - - def _set_meter(self, meter: QProgressBar, name: str, value: Optional[float]) -> None: - """Update a hardware utilization meter. - - Args: - meter: Meter to update. - name: Display name for the meter. - value: Utilization percentage. - """ - - if value is None: - meter.setValue(0) - label = self.hardware_meter_labels.get(id(meter)) - if label is not None: - label.setText(f"{name}: -") - return - bounded = max(0.0, min(100.0, float(value))) - meter.setValue(int(round(bounded))) - label = self.hardware_meter_labels.get(id(meter)) - if label is not None: - label.setText(f"{name}: {bounded:.1f}%") - - def _update_dataset_quality_report(self, summary: dict[str, Any]) -> None: - """Update dataset quality chips from a summary dictionary. - - Args: - summary: Dataset summary fields. - """ - - document_count = int(summary.get("document_count", 0) or 0) - token_count = int(summary.get("token_count", 0) or 0) - train_window_count = int(summary.get("train_window_count", 0) or 0) - val_window_count = int(summary.get("val_window_count", 0) or 0) - character_count = int(summary.get("character_count", 0) or 0) - vocab_size = int(summary.get("tokenizer_vocab_size", summary.get("vocab_size", 0)) or 0) - code_count = int(summary.get("code_sample_count", 0) or 0) - prose_count = int(summary.get("prose_sample_count", 0) or 0) - conversation_count = int(summary.get("conversation_sample_count", 0) or 0) - cached_count = int(summary.get("cached_file_count", 0) or 0) - processed_count = int(summary.get("processed_file_count", 0) or 0) - skipped_count = int(summary.get("skipped_file_count", 0) or 0) - failed_count = int(summary.get("failed_file_count", 0) or 0) - warning = str(summary.get("warning") or "none") - sequence_stats = summary.get("sequence_token_stats", {}) or {} - quality_score = float(summary.get("quality_score", 0.0) or 0.0) - quality_stars = float(summary.get("quality_stars", 0.0) or 0.0) - quality_label = str(summary.get("quality_label") or "") - corpus_block_count = int(summary.get("corpus_block_count", 0) or 0) - unique_block_count = int(summary.get("unique_block_count", 0) or 0) - duplicate_block_count = int(summary.get("duplicate_block_count", 0) or 0) - duplicate_block_ratio = float(summary.get("duplicate_block_ratio", 0.0) or 0.0) - if not quality_score and (token_count or train_window_count or vocab_size): - quality_score, quality_stars, quality_label = self._estimate_dataset_rating( - token_count, - vocab_size, - train_window_count, - val_window_count, - document_count, - code_count, - prose_count, - conversation_count, - skipped_count, - failed_count, - warning, - sequence_stats, - ) - self.dataset_quality_samples.setText(f"Documents: {document_count:,}") - self.dataset_quality_tokens.setText(f"Tokens: {token_count:,}") - if train_window_count or val_window_count: - self.dataset_quality_windows.setText(f"Windows: {train_window_count:,}/{val_window_count:,}") - else: - self.dataset_quality_windows.setText("Windows: -") - self.dataset_quality_vocab.setText(f"Vocab: {vocab_size:,}" if vocab_size else "Vocab: -") - self.dataset_quality_rating.setText( - f"Rating: {self._star_text(quality_stars)} {quality_stars:.1f}/5" - if quality_stars - else "Rating: -" - ) - self.dataset_quality_code.setText(f"Code/prose/chat: {code_count:,}/{prose_count:,}/{conversation_count:,}") - self.dataset_quality_balance.setText("Balance: prepared") - self.dataset_quality_readiness.setText("Readiness: preview needed") - self.dataset_quality_cache.setText(f"Files: {processed_count:,} ok, {cached_count:,} cached, {skipped_count:,} skipped, {failed_count:,} failed") - if corpus_block_count: - self.dataset_quality_duplicates.setText(f"Duplicates: {duplicate_block_ratio * 100:.1f}%") - self._tip( - self.dataset_quality_duplicates, - ( - f"{duplicate_block_count:,} repeated blocks out of {corpus_block_count:,}; " - f"{unique_block_count:,} unique blocks." - ), - ) - else: - self.dataset_quality_duplicates.setText("Duplicates: -") - self.dataset_quality_warning.setText(f"Warnings: {warning}") - self._tip(self.dataset_quality_samples, f"{character_count:,} source characters across prepared documents.") - if quality_stars: - self._tip( - self.dataset_quality_rating, - f"{quality_label or 'Rated'} dataset: {quality_score:.1f}/100. Higher scores usually mean more usable tokens, richer vocabulary, more windows, and fewer extraction issues.", - ) - self._tip( - self.dataset_quality_windows, - f"{train_window_count:,} training and {val_window_count:,} validation sliding windows.", - ) - self._update_dataset_stat_charts(summary, code_count, prose_count, conversation_count, sequence_stats) - if hasattr(self, "dataset_advisor") and (train_window_count or val_window_count): - advice = [ - "Documents are source items. Windows are the actual context slices used by training.", - f"This dataset can provide about {train_window_count:,} training windows and {val_window_count:,} validation windows.", - ] - if sequence_stats: - advice.append( - "Approx token distribution per source: " - f"min {int(sequence_stats.get('min', 0) or 0):,}, " - f"avg {float(sequence_stats.get('average', 0.0) or 0.0):,.0f}, " - f"median {float(sequence_stats.get('median', 0.0) or 0.0):,.0f}, " - f"max {int(sequence_stats.get('max', 0) or 0):,}." - ) - if document_count < 100 and train_window_count >= 10_000: - advice.append( - "A low document count can still be useful when each document is long, because the trainer samples many overlapping windows." - ) - if train_window_count < 1_000: - advice.append("Add more text or lower context length if training looks repetitive.") - if corpus_block_count: - advice.append( - f"Block diversity: {unique_block_count:,}/{corpus_block_count:,} unique blocks " - f"({duplicate_block_ratio * 100:.1f}% repeated)." - ) - if quality_stars: - advice.append(f"Dataset rating: {quality_stars:.1f}/5 stars ({quality_label or 'rated'}, score {quality_score:.1f}/100).") - for reason in list(summary.get("quality_reasons", []) or [])[:4]: - advice.append(f"- {reason}") - self.dataset_advisor.setPlainText("\n".join(advice)) - - def _star_text(self, stars: float) -> str: - """Return a compact five-star display string. - - Args: - stars: Rating from zero to five. - - Returns: - Unicode star display with rounded whole stars. - """ - - whole = max(0, min(5, int(round(float(stars))))) - return "★" * whole + "☆" * (5 - whole) - - def _estimate_dataset_rating( - self, - token_count: int, - vocab_size: int, - train_window_count: int, - val_window_count: int, - document_count: int, - code_count: int, - prose_count: int, - conversation_count: int, - skipped_count: int, - failed_count: int, - warning: str, - sequence_stats: dict[str, Any], - ) -> tuple[float, float, str]: - """Estimate a dataset rating for older summaries that lack saved quality fields. - - Args: - token_count: Total prepared token count. - vocab_size: Tokenizer vocabulary size. - train_window_count: Number of training windows. - val_window_count: Number of validation windows. - document_count: Number of source documents. - code_count: Code sample count. - prose_count: Prose sample count. - conversation_count: Conversation/instruction sample count. - skipped_count: Skipped source file count. - failed_count: Failed source file count. - warning: Dataset warning text. - sequence_stats: Approximate source sequence statistics. - - Returns: - Score, stars, and label. - """ - - def ratio(value: float, target: float) -> float: - return max(0.0, min(1.0, float(value) / float(target))) if target > 0 else 0.0 - - families = sum(1 for count in (code_count, prose_count, conversation_count) if count > 0) - score = ( - 30.0 * ratio(token_count, 1_000_000) - + 20.0 * ratio(train_window_count, 50_000) - + 18.0 * ratio(vocab_size, 8_000) - + 12.0 * ratio(document_count, 1_000) - + 8.0 * ratio(val_window_count, 2_000) - + 7.0 * ratio(families, 3) - + 5.0 * ratio(float(sequence_stats.get("average", 0.0) or 0.0), 256) - ) - score -= min(20.0, failed_count * 3.0 + skipped_count * 0.5) - if warning and warning != "none": - score -= 5.0 - score = max(0.0, min(100.0, score)) - stars = round(score / 20.0 * 2.0) / 2.0 - if score >= 85: - label = "Excellent" - elif score >= 70: - label = "Good" - elif score >= 50: - label = "Usable" - elif score >= 30: - label = "Weak" - else: - label = "Very weak" - return score, stars, label - - def _update_dataset_stat_charts( - self, - summary: dict[str, Any], - code_count: int, - prose_count: int, - conversation_count: int, - sequence_stats: dict[str, Any], - ) -> None: - """Update dataset statistics charts. - - Args: - summary: Dataset summary fields. - code_count: Number of code samples. - prose_count: Number of prose samples. - conversation_count: Number of conversation or instruction samples. - sequence_stats: Approximate token distribution statistics. - """ - - if not hasattr(self, "dataset_mix_chart"): - return - mixture_report = summary.get("mixture_report", {}) or {} - family_rows = list((mixture_report.get("families", {}) or {}).values()) - labels: list[str] = [] - values: list[float] = [] - for row in family_rows: - actual = float(row.get("actual_percent", 0.0) or 0.0) - selected = int(row.get("selected_documents", 0) or 0) - if actual > 0.0 or selected > 0: - labels.append(str(row.get("label") or "source")) - values.append(actual) - if not labels: - total = max(code_count + prose_count + conversation_count, 1) - labels = ["Code", "Prose", "Conversation"] - values = [ - code_count * 100.0 / total, - prose_count * 100.0 / total, - conversation_count * 100.0 / total, - ] - self.dataset_mix_chart.set_values(labels, values, "%") - if sequence_stats: - self.dataset_sequence_chart.set_values( - ["Min", "Average", "Median", "Max"], - [ - float(sequence_stats.get("min", 0) or 0), - float(sequence_stats.get("average", 0.0) or 0.0), - float(sequence_stats.get("median", 0.0) or 0.0), - float(sequence_stats.get("max", 0) or 0), - ], - ) - else: - self.dataset_sequence_chart.clear() - - def _reset_dataset_quality_report(self) -> None: - """Reset dataset quality chips to their empty state.""" - - self.dataset_quality_samples.setText("Documents: -") - self.dataset_quality_tokens.setText("Tokens: -") - self.dataset_quality_windows.setText("Windows: -") - self.dataset_quality_vocab.setText("Vocab: -") - self.dataset_quality_rating.setText("Rating: -") - self.dataset_quality_code.setText("Code/prose: -") - self.dataset_quality_balance.setText("Balance: -") - self.dataset_quality_readiness.setText("Readiness: -") - self.dataset_quality_cache.setText("Cache: -") - self.dataset_quality_duplicates.setText("Duplicates: -") - self.dataset_quality_extraction.setText("Extraction: -") - self.dataset_quality_warning.setText("Warnings: none") - if hasattr(self, "dataset_mix_chart"): - self.dataset_mix_chart.clear() - if hasattr(self, "dataset_sequence_chart"): - self.dataset_sequence_chart.clear() - if hasattr(self, "dataset_advisor"): - self.dataset_advisor.setPlainText("Run Preview Dataset to get cleanup suggestions.") - - def _card(self, title: str, content_layout: Union[QVBoxLayout, QFormLayout, QGridLayout, QHBoxLayout]) -> QWidget: - """Create a neon module card. - - Args: - title: Card heading. - content_layout: Layout to place inside the card. - - Returns: - Card widget. - """ - - card = QWidget() - card.setObjectName("Card") - layout = QVBoxLayout(card) - layout.setContentsMargins(14, 12, 14, 12) - layout.setSpacing(8) - title_label = QLabel(title) - title_label.setObjectName("SectionLabel") - layout.addWidget(title_label) - layout.addLayout(content_layout) - return card - - def _spin(self, minimum: int, maximum: int, value: int) -> QSpinBox: - """Create a bounded integer input. - - Args: - minimum: Minimum value. - maximum: Maximum value. - value: Initial value. - - Returns: - Configured spin box. - """ - - spin = QSpinBox() - spin.setRange(minimum, maximum) - spin.setValue(value) - spin.setMaximumWidth(220) - return spin - - def _double_spin(self, minimum: float, maximum: float, value: float, step: float, decimals: int) -> QDoubleSpinBox: - """Create a bounded float input. - - Args: - minimum: Minimum value. - maximum: Maximum value. - value: Initial value. - step: Increment step. - decimals: Number of displayed decimal places. - - Returns: - Configured double spin box. - """ - - spin = QDoubleSpinBox() - spin.setRange(minimum, maximum) - spin.setDecimals(decimals) - spin.setSingleStep(step) - spin.setValue(value) - spin.setMaximumWidth(220) - return spin - - def _path_row(self, field: QLineEdit, directory: bool = True, file_filter: str = "Checkpoints (*.pt)") -> QWidget: - """Create a path field with a browse button. - - Args: - field: Path input widget. - directory: Whether the browse dialog selects folders. - file_filter: File dialog filter used when ``directory`` is false. - - Returns: - Row widget containing the path input and button. - """ - - row = QWidget() - row.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed) - layout = QHBoxLayout(row) - layout.setContentsMargins(0, 0, 0, 0) - layout.setSpacing(8) - browse = QPushButton("Browse") - browse.setFixedWidth(88) - self._tip(browse, "Open a file/folder picker for this path.") - field.setMinimumWidth(180) - field.setSizePolicy(QSizePolicy.Ignored, QSizePolicy.Fixed) - browse.clicked.connect(lambda: self._browse(field, directory, file_filter)) - layout.addWidget(field, 1) - layout.addWidget(browse) - return row - - def _multi_file_path_row(self, field: QLineEdit, file_filter: str = "All files (*)") -> QWidget: - """Create a path field with a multi-file browse button. - - Args: - field: Path input widget. Multiple paths are separated with semicolons. - file_filter: File dialog filter. - - Returns: - Row widget containing the path input and browse button. - """ - - row = QWidget() - row.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed) - layout = QHBoxLayout(row) - layout.setContentsMargins(0, 0, 0, 0) - layout.setSpacing(8) - browse = QPushButton("Browse") - browse.setFixedWidth(88) - self._tip(browse, "Choose one or more JSON/JSONL files. You can also paste a folder path.") - field.setMinimumWidth(180) - field.setSizePolicy(QSizePolicy.Ignored, QSizePolicy.Fixed) - browse.clicked.connect(lambda: self._browse_multiple_files(field, file_filter)) - layout.addWidget(field, 1) - layout.addWidget(browse) - return row - - def _configure_form(self, form: QFormLayout) -> None: - """Apply common form spacing and growth policy. - - Args: - form: Form layout to configure. - """ - - form.setLabelAlignment(Qt.AlignRight | Qt.AlignVCenter) - form.setFormAlignment(Qt.AlignLeft | Qt.AlignTop) - form.setFieldGrowthPolicy(QFormLayout.AllNonFixedFieldsGrow) - form.setHorizontalSpacing(12) - form.setVerticalSpacing(7) - - def _configure_device_options(self) -> None: - """Populate training device choices without duplicate CPU entries.""" - - self.device.clear() - if torch.cuda.is_available(): - device_name = torch.cuda.get_device_name(0) - self.device.addItem("cuda") - self.device.addItem("cpu") - self.device_info.setText(f"CUDA ready: {device_name}") - self.use_amp_default = True - else: - self.device.addItem("cpu") - cuda_build = getattr(torch.backends, "cuda", None) - built_with_cuda = bool(cuda_build and torch.backends.cuda.is_built()) - if built_with_cuda: - detail = "CUDA build found, but no usable NVIDIA GPU/driver was detected." - else: - detail = "CUDA is not available in this PyTorch install." - self.device_info.setText(detail) - self.use_amp_default = False - - def _thin_progress(self) -> QProgressBar: - """Create a thin bottom progress bar. - - Returns: - Configured progress bar. - """ - - progress = QProgressBar() - progress.setRange(0, 100) - progress.setTextVisible(False) - progress.setFixedHeight(4) - progress.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed) - self._tip(progress, "Progress indicator for the current page operation.") - return progress - - def _tip(self, widget: QWidget, text: str) -> None: - """Attach tooltip and status tip text. - - Args: - widget: Widget receiving the tip. - text: Tooltip text. - """ - - widget.setToolTip(text) - widget.setStatusTip(text) - - def _render_chat_markdown(self, markdown_text: str) -> None: - """Render chat Markdown with highlighted fenced code blocks when possible. - - Args: - markdown_text: Markdown transcript to render. - """ - - if not hasattr(self, "current_assistant_message") or self.current_assistant_message is None: - return - self.current_assistant_message.set_content(markdown_text) - - def _add_chat_message( - self, - role: str, - content: str, - metrics: str = "", - resend_prompt: Optional[str] = None, - ) -> QTextBrowser: - """Add one chat bubble. - - Args: - role: Message role, either ``user`` or ``assistant``. - content: Markdown message content. - metrics: Optional metric text shown under assistant replies. - resend_prompt: Prompt to resend from the bubble. - - Returns: - Text browser used by the bubble. - """ - - should_follow = self._is_chat_near_bottom() - max_width = max(320, int(self.chat_scroll.viewport().width() * 0.78)) if hasattr(self, "chat_scroll") else 900 - message = ChatMessageWidget( - role, - content, - markdown_to_html, - self._resend_chat_message, - metrics=metrics, - resend_prompt=resend_prompt, - max_width=max_width, - ) - self.chat_messages.insertWidget(max(self.chat_messages.count() - 1, 0), message) - if should_follow: - message.scroll_later(lambda: self.chat_scroll.verticalScrollBar().setValue(self.chat_scroll.verticalScrollBar().maximum())) - if role == "assistant": - self.current_assistant_message = message - self.current_assistant_browser = message.browser - self.current_assistant_meta = message.meta_label - return message.browser - - def _is_chat_near_bottom(self) -> bool: - """Return whether the chat scroll is close enough to follow streaming. - - Returns: - True when the view should auto-scroll. - """ - - if not hasattr(self, "chat_scroll"): - return True - bar = self.chat_scroll.verticalScrollBar() - return bar.maximum() - bar.value() < 48 - - def _clear_chat_messages(self) -> None: - """Remove all message bubbles.""" - - while self.chat_messages.count() > 1: - item = self.chat_messages.takeAt(0) - widget = item.widget() - if widget is not None: - widget.deleteLater() - self.current_assistant_message = None - self.current_assistant_browser = None - self.current_assistant_meta = None - - def _resend_chat_message(self, prompt: str) -> None: - """Resend text from a message bubble. - - Args: - prompt: Prompt text to send. - """ - - self.chat_input.setPlainText(prompt) - self.send_chat_message() - - def _set_chat_stats(self, elapsed_seconds: float, token_count: int, tokens_per_second: float) -> None: - """Update live chat generation metrics. - - Args: - elapsed_seconds: Elapsed generation time. - token_count: Generated token count. - tokens_per_second: Approximate token speed. - """ - - text = f"Time: {elapsed_seconds:.2f}s | Tokens: {token_count:,} | Speed: {tokens_per_second:.2f} tok/s" - self.chat_stats.setText(text) - if self.current_assistant_meta is not None: - self.current_assistant_meta.setText(text) - self.current_assistant_meta.setVisible(True) - - def _chat_backend_value(self) -> str: - """Return the selected chat model backend. - - Returns: - Stable chat backend identifier. - """ - - if not hasattr(self, "chat_model_backend"): - return "gguf" - return "microgpt" if self.chat_model_backend.currentText() == "MicroGPT checkpoint" else "gguf" - - def _update_chat_backend_controls(self) -> None: - """Show controls relevant to the selected chat backend.""" - - if not hasattr(self, "chat_model_backend"): - return - native = self._chat_backend_value() == "microgpt" - self.gguf_path_row.setVisible(not native) - self.microgpt_path_row.setVisible(native) - self.llama_gpu_layers.setEnabled(not native) - self.llama_threads.setEnabled(not native) - self.llama_context.setEnabled(not native) - if native: - self._tip(self.load_llm_button, "Load the native MicroGPT checkpoint into memory once for repeated chat messages.") - else: - self._tip(self.load_llm_button, "Load the GGUF model into memory once for repeated chat messages.") - - def _app_icon(self) -> QIcon: - """Create the application icon. - - Returns: - Application icon. - """ - - return self._static_app_icon() - - @staticmethod - def _app_logo_path() -> Path: - """Return the bundled logo path. - - Returns: - Logo path. - """ - - 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"): - candidates.insert(0, Path(sys._MEIPASS) / "drunken_bot_logo_small.png") - app_root = os.environ.get("DRUNKENBOT_APP_ROOT") - if app_root: - candidates.insert(0, Path(app_root) / "drunken_bot_logo_small.png") - return next((path for path in candidates if path.exists()), candidates[0]) - - @staticmethod - def _app_logo_pixmap(size: int = 64) -> QPixmap: - """Load the bundled logo as a pixmap. - - Args: - size: Maximum square size. - - Returns: - Logo pixmap, or null pixmap when the file is missing. - """ - - pixmap = QPixmap(str(MainWindow._app_logo_path())) - if pixmap.isNull(): - return pixmap - return pixmap.scaled(size, size, Qt.KeepAspectRatio, Qt.SmoothTransformation) - - @staticmethod - def _static_app_icon() -> QIcon: - """Create the static app icon. - - Returns: - Application icon. - """ - - logo_path = MainWindow._app_logo_path() - if logo_path.exists(): - icon = QIcon(str(logo_path)) - if not icon.isNull(): - return icon - pixmap = QPixmap(64, 64) - pixmap.fill(Qt.transparent) - painter = QPainter(pixmap) - try: - painter.setRenderHint(QPainter.Antialiasing) - painter.setBrush(QBrush(QColor("#1f1f1f"))) - painter.setPen(QPen(QColor("#f5b041"), 3)) - painter.drawRoundedRect(4, 4, 56, 56, 12, 12) - bolt = QPolygon([ - QPoint(36, 8), - QPoint(17, 35), - QPoint(31, 35), - QPoint(25, 56), - QPoint(48, 25), - QPoint(33, 25), - ]) - painter.setPen(QPen(QColor("#ffd27a"), 2)) - painter.setBrush(QBrush(QColor("#f5b041"))) - painter.drawPolygon(bolt) - finally: - painter.end() - return QIcon(pixmap) - - @staticmethod - def _windows_icon_path() -> Path: - """Return the Windows icon file path. - - Returns: - Path to the generated ``.ico`` file. - """ - - return Path(__file__).with_name("drunkenbot_llm_ide.ico") - - @staticmethod - def _ensure_windows_icon_file() -> Optional[Path]: - """Ensure the generated Windows ``.ico`` file exists. - - Returns: - Icon path on Windows, otherwise ``None``. - """ - - if sys.platform != "win32": - return None - icon_path = MainWindow._windows_icon_path() - if icon_path.exists(): - return icon_path - icon = MainWindow._static_app_icon() - pixmap = icon.pixmap(256, 256) - if pixmap.isNull() or not pixmap.save(str(icon_path), "ICO"): - return None - return icon_path - - def apply_windows_taskbar_icon(self) -> None: - """Apply the app icon to the native Windows window handle.""" - - if sys.platform != "win32": - return - icon_path = self._ensure_windows_icon_file() - if icon_path is None: - return - - hwnd = int(self.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) - self._windows_icon_handles.append(hicon_big) - if hicon_small: - user32.SendMessageW(hwnd, wm_seticon, icon_small, hicon_small) - self._windows_icon_handles.append(hicon_small) - - def _browse(self, field: QLineEdit, directory: bool, file_filter: str = "Checkpoints (*.pt)") -> None: - """Open a file or folder picker for a path field. - - Args: - field: Path input to update. - directory: Whether to select a folder instead of a file. - file_filter: File dialog filter used for files. - """ - - start_dir = self._browse_start_dir(field, directory) - if directory: - value = QFileDialog.getExistingDirectory(self, "Choose folder", start_dir) - else: - value, _ = QFileDialog.getOpenFileName(self, "Choose file", start_dir, file_filter) - if value: - field.setText(value) - - def _browse_multiple_files(self, field: QLineEdit, file_filter: str) -> None: - """Open a multi-file picker and write selected paths to a field. - - Args: - field: Path field to update. - file_filter: File dialog filter. - """ - - values, _ = QFileDialog.getOpenFileNames(self, "Choose files", self._browse_start_dir(field, False), file_filter) - if values: - field.setText("; ".join(values)) - - def _browse_start_dir(self, field: QLineEdit, directory: bool) -> str: - """Return the best initial folder for a browse dialog. - - Args: - field: Path field being browsed. - directory: Whether the dialog selects a folder. - - Returns: - Existing field path, active project folder, or current folder. - """ - - text = field.text().strip() - if text: - path = Path(text) - if path.exists(): - if path.is_dir(): - return str(path) - return str(path.parent) - parent = path if directory else path.parent - if parent.exists(): - return str(parent) - if self.current_project_file is not None: - return str(self.current_project_file.parent) - return str(Path.cwd()) - - def save_project(self) -> None: - """Save the current project settings into a named project folder.""" - - project_name = self.search_box.text().strip() or "MicroLLMProject" - safe_name = self._safe_project_name(project_name) - if self.current_project_file is None: - base_dir = QFileDialog.getExistingDirectory(self, "Choose parent folder for project", self._project_dialog_start_dir()) - if not base_dir: - return - project_dir = Path(base_dir) / safe_name - project_file = project_dir / "project.json" - else: - project_file = self.current_project_file - project_dir = project_file.parent - project_dir.mkdir(parents=True, exist_ok=True) - self._ensure_project_workspace(project_dir) - if self.current_project_file is None: - self._apply_project_workspace_paths(project_dir) - project_file.write_text(json.dumps(self._project_state_dict(project_name, project_dir), indent=2), encoding="utf-8") - self.current_project_file = project_file - _register_recent_project(project_file) - self._apply_project_runtime_environment(project_dir) - self._refresh_notification_manager(project_dir) - if hasattr(self, "runpod_api_key"): - self.load_runpod_settings() - self.project_state.setText("Project saved") - LOGGER.info("Project saved: %s", project_file) - if self.current_project_file == project_file: - self.dataset_log.append(f"Project saved: {project_file}") - self.dataset_log.append(f"Project workspace: {project_dir}") - self.dataset_log.append(f"Notifier config: {project_dir / 'notifier_config.json'}") - - def new_project(self) -> None: - """Start a fresh project and clear the active project file binding.""" - - if self.thread is not None: - QMessageBox.information(self, "Task running", "Please stop or wait for the current task before creating a new project.") - return - if self.current_project_file is not None or self.search_box.text().strip(): - choice = QMessageBox.question( - self, - "New project", - "Start a new project? Unsaved changes in the current project will not be saved automatically.", - QMessageBox.Yes | QMessageBox.No, - QMessageBox.No, - ) - if choice != QMessageBox.Yes: - return - - base_dir = QFileDialog.getExistingDirectory( - self, - "Choose folder where the new project will be created", - self._project_dialog_start_dir(), - ) - if not base_dir: - self.project_state.setText("Project creation cancelled") - return - if self.chat_session is not None and hasattr(self.chat_session, "reset"): - self.chat_session.reset() - self.chat_session = None - project_name = self.search_box.text().strip() or "MicroLLMProject" - try: - self._create_project_at(project_name, Path(base_dir)) - except Exception as exc: - QMessageBox.warning(self, "New project failed", f"Could not create project:\n{exc}") - - def open_project(self) -> None: - """Open a saved project file and restore UI settings.""" - - project_file, _ = QFileDialog.getOpenFileName( - self, - "Open Micro LLM project", - self._project_dialog_start_dir(), - "Micro LLM project (project.json *.json);;All files (*)", - ) - if not project_file: - return - try: - self._open_project_file(Path(project_file)) - except Exception as exc: - QMessageBox.warning(self, "Open failed", f"Could not open project:\n{exc}") - return - - def _create_project_at(self, project_name: str, base_dir: Path) -> Path: - """Create and activate a new project at the selected folder. - - Args: - project_name: User-facing project name. - base_dir: Parent folder for the new project. - - Returns: - Path to the created project.json file. - """ - - if self.chat_session is not None and hasattr(self.chat_session, "reset"): - self.chat_session.reset() - self.chat_session = None - project_dir = base_dir / self._safe_project_name(project_name) - project_file = project_dir / "project.json" - project_dir.mkdir(parents=True, exist_ok=True) - self._ensure_project_workspace(project_dir) - copied_count = self._ensure_project_training_data(project_dir) - self.current_project_file = project_file - self._apply_project_state(self._default_project_state()) - self.search_box.setText(project_name) - self._apply_project_workspace_paths(project_dir) - self._reset_dataset_blueprint_source(project_dir / "training_data") - self._apply_project_runtime_environment(project_dir) - self._refresh_notification_manager(project_dir) - if hasattr(self, "runpod_api_key"): - self.load_runpod_settings() - self._reset_project_runtime_state() - project_file.write_text(json.dumps(self._project_state_dict(project_name, project_dir), indent=2), encoding="utf-8") - _register_recent_project(project_file) - self.project_state.setText("New project") - LOGGER.info("New project created: %s", project_file) - self.dataset_log.append(f"Started a new project: {project_file}") - self.dataset_log.append(f"Project workspace: {project_dir}") - self.dataset_log.append( - "Bundled training data is no longer included; use Dataset Sources to download or select a dataset." - ) - self.dataset_log.append(f"Notifier config: {project_dir / 'notifier_config.json'}") - return project_file - - def _reset_dataset_blueprint_source(self, data_root: Path) -> None: - """Reset the current Dataset Sources tree without replacing its widget.""" - self.blueprint_data_root = Path(data_root) - if hasattr(self, "external_dataset_dir"): - self.external_dataset_dir.setText(str(data_root)) - if hasattr(self, "dataset_plan_source_label"): - self.dataset_plan_source_label.setText(f"Source: {data_root}") - if hasattr(self, "default_data_tree"): - self.default_data_tree.clear() - self.default_data_actions.clear() - self.default_data_category_items.clear() - self.default_data_tree.addTopLevelItem( - QTreeWidgetItem(["No project data files were found.", "", ""]) - ) - - def _open_project_file(self, project_file: Path) -> None: - """Open and activate a project file. - - Args: - project_file: Path to ``project.json``. - """ - - data = json.loads(project_file.read_text(encoding="utf-8")) - self.current_project_file = project_file - _register_recent_project(project_file) - self._ensure_project_workspace(self.current_project_file.parent) - dataset_state = data.get("dataset", {}) if isinstance(data, dict) else {} - saved_default_data_paths = dataset_state.get("default_data_paths") - self._refresh_dataset_blueprint_source( - self.current_project_file.parent / "training_data", - saved_paths=(list(saved_default_data_paths) if saved_default_data_paths is not None else None), - saved_plan=dict(dataset_state.get("domain_plan", {})), - preset=str(dataset_state.get("domain_plan_preset", "Balanced Tiny LLM")), - ) - self._apply_project_state(data) - self._apply_project_runtime_environment(self.current_project_file.parent) - self._refresh_notification_manager(self.current_project_file.parent) - if hasattr(self, "runpod_api_key"): - self.load_runpod_settings() - if self.model_dir.text().strip(): - self._load_existing_telemetry(Path(self.model_dir.text())) - self.project_state.setText("Project opened") - LOGGER.info("Project opened: %s", project_file) - self.dataset_log.append(f"Opened project: {project_file}") - self.dataset_log.append(f"Notifier config: {self.current_project_file.parent / 'notifier_config.json'}") - self.refresh_model_estimate() - - def _project_dialog_start_dir(self) -> str: - """Return the best initial folder for project dialogs. - - Returns: - Active project folder, its parent, or the current folder. - """ - - if self.current_project_file is not None: - return str(self.current_project_file.parent) - text = self.dataset_dir.text().strip() if hasattr(self, "dataset_dir") else "" - if text: - path = Path(text) - for candidate in (path, path.parent): - if candidate.exists(): - return str(candidate) - return str(Path.cwd()) - - def _ensure_project_workspace(self, project_dir: Path) -> None: - """Create standard folders inside a project. - - Args: - project_dir: Project root folder. - """ - - for name in ("datasets", "models", "fine_tunes", "exports", "training_data", "cache", "temp"): - (project_dir / name).mkdir(parents=True, exist_ok=True) - ensure_notifier_config(project_dir / "notifier_config.json") - ensure_runpod_config(project_dir / "runpod_config.json") - - def _ensure_project_training_data(self, project_dir: Path) -> int: - """Create the project training-data folder without bundling corpus files.""" - target_root = project_dir / "training_data" - target_root.mkdir(parents=True, exist_ok=True) - return 0 - - def _refresh_dataset_blueprint_source( - self, - data_root: Path, - saved_paths: Optional[list[Any]] = None, - saved_plan: Optional[dict[str, Any]] = None, - preset: str = "Balanced Tiny LLM", - ) -> None: - """Rebuild the Dataset Blueprint tab from a source data folder. - - Args: - data_root: Project-local training data folder. - saved_paths: Optional selected file paths to restore. - saved_plan: Optional saved domain recipe. - preset: Saved recipe preset. - """ - - if not hasattr(self, "pages"): - self.blueprint_data_root = Path(data_root) - return - self.blueprint_data_root = Path(data_root) - if hasattr(self, "default_data_tree"): - selected = saved_paths if saved_paths is not None else self._selected_default_data_paths() - populate_default_data_tree(self, self.blueprint_data_root) - self._set_selected_default_data_paths(selected) - if saved_plan is not None: - self._set_dataset_plan(saved_plan, preset) - self.dataset_plan_source_label.setText(f"Source: {self.blueprint_data_root}") - return - current_index = self.pages.currentIndex() - old_page = self.pages.widget(0) - old_page.hide() - QApplication.processEvents() - new_page = self._build_dataset_plan_tab() - self.pages.removeWidget(old_page) - old_page.setParent(None) - old_page.deleteLater() - self.pages.insertWidget(0, new_page) - if saved_plan is not None: - self._set_dataset_plan(saved_plan, preset) - if saved_paths is not None: - self._set_selected_default_data_paths(saved_paths) - elif self.current_project_file is not None: - self._set_selected_default_data_paths(None) - self.pages.setCurrentIndex(current_index) - - def download_latest_external_dataset(self) -> None: - """Download the latest managed dataset into the selected install folder.""" - destination = Path(self.external_dataset_dir.text()).expanduser() - try: - manifest = load_manifest() - except Exception as exc: - self.external_dataset_version.setText(f"Could not load dataset options: {exc}") - return - dialog = QDialog(self) - dialog.setWindowTitle("Select dataset components") - dialog.setMinimumWidth(480) - dialog_layout = QVBoxLayout(dialog) - dialog_layout.addWidget(QLabel(f"Dataset version {manifest.version}")) - version_selector = QComboBox() - version_selector.addItem(manifest.version, DEFAULT_MANIFEST_URL) - installed_version_file = destination / "version.txt" - if installed_version_file.is_file(): - installed_version = installed_version_file.read_text(encoding="utf-8").strip() - if installed_version and installed_version != manifest.version: - version_selector.addItem( - installed_version, - f"https://github.com/drunkenbot-ai/dataset/releases/download/dataset-v{installed_version}/manifest.json", - ) - dialog_layout.addWidget(QLabel("Dataset version")) - dialog_layout.addWidget(version_selector) - component_tree = QTreeWidget() - component_tree.setHeaderLabels(["Component", "Files"]) - for category in manifest.categories: - if category.file_count <= 0: - continue - item = QTreeWidgetItem([category.name, str(category.file_count)]) - item.setFlags(item.flags() | Qt.ItemIsUserCheckable) - existing = destination / category.name - has_existing = existing.exists() and any(existing.rglob("*")) - item.setCheckState(0, Qt.Checked if has_existing else Qt.Unchecked) - component_tree.addTopLevelItem(item) - dialog_layout.addWidget(component_tree) - buttons = QHBoxLayout() - download_button = QPushButton("Download selected") - cancel_button = QPushButton("Cancel") - buttons.addStretch(1) - buttons.addWidget(cancel_button) - buttons.addWidget(download_button) - dialog_layout.addLayout(buttons) - cancel_button.clicked.connect(dialog.reject) - download_button.clicked.connect(dialog.accept) - if dialog.exec() != QDialog.Accepted: - return - categories = [ - component_tree.topLevelItem(index).text(0) - for index in range(component_tree.topLevelItemCount()) - if component_tree.topLevelItem(index).checkState(0) == Qt.Checked - ] - if not categories: - self.external_dataset_version.setText("Select at least one dataset component.") - return - - self.dataset_log.append(f"Downloading latest external dataset to {destination}...") - self.external_dataset_version.setText("Downloading latest dataset...") - self.dataset_plan_progress.setVisible(True) - self._run_task( - partial(download_latest_dataset, manifest_url=version_selector.currentData()), - (destination, categories), - self._external_dataset_download_finished, - self.dataset_log, - self.dataset_plan_progress, - with_progress=True, - button=self.external_dataset_download_button, - busy_text="Downloading dataset", - task_kind="dataset_download", - ) - - def _refresh_external_dataset_status(self) -> None: - """Restore the installed dataset version from the selected folder.""" - if not hasattr(self, "external_dataset_dir"): - return - version_file = Path(self.external_dataset_dir.text()).expanduser() / "version.txt" - if not version_file.is_file(): - self.external_dataset_version.setText("Installed version: not installed") - self.external_dataset_download_button.setEnabled(True) - return - try: - version = version_file.read_text(encoding="utf-8").strip() - except OSError: - version = "" - if version: - has_dataset_files = any( - path.is_file() and path.name not in {"version.txt", "manifest.json"} - for path in Path(self.external_dataset_dir.text()).expanduser().rglob("*") - ) - if not has_dataset_files: - self.external_dataset_version.setText("Installed version: not installed") - self.external_dataset_download_button.setEnabled(True) - return - self.external_dataset_version.setText(f"Installed version: {version}") - try: - latest = load_manifest() - except Exception as exc: - LOGGER.warning("Could not check for a newer dataset release: %s", exc) - self.external_dataset_download_button.setEnabled(True) - return - installed_root = Path(self.external_dataset_dir.text()).expanduser() - missing_components = [ - category.name - for category in latest.categories - if category.file_count > 0 - and not ( - (installed_root / category.name).exists() - and any((installed_root / category.name).rglob("*")) - ) - ] - if is_newer_version(latest.version, version): - self.external_dataset_version.setText( - f"Installed version: {version} (update available: {latest.version})" - ) - self.external_dataset_download_button.setEnabled(True) - elif missing_components: - self.external_dataset_version.setText( - f"Installed version: {version} ({len(missing_components)} components missing)" - ) - self.external_dataset_download_button.setEnabled(True) - else: - self.external_dataset_download_button.setEnabled(False) - else: - self.external_dataset_version.setText("Installed version: not installed") - self.external_dataset_download_button.setEnabled(True) - - @Slot(object) - def _external_dataset_download_finished(self, manifest: Any) -> None: - """Apply the downloaded dataset as the active source vault.""" - self.dataset_log.append(f"Installed external dataset version {manifest.version}.") - self.external_dataset_version.setText(f"Installed version: {manifest.version}") - self.external_dataset_download_button.setEnabled(False) - self.dataset_plan_progress.setVisible(False) - self._refresh_dataset_blueprint_source( - Path(self.external_dataset_dir.text()), - saved_plan=self._dataset_plan_from_ui(), - preset=( - self.dataset_plan_preset.currentText() - if hasattr(self, "dataset_plan_preset") - else "Balanced Tiny LLM" - ), - ) - self.input_dir.setText(self.external_dataset_dir.text()) - - def _refresh_notification_manager(self, project_dir: Optional[Path] = None) -> None: - """Load notification settings for the current project. - - Args: - project_dir: Optional project root folder. - """ - - if project_dir is None and self.current_project_file is not None: - project_dir = self.current_project_file.parent - config_path = default_notifier_config_path(project_dir) - self.notification_manager = NotificationManager(config_path) - LOGGER.info("Notifier config active: %s", config_path) - - def _apply_project_workspace_paths(self, project_dir: Path) -> None: - """Point project output fields at the standard project folders. - - Args: - project_dir: Project root folder. - """ - - dataset_dir = project_dir / "datasets" - model_dir = project_dir / "models" - fine_tune_dir = project_dir / "fine_tunes" - export_dir = project_dir / "exports" - training_data_dir = project_dir / "training_data" - self.dataset_dir.setText(str(dataset_dir)) - self.train_data_dir.setText(str(dataset_dir)) - self.model_dir.setText(str(model_dir)) - self.fine_tune_checkpoint.setText(str(model_dir / "final_model.pt")) - self.fine_tune_output_dir.setText(str(fine_tune_dir / "latest")) - self.export_model_dir.setText(str(model_dir)) - self.export_dir.setText(str(export_dir)) - self.gguf_output_path.setText(str(export_dir / "model.gguf")) - if not self.input_dir.text().strip(): - self.input_dir.setText(str(training_data_dir)) - - def _apply_project_runtime_environment(self, project_dir: Path) -> None: - """Prefer project-local cache/temp folders for runtime work. - - Args: - project_dir: Project root folder. - """ - - cache_dir = project_dir / "cache" - temp_dir = project_dir / "temp" - cache_dir.mkdir(parents=True, exist_ok=True) - temp_dir.mkdir(parents=True, exist_ok=True) - for key in ("TMPDIR", "TEMP", "TMP"): - os.environ[key] = str(temp_dir) - for key in ("TORCH_HOME", "HF_HOME", "TRANSFORMERS_CACHE", "PYTORCH_KERNEL_CACHE"): - os.environ[key] = str(cache_dir / key.lower()) - Path(os.environ[key]).mkdir(parents=True, exist_ok=True) - - def _default_project_state(self) -> dict[str, Any]: - """Build the default state used for a newly created project. - - Returns: - JSON-style project state with fresh paths and default settings. - """ - - runs_dir = Path.cwd() / "runs" - dataset_dir = runs_dir / "dataset" - model_dir = runs_dir / "model" - fine_tune_dir = runs_dir / "fine_tune" - export_dir = runs_dir / "export" - return { - "schema": "drunkenbot_ide_project", - "version": 1, - "project_name": "", - "project_dir": "", - "paths": { - "source_vault": "", - "dataset_core": str(dataset_dir), - "training_dataset": str(dataset_dir), - "model_output": str(model_dir), - "export_model_core": str(model_dir), - "export_output": str(export_dir), - "llama_cpp_dir": "", - "gguf_output_path": str(export_dir / "model.gguf"), - "gguf_model": "", - "microgpt_chat_model": "", - "tokenizer_import": "", - "resume_checkpoint": "", - "fine_tune_checkpoint": "", - "fine_tune_output": str(fine_tune_dir), - }, - "dataset": { - "domain_plan_preset": "Balanced Tiny LLM", - "domain_plan": dataset_plan_defaults(), - "default_data_paths": [str(path) for path, _category in iter_default_data_files()], - "auto_vocab": True, - "manual_vocab_size": 8000, - "include_conversation_datasets": False, - "dataset_stage": "base", - "conversation_datasets": [], - "conversation_sample_limit": 20000, - "conversation_dataset_path": "", - "instruction_dataset_path": "", - "mixture_weights": {}, - "min_frequency": 2, - "context_length": 128, - "validation_split": 0.1, - "lowercase": False, - "max_workers": 4, - "prepare_mode": "incremental", - "tokenizer_strategy": "auto", - "code_training_mode": True, - "include_prose": True, - "include_source_code": True, - "extract_code_blocks": True, - "preserve_indentation": True, - "instruction_samples": True, - "reasoning_sample_mode": "scaffold", - }, - "training": { - "preset": "Tiny", - "architecture_style": "Classic GPT", - "launch_target": "local", - "training_mode": "pretrain", - "training_stage": "base", - "peft_method": "none", - "lora_rank": 8, - "lora_alpha": 16.0, - "lora_dropout": 0.05, - "lora_target_modules": "attention", - "n_embd": 128, - "n_head": 4, - "n_layer": 4, - "context_length": 128, - "dropout": 0.1, - "training_profile": "Stable LLM", - "epochs": 5, - "batch_size": 16, - "learning_rate": 0.0003, - "weight_decay": 0.1, - "gradient_accumulation": 1, - "warmup_steps": 100, - "eval_interval": 100, - "max_eval_batches": 50, - "save_interval": 500, - "data_loader_workers": 0, - "max_grad_norm": 1.0, - "activation_checkpointing": False, - "seed": 1337, - "device": self.device.currentText(), - "use_amp": self.use_amp_default, - "resume": True, - "require_compatible_resume": True, - "benchmark_prompts": "\n\n".join(DEFAULT_BENCHMARK_PROMPTS), - "benchmark_tokens": 128, - "benchmark_temperature": 0.7, - "benchmark_kv_cache": True, - }, - "export": { - "quantization": "FP16 checkpoint", - "gguf_outtype": "f16", - }, - "chat": { - "model_backend": "gguf", - "context": 2048, - "cpu_threads": 4, - "gpu_layers": -1, - "thinking_enabled": True, - "reasoning_effort": "Balanced", - "max_tokens": 512, - "temperature": 0.7, - "top_p": 0.9, - "repeat_penalty": 1.1, - "system_prompt": "", - }, - "distributed": { - "host": "0.0.0.0", - "port": 8765, - "artifact_root": str(Path.home() / ".drunkenbot_ide" / "artifacts"), - "public_url": "http://127.0.0.1:8765", - }, - "artifacts": {}, - } - - def _reset_project_runtime_state(self) -> None: - """Clear logs, progress, charts, and status labels for a new project.""" - - self.dataset_log.clear() - self.training_log.clear() - self.fine_tune_log.clear() - self.benchmark_log.clear() - self.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" - ) - for progress in ( - self.dataset_progress, - self.training_progress, - self.fine_tune_progress, - self.benchmark_progress, - self.export_progress, - self.chat_progress, - ): - progress.setRange(0, 100) - progress.setValue(0) - self.dataset_status.setText("Dataset: not prepared") - self.train_status.setText("Training: idle") - self.export_status.setText("Export: waiting") - self.chat_status.setText("Chat: no model loaded") - self.prepare_button.setText("Prepare Dataset") - self.train_button.setText("Start Training") - self.fine_tune_button.setText("Start Fine-Tune") - self.stop_dataset_button.setEnabled(False) - self.stop_training_button.setEnabled(False) - self.stop_fine_tune_button.setEnabled(False) - self.stop_benchmark_button.setEnabled(False) - self.stop_chat_button.setEnabled(False) - self.load_llm_button.setText("Load Model") - self._update_chat_backend_controls() - self._reset_dataset_quality_report() - self.training_epoch_metric.setText("Epoch: -") - self.training_step_metric.setText("Step: -") - self.training_loss_metric.setText("Train loss: -") - self.training_val_metric.setText("Val loss: -") - self.training_health_metric.setText("Health: -") - self.training_health_points = [] - self.training_lr_metric.setText("LR: -") - self.training_speed_metric.setText("Speed: -") - self.training_grad_metric.setText("Grad: -") - self.training_vram_metric.setText("VRAM: -") - self.training_eta_metric.setText("ETA: -") - self.model_size_metric.setText("Model: -") - self.vram_estimate_metric.setText("VRAM est: -") - self.parameter_breakdown_metric.setText("Params: -") - self.memory_breakdown_metric.setText("Memory: -") - self.architecture_advisor_metric.setText("Advisor: -") - self.history_metric.setText(f"Runs: {len(self._load_training_history())}") - self.loss_chart.clear() - self.optimization_chart.clear() - self.stability_chart.clear() - self.throughput_chart.clear() - self.memory_chart.clear() - self.live_prediction_chart.update_distribution(0, None) - self.live_attention_chart.update_heatmap(0, None) - self.live_activation_chart.update_histogram(0, None) - self.live_gradient_chart.update_flow(self.n_layer.value(), None, 0) - self.live_sample_text.setText("Training text: -") - self.telemetry_db_path = None - self.telemetry_run_id = "" - self.telemetry_latest_id = 0 - self.telemetry_latest_index = 0 - self.live_scrub_active = False - self.live_time_slider.blockSignals(True) - self.live_time_slider.setRange(0, 0) - self.live_time_slider.setValue(0) - self.live_time_slider.blockSignals(False) - self.live_timeline_label.setText("Timeline: no saved telemetry") - self._set_meter(self.live_cpu_bar, "CPU", self._system_cpu_value()) - self._set_meter(self.live_gpu_bar, "GPU memory", None) - self._set_meter(self.live_vram_bar, "VRAM reserved", None) - self._set_meter(self.live_ram_bar, "System RAM", self._system_ram_value()) - self.live_worker_status.setText(f"CPU workers: {self.data_loader_workers.value()}") - self._clear_chat_messages() - self.chat_markdown = "" - self.chat_stream_prefix = "" - self.chat_stream_reply = "" - self.chat_stats.setText("Idle") - self._add_chat_message("assistant", "Load a GGUF or MicroGPT model to start testing.") - - def _project_state_dict(self, project_name: str, project_dir: Path) -> dict[str, Any]: - """Collect all UI state that defines a Micro LLM project. - - Args: - project_name: User-facing project name. - project_dir: Folder where the project file will live. - - Returns: - JSON-serializable project state. - """ - - dataset_dir = Path(self.dataset_dir.text()) if self.dataset_dir.text().strip() else None - model_dir = Path(self.model_dir.text()) if self.model_dir.text().strip() else None - export_dir = Path(self.export_dir.text()) if self.export_dir.text().strip() else None - now_iso = datetime.now().isoformat(timespec="seconds") - created_at = now_iso - existing_project_file = project_dir / "project.json" - if existing_project_file.exists(): - try: - existing_data = json.loads(existing_project_file.read_text(encoding="utf-8")) - except Exception: - existing_data = {} - if isinstance(existing_data, dict): - # Preserve the original creation timestamp across saves. - # "saved_at" below is overwritten every save, so it cannot be - # used as a creation date; fall back to it only for projects - # saved before this field existed. - created_at = str(existing_data.get("created_at") or existing_data.get("saved_at") or now_iso) - return { - "schema": "drunkenbot_ide_project", - "version": 1, - "project_name": project_name, - "project_dir": str(project_dir), - "created_at": created_at, - "saved_at": now_iso, - "paths": { - "source_vault": self.input_dir.text(), - "dataset_core": self.dataset_dir.text(), - "training_dataset": self.train_data_dir.text(), - "model_output": self.model_dir.text(), - "export_model_core": self.export_model_dir.text(), - "export_output": self.export_dir.text(), - "llama_cpp_dir": self.llama_cpp_dir.text(), - "gguf_output_path": self.gguf_output_path.text(), - "gguf_model": self.gguf_path.text(), - "microgpt_chat_model": self.microgpt_chat_path.text(), - "tokenizer_import": self.tokenizer_path.text(), - "resume_checkpoint": self.resume_checkpoint.text(), - "fine_tune_checkpoint": self.fine_tune_checkpoint.text(), - "fine_tune_output": self.fine_tune_output_dir.text(), - }, - "dataset": { - "domain_plan_preset": self.dataset_plan_preset.currentText() if hasattr(self, "dataset_plan_preset") else "Balanced Tiny LLM", - "domain_plan": self._dataset_plan_from_ui(), - "default_data_paths": [str(path) for path in self._selected_default_data_paths()], - "external_dataset_dir": self.external_dataset_dir.text() if hasattr(self, "external_dataset_dir") else "", - "auto_vocab": self.auto_vocab.isChecked(), - "manual_vocab_size": self.manual_vocab_size.value(), - "include_conversation_datasets": self.include_conversation_datasets.isChecked(), - "dataset_stage": self._dataset_stage_value(), - "conversation_datasets": self._selected_conversation_datasets(), - "conversation_sample_limit": self.conversation_sample_limit.value(), - "mixture_weights": self._mixture_weights_from_ui(), - "min_frequency": self.min_frequency.value(), - "context_length": self.context_length.value(), - "validation_split": self.validation_split.value(), - "lowercase": False, - "max_workers": self.max_workers.value(), - "prepare_mode": self._prepare_mode_value(), - "tokenizer_strategy": self._tokenizer_strategy_value(), - "code_training_mode": self.code_training_mode.isChecked(), - "include_prose": self.include_prose.isChecked(), - "include_source_code": self.include_source_code.isChecked(), - "extract_code_blocks": self.extract_code_blocks.isChecked(), - "preserve_indentation": self.preserve_indentation.isChecked(), - "instruction_samples": self.instruction_samples.isChecked(), - "reasoning_sample_mode": self._reasoning_sample_mode_value(), - }, - "training": { - "preset": self.preset.currentText(), - "architecture_style": self.architecture_style.currentText(), - "launch_target": self._training_launch_target_value(), - "fine_tune_launch_target": self._fine_tune_launch_target_value(), - "training_stage": self._training_stage_value(), - "n_embd": self.n_embd.value(), - "n_head": self.n_head.value(), - "attention_type": self._attention_type_value(), - "kv_head_count": self.kv_head_count.value(), - "attention_backend": self._attention_backend_value(), - "attention_window": self.attention_window.value(), - "training_mode": self._training_mode_value(), - "peft_method": self._peft_method_value(), - "lora_rank": self.lora_rank.value(), - "lora_alpha": self.lora_alpha.value(), - "lora_dropout": self.lora_dropout.value(), - "lora_target_modules": self._lora_target_value(), - "n_layer": self.n_layer.value(), - "context_length": self.train_context_length.value(), - "dropout": self.dropout.value(), - "training_profile": self.training_profile.currentText(), - "epochs": self.epochs.value(), - "batch_size": self.batch_size.value(), - "learning_rate": self.learning_rate.value(), - "weight_decay": self.weight_decay.value(), - "optimizer_name": self._optimizer_value(), - "scheduler_name": self._scheduler_value(), - "scheduler_min_lr_ratio": self.min_lr_ratio.value(), - "polynomial_power": self.polynomial_power.value(), - "gradient_accumulation": self.gradient_accumulation.value(), - "sample_stride": self.sample_stride.value(), - "warmup_steps": self.warmup_steps.value(), - "eval_interval": self.eval_interval.value(), - "max_eval_batches": self.max_eval_batches.value(), - "save_interval": self.save_interval.value(), - "data_loader_workers": self.data_loader_workers.value(), - "max_grad_norm": self.max_grad_norm.value(), - "activation_checkpointing": self.activation_checkpointing.isChecked(), - "seed": self.seed.value(), - "device": self.device.currentText(), - "use_amp": self.use_amp.isChecked(), - "precision": self._precision_value(), - "resume": self.resume_training.isChecked(), - "require_compatible_resume": self.resume_safety.isChecked(), - "early_stopping": self.early_stopping.isChecked(), - "benchmark_prompts": self.benchmark_prompts.toPlainText(), - "benchmark_tokens": self.benchmark_tokens.value(), - "benchmark_temperature": self.benchmark_temperature.value(), - "benchmark_kv_cache": self.benchmark_kv_cache.isChecked(), - }, - "export": { - "quantization": self.quant_mode.currentText(), - "gguf_outtype": self.gguf_outtype.currentText(), - }, - "chat": { - "model_backend": self._chat_backend_value(), - "context": self.llama_context.value(), - "cpu_threads": self.llama_threads.value(), - "gpu_layers": self.llama_gpu_layers.value(), - "thinking_enabled": self.thinking_enabled.isChecked(), - "reasoning_effort": self.reasoning_effort.currentText(), - "max_tokens": self.chat_max_tokens.value(), - "temperature": self.chat_temperature.value(), - "top_p": self.chat_top_p.value(), - "repeat_penalty": self.chat_repeat_penalty.value(), - "system_prompt": self.system_prompt.toPlainText(), - }, - "distributed": { - "host": self.coordinator_host.text(), - "port": self.coordinator_port.value(), - "artifact_root": self.coordinator_artifact_root.text(), - "public_url": self.coordinator_public_url.text(), - }, - "artifacts": { - "dataset_summary": self._read_json_if_exists(dataset_dir / "dataset_summary.json") if dataset_dir else None, - "training_summary": self._read_json_if_exists(model_dir / "training_summary.json") if model_dir else None, - "export_summary": self._read_json_if_exists(export_dir / "export_summary.json") if export_dir else None, - }, - } - - def _apply_project_state(self, data: dict[str, Any]) -> None: - """Restore UI state from a saved project dictionary. - - Args: - data: Project state loaded from JSON. - """ - - self.search_box.setText(str(data.get("project_name", ""))) - paths = data.get("paths", {}) - dataset = data.get("dataset", {}) - training = data.get("training", {}) - export = data.get("export", {}) - chat = data.get("chat", {}) - distributed = data.get("distributed", {}) - - self.input_dir.setText(str(paths.get("source_vault", ""))) - self.dataset_dir.setText(str(paths.get("dataset_core", ""))) - self.train_data_dir.setText(str(paths.get("training_dataset", ""))) - self.model_dir.setText(str(paths.get("model_output", ""))) - self.export_model_dir.setText(str(paths.get("export_model_core", ""))) - self.export_dir.setText(str(paths.get("export_output", ""))) - self.llama_cpp_dir.setText(str(paths.get("llama_cpp_dir", ""))) - self.gguf_output_path.setText(str(paths.get("gguf_output_path", ""))) - self.gguf_path.setText(str(paths.get("gguf_model", ""))) - self.microgpt_chat_path.setText(str(paths.get("microgpt_chat_model", ""))) - self.tokenizer_path.setText(str(paths.get("tokenizer_import", ""))) - self.resume_checkpoint.setText(str(paths.get("resume_checkpoint", ""))) - self.fine_tune_checkpoint.setText(str(paths.get("fine_tune_checkpoint", ""))) - self.fine_tune_output_dir.setText(str(paths.get("fine_tune_output", ""))) - - self._set_dataset_plan( - dict(dataset.get("domain_plan", {})), - str(dataset.get("domain_plan_preset", "Balanced Tiny LLM")), - ) - saved_default_data_paths = dataset.get("default_data_paths") - self._set_selected_default_data_paths( - list(saved_default_data_paths) if saved_default_data_paths is not None else None - ) - if hasattr(self, "external_dataset_dir") and dataset.get("external_dataset_dir"): - self.external_dataset_dir.setText(str(dataset["external_dataset_dir"])) - self._refresh_external_dataset_status() - self.auto_vocab.setChecked(bool(dataset.get("auto_vocab", True))) - self.manual_vocab_size.setValue(int(dataset.get("manual_vocab_size", self.manual_vocab_size.value()))) - include_conversation = bool(dataset.get("include_conversation_datasets", False)) - self._set_dataset_stage(str(dataset.get("dataset_stage", "base"))) - self.include_conversation_datasets.setChecked(include_conversation) - self._set_selected_conversation_datasets(list(dataset.get("conversation_datasets", []))) - self.conversation_sample_limit.setValue(int(dataset.get("conversation_sample_limit", self.conversation_sample_limit.value()))) - self._set_mixture_weights(dict(dataset.get("mixture_weights", {}))) - self.min_frequency.setValue(int(dataset.get("min_frequency", self.min_frequency.value()))) - self.context_length.setValue(int(dataset.get("context_length", self.context_length.value()))) - self.validation_split.setValue(float(dataset.get("validation_split", self.validation_split.value()))) - self.max_workers.setValue(int(dataset.get("max_workers", self.max_workers.value()))) - self._set_combo_by_data(self.prepare_mode, str(dataset.get("prepare_mode", "incremental")), { - "incremental": "Incremental update", - "full_rebuild": "Full rebuild", - "force_reprocess": "Force reprocess", - }) - self._set_combo_by_data(self.tokenizer_strategy, str(dataset.get("tokenizer_strategy", "auto")), { - "auto": "Auto", - "train_new": "Train new tokenizer", - "reuse_dataset": "Reuse dataset tokenizer", - "import_tokenizer": "Import tokenizer.json", - }) - self.code_training_mode.setChecked(bool(dataset.get("code_training_mode", True))) - self.include_prose.setChecked(bool(dataset.get("include_prose", True))) - self.include_source_code.setChecked(bool(dataset.get("include_source_code", True))) - self.extract_code_blocks.setChecked(bool(dataset.get("extract_code_blocks", True))) - self.preserve_indentation.setChecked(bool(dataset.get("preserve_indentation", True))) - self.instruction_samples.setChecked(bool(dataset.get("instruction_samples", True))) - self._set_combo_by_data(self.reasoning_sample_mode, str(dataset.get("reasoning_sample_mode", "scaffold")), { - "scaffold": "Reasoning scaffold", - "detailed": "Detailed code reasoning", - "none": "No reasoning wrapper", - }) - - self._set_combo_text(self.preset, str(training.get("preset", self.preset.currentText()))) - self._set_combo_text(self.architecture_style, str(training.get("architecture_style", self.architecture_style.currentText()))) - self._set_combo_by_data(self.training_launch_target, str(training.get("launch_target", "local")), { - "local": "Local machine", - "remote": "Remote workers", - "runpod": "RunPod cloud", - }) - if hasattr(self, "fine_tune_launch_target"): - self._set_combo_by_data(self.fine_tune_launch_target, str(training.get("fine_tune_launch_target", "local")), { - "local": "Local machine", - "remote": "Remote workers", - "runpod": "RunPod cloud", - }) - self.n_embd.setValue(int(training.get("n_embd", self.n_embd.value()))) - self.n_head.setValue(int(training.get("n_head", self.n_head.value()))) - self._set_combo_by_data(self.attention_type, str(training.get("attention_type", "mha")), { - "mha": "Multi-head", - "gqa": "Grouped-query", - "mqa": "Multi-query", - }) - self.kv_head_count.setValue(int(training.get("kv_head_count", self.kv_head_count.value()))) - self._set_combo_by_data(self.attention_backend, str(training.get("attention_backend", "sdpa")), { - "sdpa": "SDPA / Flash when available", - "manual": "Manual", - }) - self.attention_window.setValue(int(training.get("attention_window", self.attention_window.value()))) - self._set_combo_by_data(self.training_mode, str(training.get("training_mode", "pretrain")), { - "pretrain": "Pretrain from scratch", - "fine_tune": "Fine-tune checkpoint", - "instruction_fine_tune": "Instruction fine-tune", - "conversation_fine_tune": "Conversation fine-tune", - "code_fine_tune": "Code fine-tune", - }) - training_stage = str(training.get("training_stage", "")) - if training_stage == "instruction": - self._set_combo_text(self.training_mode, "Instruction fine-tune") - elif training_stage == "conversation": - self._set_combo_text(self.training_mode, "Conversation fine-tune") - elif training_stage == "code": - self._set_combo_text(self.training_mode, "Code fine-tune") - self._set_combo_by_data(self.peft_method, str(training.get("peft_method", "none")), { - "none": "Full fine-tune", - "lora": "LoRA adapters", - }) - self.lora_rank.setValue(int(training.get("lora_rank", self.lora_rank.value()))) - self.lora_alpha.setValue(float(training.get("lora_alpha", self.lora_alpha.value()))) - self.lora_dropout.setValue(float(training.get("lora_dropout", self.lora_dropout.value()))) - self._set_combo_by_data(self.lora_targets, str(training.get("lora_target_modules", "attention")), { - "attention": "Attention projections", - "mlp": "MLP projections", - "attention,mlp": "Attention + MLP", - }) - self.n_layer.setValue(int(training.get("n_layer", self.n_layer.value()))) - self.train_context_length.setValue(int(training.get("context_length", self.train_context_length.value()))) - self.dropout.setValue(float(training.get("dropout", self.dropout.value()))) - self._set_combo_text(self.training_profile, str(training.get("training_profile", self.training_profile.currentText()))) - self.epochs.setValue(int(training.get("epochs", self.epochs.value()))) - self.batch_size.setValue(int(training.get("batch_size", self.batch_size.value()))) - self.learning_rate.setValue(float(training.get("learning_rate", self.learning_rate.value()))) - self.weight_decay.setValue(float(training.get("weight_decay", self.weight_decay.value()))) - self._set_combo_by_data(self.optimizer_name, str(training.get("optimizer_name", "adamw")), { - "adamw": "AdamW", - "adam": "Adam", - "lion": "Lion", - "adafactor": "Adafactor", - }) - self._set_combo_by_data(self.scheduler_name, str(training.get("scheduler_name", "warmup_linear")), { - "warmup_linear": "Warmup linear", - "cosine": "Cosine decay", - "polynomial": "Polynomial decay", - "one_cycle": "One-cycle", - "constant": "Constant", - }) - self.min_lr_ratio.setValue(float(training.get("scheduler_min_lr_ratio", self.min_lr_ratio.value()))) - self.polynomial_power.setValue(float(training.get("polynomial_power", self.polynomial_power.value()))) - self.gradient_accumulation.setValue(int(training.get("gradient_accumulation", self.gradient_accumulation.value()))) - self.sample_stride.setValue(int(training.get("sample_stride", self.sample_stride.value()))) - self.warmup_steps.setValue(int(training.get("warmup_steps", self.warmup_steps.value()))) - self.eval_interval.setValue(int(training.get("eval_interval", self.eval_interval.value()))) - self.max_eval_batches.setValue(int(training.get("max_eval_batches", self.max_eval_batches.value()))) - self.save_interval.setValue(int(training.get("save_interval", self.save_interval.value()))) - self.data_loader_workers.setValue(int(training.get("data_loader_workers", self.data_loader_workers.value()))) - self.max_grad_norm.setValue(float(training.get("max_grad_norm", self.max_grad_norm.value()))) - self.activation_checkpointing.setChecked(bool(training.get("activation_checkpointing", False))) - self.seed.setValue(int(training.get("seed", self.seed.value()))) - self._set_combo_text(self.device, str(training.get("device", self.device.currentText()))) - self.use_amp.setChecked(bool(training.get("use_amp", self.use_amp.isChecked()))) - self._set_combo_by_data(self.precision, str(training.get("precision", "fp16")), { - "fp16": "FP16", - "bf16": "BF16", - "fp32": "FP32", - }) - self.resume_training.setChecked(bool(training.get("resume", self.resume_training.isChecked()))) - self.resume_safety.setChecked(bool(training.get("require_compatible_resume", True))) - self.early_stopping.setChecked(bool(training.get("early_stopping", True))) - self.benchmark_prompts.setPlainText(str(training.get("benchmark_prompts", self.benchmark_prompts.toPlainText()))) - self.benchmark_tokens.setValue(int(training.get("benchmark_tokens", self.benchmark_tokens.value()))) - self.benchmark_temperature.setValue(float(training.get("benchmark_temperature", self.benchmark_temperature.value()))) - self.benchmark_kv_cache.setChecked(bool(training.get("benchmark_kv_cache", True))) - - self._set_combo_text(self.quant_mode, str(export.get("quantization", self.quant_mode.currentText()))) - self._set_combo_text(self.gguf_outtype, str(export.get("gguf_outtype", self.gguf_outtype.currentText()))) - self.llama_context.setValue(int(chat.get("context", self.llama_context.value()))) - self._set_combo_by_data(self.chat_model_backend, str(chat.get("model_backend", "gguf")), { - "gguf": "GGUF / llama.cpp", - "microgpt": "MicroGPT checkpoint", - }) - self.llama_threads.setValue(int(chat.get("cpu_threads", self.llama_threads.value()))) - self.llama_gpu_layers.setValue(int(chat.get("gpu_layers", self.llama_gpu_layers.value()))) - self.thinking_enabled.setChecked(bool(chat.get("thinking_enabled", True))) - self._set_combo_text(self.reasoning_effort, str(chat.get("reasoning_effort", self.reasoning_effort.currentText()))) - self.reasoning_effort.setEnabled(self.thinking_enabled.isChecked()) - self.chat_max_tokens.setValue(int(chat.get("max_tokens", self.chat_max_tokens.value()))) - self.chat_temperature.setValue(float(chat.get("temperature", self.chat_temperature.value()))) - self.chat_top_p.setValue(float(chat.get("top_p", self.chat_top_p.value()))) - self.chat_repeat_penalty.setValue(float(chat.get("repeat_penalty", self.chat_repeat_penalty.value()))) - self.system_prompt.setPlainText(str(chat.get("system_prompt", ""))) - if hasattr(self, "coordinator_host"): - self.coordinator_host.setText(str(distributed.get("host", self.coordinator_host.text()))) - self.coordinator_port.setValue(int(distributed.get("port", self.coordinator_port.value()))) - self.coordinator_artifact_root.setText(str(distributed.get("artifact_root", self.coordinator_artifact_root.text()))) - self.coordinator_public_url.setText(str(distributed.get("public_url", self.coordinator_public_url.text()))) - self._update_tokenizer_strategy_controls() - self._update_training_mode_controls() - self._restore_artifact_status(data.get("artifacts", {})) - self.refresh_fine_tune_workflow() - - def _restore_artifact_status(self, artifacts: dict[str, Any]) -> None: - """Refresh top-bar and button state from saved or existing artifacts. - - Args: - artifacts: Saved artifact summary dictionary. - """ - - dataset_dir = Path(self.dataset_dir.text()) if self.dataset_dir.text().strip() else None - if dataset_dir and self._dataset_artifacts_exist(dataset_dir): - summary = self._read_json_if_exists(dataset_dir / "dataset_summary.json") or artifacts.get("dataset_summary") or {} - document_count = int(summary.get("document_count", 0) or 0) - token_count = int(summary.get("token_count", 0) or 0) - code_count = int(summary.get("code_sample_count", 0) or 0) - prose_count = int(summary.get("prose_sample_count", 0) or 0) - conversation_count = int(summary.get("conversation_sample_count", 0) or 0) - vocab_size = int(summary.get("tokenizer_vocab_size", 0) or 0) - self._update_dataset_quality_report(summary) - self.prepare_button.setText("DataSet Prepared") - self.dataset_progress.setValue(100) - if vocab_size: - self.auto_vocab_label.setText(f"{vocab_size:,}") - if code_count or prose_count or conversation_count: - self.dataset_status.setText( - f"Dataset: {code_count:,} code, {prose_count:,} prose, {conversation_count:,} chat, {token_count:,} tokens" - ) - elif document_count or token_count: - self.dataset_status.setText(f"Dataset: {document_count:,} files, {token_count:,} tokens") - else: - self.dataset_status.setText("Dataset: prepared") - version = summary.get("dataset_version", {}) - if isinstance(version, dict) and version.get("version_id"): - self.dataset_log.append(f"Dataset version: {version['version_id']}") - self.train_data_dir.setText(str(dataset_dir)) - self.dataset_log.append(f"Dataset already prepared: {dataset_dir}") - else: - self.prepare_button.setText("Prepare Dataset") - self.dataset_progress.setValue(0) - self.dataset_status.setText("Dataset: not prepared") - self.auto_vocab_label.setText("Auto after reading files") - self._reset_dataset_quality_report() - - model_dir = Path(self.model_dir.text()) if self.model_dir.text().strip() else None - if model_dir and (model_dir / "final_model.pt").exists(): - summary = self._read_json_if_exists(model_dir / "training_summary.json") or artifacts.get("training_summary") or {} - loss = summary.get("final_train_loss") - self.train_status.setText(f"Training: loss {float(loss):.4f}" if loss is not None else "Training: model ready") - self.export_model_dir.setText(str(model_dir)) - - export_dir = Path(self.export_dir.text()) if self.export_dir.text().strip() else None - if export_dir and export_dir.exists() and any(export_dir.iterdir()): - self.export_status.setText("Export: artifacts found") - - @staticmethod - def _dataset_artifacts_exist(dataset_dir: Path) -> bool: - """Return whether a dataset folder has the required prepared files. - - Args: - dataset_dir: Dataset folder. - - Returns: - True if required dataset artifacts exist. - """ - - if not dataset_dir.exists(): - return False - if not (dataset_dir / "tokenizer.json").exists(): - return False - has_npy_tokens = (dataset_dir / "train_tokens.npy").exists() and (dataset_dir / "val_tokens.npy").exists() - has_json_tokens = (dataset_dir / "train_tokens.json").exists() and (dataset_dir / "val_tokens.json").exists() - return has_npy_tokens or has_json_tokens - - @staticmethod - def _safe_project_name(project_name: str) -> str: - """Return a filesystem-safe project folder name. - - Args: - project_name: Raw user project name. - - Returns: - Safe folder name. - """ - - return re.sub(r"[^A-Za-z0-9_.-]+", "_", project_name).strip("._") or "MicroLLMProject" - - @staticmethod - def _read_json_if_exists(path: Path) -> Optional[Any]: - """Read a JSON file when it exists. - - Args: - path: JSON file path. - - Returns: - Parsed JSON or ``None``. - """ - - if not path.exists(): - return None - try: - return json.loads(path.read_text(encoding="utf-8")) - except Exception: - return None - - @staticmethod - def _set_combo_text(combo: QComboBox, text: str) -> None: - """Set combo text when the value exists. - - Args: - combo: Combo box to update. - text: Display text to select. - """ - - index = combo.findText(text) - if index >= 0: - combo.setCurrentIndex(index) - elif combo.isEditable(): - combo.setEditText(text) - - def _set_combo_by_data(self, combo: QComboBox, value: str, labels: dict[str, str]) -> None: - """Set a combo by internal saved value. - - Args: - combo: Combo box to update. - value: Internal saved value. - labels: Mapping from saved value to display label. - """ - - self._set_combo_text(combo, labels.get(value, value)) - - def _run_task( - self, - fn, - args, - on_finished, - log: QTextEdit, - progress_bar: QProgressBar, - with_progress: bool = False, - button: Optional[QPushButton] = None, - stop_button: Optional[QPushButton] = None, - busy_text: str = "Working", - task_kind: str = "", - isolate_process: bool = False, - ) -> None: - """Run a long task on a background thread. - - Args: - fn: Callable to execute. - args: Positional arguments for the callable. - on_finished: Slot called with the task result. - log: Log widget receiving progress messages. - progress_bar: Progress bar receiving percent updates. - with_progress: Whether to pass a progress callback to the task. - button: Optional button to disable while running. - stop_button: Optional stop button to enable while running. - busy_text: Button text shown while running. - task_kind: Optional notification stage key. - isolate_process: Run the task inside a child process. - """ - - if self.thread is not None: - QMessageBox.information(self, "Task running", "Please wait for the current task to finish.") - return - - LOGGER.info("Starting background task: %s", getattr(fn, "__name__", str(fn))) - self.active_task_kind = task_kind - if button: - self._set_button_busy(button, busy_text) - if stop_button: - stop_button.setEnabled(True) - self.active_stop_button = stop_button - - self.stop_event = Event() - self.progress_queue = Queue() - self.active_log = log - self.active_progress_bar = progress_bar - self.thread = QThread(self) - worker_class = ProcessTaskWorker if isolate_process else TaskWorker - self.worker = worker_class( - fn, - *args, - progress_queue=self.progress_queue, - with_progress=with_progress, - stop_event=self.stop_event, - ) - self.worker.moveToThread(self.thread) - self.thread.started.connect(self.worker.run) - self.worker.finished.connect(on_finished) - self.worker.finished.connect(self.worker.deleteLater) - self.worker.finished.connect(self.thread.quit) - self.worker.failed.connect(self._task_failed_from_worker) - self.worker.failed.connect(self.worker.deleteLater) - self.worker.failed.connect(self.thread.quit) - self.thread.finished.connect(self.thread.deleteLater) - self.thread.finished.connect(self._thread_finished) - self.progress_timer.start(100) - self.thread.start() - - @Slot(str) - def _task_failed_from_worker(self, message: str) -> None: - """Handle a worker failure on the UI thread. - - Args: - message: Error message emitted by the worker. - """ - - if self.active_log is None or self.active_progress_bar is None: - return - LOGGER.error("Background task failed: %s", message) - if self.active_task_kind == "chat": - self.chat_status.setText(f"Chat: load failed - {message}") - elif self.active_task_kind == "dataset_download": - self.external_dataset_version.setText(f"Download failed: {message}") - self.dataset_plan_progress.setVisible(False) - self._task_failed(message, self.active_log, self.active_progress_bar) - - def stop_active_task(self) -> None: - """Request a graceful stop for the active background task.""" - - if self.stop_event is None: - return - LOGGER.info("Stop requested for active background task") - self.stop_event.set() - self._notify_failure("Stop requested", "The task is stopping at the next safe point.") - if self.active_log is not None: - self.active_log.append("Stop requested. Finishing the current safe point...") - if self.active_stop_button is not None: - self.active_stop_button.setEnabled(False) - if torch.cuda.is_available(): - try: - torch.cuda.empty_cache() - except Exception: - LOGGER.exception( - "Failed to empty CUDA cache in _thread_finished") - - @Slot() - def request_shutdown_from_signal(self) -> None: - """Handle Ctrl+C from a terminal without leaving Qt threads wedged.""" - - self.interrupt_count += 1 - if self.interrupt_count > 1: - os._exit(130) - if self.stop_event is not None: - self.stop_event.set() - if self.active_log is not None: - self.active_log.append("Interrupt received. Requesting stop...") - self.project_state.setText("Stopping") - if self.thread is None: - QApplication.quit() - return - QTimer.singleShot(3000, lambda: os._exit(130) if self.thread is not None else QApplication.quit()) - - def closeEvent(self, event: Any) -> None: - """Clean up background services before the window closes. - - Args: - event: Qt close event. - """ - - if self.thread is not None: - if self.stop_event is not None: - self.stop_event.set() - if self.active_log is not None: - self.active_log.append("Close requested. Stopping active task first...") - self.project_state.setText("Stopping") - LOGGER.info("Close requested while background task is running; waiting for task shutdown") - event.ignore() - QTimer.singleShot(500, self.close) - return - if self.coordinator_server is not None: - self.stop_coordinator_server() - super().closeEvent(event) - - def _handle_progress(self, event: object, log: QTextEdit, progress_bar: QProgressBar) -> None: - """Apply one progress event to UI widgets. - - Args: - event: Progress dictionary or message. - log: Log widget to append messages to. - progress_bar: Progress bar to update. - """ - - if isinstance(event, dict): - if event.get("type") == "chat_delta": - self._apply_chat_delta(event) - return - message = event.get("message") - percent = event.get("percent") - if log in (self.training_log, getattr(self, "fine_tune_log", None)): - self._update_training_metrics(event, update_fine_tune=log is getattr(self, "fine_tune_log", None)) - if message: - log.append(str(message)) - if log in (self.training_log, getattr(self, "fine_tune_log", None)) and hasattr(self, "live_log"): - self.live_log.append(str(message)) - if percent is not None: - progress_bar.setValue(max(0, min(100, int(percent)))) - if log in (self.training_log, getattr(self, "fine_tune_log", None)) and hasattr(self, "live_progress"): - self.live_progress.setValue(max(0, min(100, int(percent)))) - else: - log.append(str(event)) - - def _notify_progress(self, event: dict[str, Any]) -> None: - """Send throttled external progress notifications for long tasks. - - Args: - event: Progress event emitted by a worker. - """ - - if not self.active_task_kind or self.notification_manager is None: - return - if self.active_task_kind not in {"dataset", "training", "fine_tune"}: - return - title = { - "dataset": "Dataset preparation", - "training": "Model training", - "fine_tune": "Fine-tuning", - }[self.active_task_kind] - percent = event.get("percent") - self.notification_manager.notify_progress( - self.active_task_kind, - title, - self._notification_lines_from_event(event), - int(percent) if percent is not None else None, - ) - - def _notify_complete(self, stage_key: str, title: str, lines: list[str]) -> None: - """Send an external completion notification when configured. - - Args: - stage_key: Notification stage key. - title: User-facing title. - lines: Plain-text summary lines. - """ - - if self.notification_manager is not None: - self.notification_manager.notify_complete(stage_key, title, lines) - - def _notify_failure(self, title: str, message: str) -> None: - """Send an external failure or stop notification for the active task. - - Args: - title: User-facing title. - message: Failure details. - """ - - if self.active_task_kind and self.notification_manager is not None: - self.notification_manager.notify_failure(self.active_task_kind, title, message) - - def _notification_lines_from_event(self, event: dict[str, Any]) -> list[str]: - """Build compact notification text from a worker progress event. - - Args: - event: Progress event emitted by a worker. - - Returns: - Body lines for the notification message. - """ - - lines: list[str] = [] - if event.get("message"): - lines.append(str(event["message"])) - if "epoch" in event and "total_epochs" in event: - lines.append(f"Epoch: {event['epoch']}/{event['total_epochs']}") - if "step" in event and "total_steps" in event: - lines.append(f"Step: {event['step']}/{event['total_steps']}") - train_loss = self._finite_metric(event.get("train_loss")) - if train_loss is not None: - lines.append(f"Train loss: {float(train_loss):.4f}") - val_loss = self._finite_metric(event.get("val_loss")) - if val_loss is not None: - lines.append(f"Validation loss: {float(val_loss):.4f}") - learning_rate = self._finite_metric(event.get("learning_rate")) - if learning_rate is not None: - lines.append(f"Learning rate: {float(learning_rate):.2e}") - tokens_per_second = self._finite_metric(event.get("tokens_per_second")) - if tokens_per_second is not None: - lines.append(f"Speed: {float(tokens_per_second):.0f} tokens/sec") - eta_seconds = self._finite_metric(event.get("eta_seconds")) - if eta_seconds is not None: - lines.append(f"ETA: {self._format_duration(float(eta_seconds))}") - vram_allocated = self._finite_metric(event.get("vram_allocated_gb")) - vram_reserved = self._finite_metric(event.get("vram_reserved_gb")) - if vram_allocated is not None or vram_reserved is not None: - allocated = "-" if vram_allocated is None else f"{float(vram_allocated):.2f} GB" - reserved = "-" if vram_reserved is None else f"{float(vram_reserved):.2f} GB" - lines.append(f"VRAM: {allocated} allocated, {reserved} reserved") - return lines[:10] - - def _update_training_metrics(self, event: dict[str, Any], update_fine_tune: bool = False) -> None: - """Update training metric chips from a progress event. - - Args: - event: Progress event emitted by the training backend. - update_fine_tune: Whether to mirror metrics into the Fine-Tuning tab chips. - """ - - if "epoch" in event and "total_epochs" in event: - self.training_epoch_metric.setText(f"Epoch: {event['epoch']}/{event['total_epochs']}") - if update_fine_tune and hasattr(self, "fine_tune_epoch_metric"): - self.fine_tune_epoch_metric.setText(f"Epoch: {event['epoch']}/{event['total_epochs']}") - if "step" in event and "total_steps" in event: - self.training_step_metric.setText(f"Step: {event['step']}/{event['total_steps']}") - if update_fine_tune and hasattr(self, "fine_tune_step_metric"): - self.fine_tune_step_metric.setText(f"Step: {event['step']}/{event['total_steps']}") - train_loss = self._finite_metric(event.get("train_loss")) - if train_loss is not None: - self.training_loss_metric.setText(f"Train loss: {float(train_loss):.4f}") - if update_fine_tune and hasattr(self, "fine_tune_loss_metric"): - self.fine_tune_loss_metric.setText(f"Train loss: {float(train_loss):.4f}") - val_loss = self._finite_metric(event.get("val_loss")) - if val_loss is not None: - self.training_val_metric.setText(f"Val loss: {float(val_loss):.4f}") - if update_fine_tune and hasattr(self, "fine_tune_val_metric"): - self.fine_tune_val_metric.setText(f"Val loss: {float(val_loss):.4f}") - step = event.get("step") - if step is not None and (train_loss is not None or val_loss is not None): - step_int_for_loss = int(step) - self.loss_chart.add_metrics(step_int_for_loss, train_loss, val_loss) - self._update_training_health(step_int_for_loss, train_loss, val_loss) - if step is None: - return - step_int = int(step) - self._record_live_metric(event) - learning_rate = self._finite_metric(event.get("learning_rate")) - grad_norm = self._finite_metric(event.get("grad_norm")) - weight_norm = self._finite_metric(event.get("weight_norm")) - update_ratio = self._finite_metric(event.get("update_ratio")) - tokens_per_second = self._finite_metric(event.get("tokens_per_second")) - samples_per_second = self._finite_metric(event.get("samples_per_second")) - vram_allocated = self._finite_metric(event.get("vram_allocated_gb")) - vram_reserved = self._finite_metric(event.get("vram_reserved_gb")) - gpu_memory = self._finite_metric(event.get("gpu_memory_percent")) - system_cpu = self._finite_metric(event.get("system_cpu_percent")) - system_ram = self._finite_metric(event.get("system_ram_percent")) - data_workers = event.get("data_loader_workers") - eta_seconds = self._finite_metric(event.get("eta_seconds")) - if learning_rate is not None: - self.training_lr_metric.setText(f"LR: {float(learning_rate):.2e}") - if update_fine_tune and hasattr(self, "fine_tune_lr_metric"): - self.fine_tune_lr_metric.setText(f"LR: {float(learning_rate):.2e}") - if grad_norm is not None: - self.training_grad_metric.setText(f"Grad: {float(grad_norm):.3f}") - if update_fine_tune and hasattr(self, "fine_tune_grad_metric"): - self.fine_tune_grad_metric.setText(f"Grad: {float(grad_norm):.3f}") - if tokens_per_second is not None: - self.training_speed_metric.setText(f"Speed: {float(tokens_per_second):.0f} tok/s") - if update_fine_tune and hasattr(self, "fine_tune_speed_metric"): - self.fine_tune_speed_metric.setText(f"Speed: {float(tokens_per_second):.0f} tok/s") - if vram_allocated is not None: - self.training_vram_metric.setText(f"VRAM: {float(vram_allocated):.2f} GB") - if eta_seconds is not None: - self.training_eta_metric.setText(f"ETA: {self._format_duration(float(eta_seconds))}") - if update_fine_tune and hasattr(self, "fine_tune_eta_metric"): - self.fine_tune_eta_metric.setText(f"ETA: {self._format_duration(float(eta_seconds))}") - if learning_rate is not None or grad_norm is not None: - self.optimization_chart.add_values(step_int, learning_rate, grad_norm) - if weight_norm is not None or update_ratio is not None: - self.stability_chart.add_values(step_int, weight_norm, update_ratio) - if tokens_per_second is not None or samples_per_second is not None: - self.throughput_chart.add_values(step_int, tokens_per_second, samples_per_second) - if vram_allocated is not None or vram_reserved is not None: - self.memory_chart.add_values(step_int, vram_allocated, vram_reserved) - if hasattr(self, "live_epoch_metric"): - self._update_live_training_metrics( - step_int, - event, - train_loss, - learning_rate, - grad_norm, - update_ratio, - tokens_per_second, - samples_per_second, - vram_allocated, - vram_reserved, - gpu_memory, - system_cpu, - system_ram, - data_workers, - ) - - def _update_training_health( - self, - step: int, - train_loss: Optional[float], - val_loss: Optional[float], - ) -> None: - """Update the training health advisor from recent loss values. - - Args: - step: Current optimizer step. - train_loss: Latest training loss. - val_loss: Latest validation loss. - """ - - self.training_health_points.append((step, train_loss, val_loss)) - self.training_health_points = self.training_health_points[-12:] - latest_train = next((item[1] for item in reversed(self.training_health_points) if item[1] is not None), None) - latest_val = next((item[2] for item in reversed(self.training_health_points) if item[2] is not None), None) - val_points = [(item[0], item[2]) for item in self.training_health_points if item[2] is not None] - if latest_train is None and latest_val is None: - label = "Health: collecting" - tip = "Waiting for train and validation loss." - elif latest_train is not None and latest_val is not None and latest_train < 0.2 and latest_val > max(2.0, latest_train * 8.0): - label = "Health: validation gap" - tip = "Training loss is very low while validation loss is high. Check overfitting, validation split, tokenizer match, or eval settings." - elif len(val_points) >= 3 and val_points[-1][1] > val_points[-2][1] > val_points[-3][1]: - label = "Health: overfitting?" - tip = "Validation loss has increased for three checks. Consider stopping, reducing epochs, or improving validation data." - elif latest_train is not None and (latest_train > 20.0 or not math.isfinite(latest_train)): - label = "Health: diverging" - tip = "Training loss is unstable or extremely high. Lower learning rate and check gradients/data." - elif latest_val is not None and latest_val > 10.0: - label = "Health: high val loss" - tip = "Validation loss is high. This may be early training, a difficult validation split, or a dataset/tokenizer mismatch." - elif latest_train is not None and latest_val is not None and latest_val <= latest_train * 1.8: - label = "Health: stable" - tip = "Training and validation loss are reasonably close." - else: - label = "Health: watching" - tip = "Collecting more loss points before making a stronger diagnosis." - self.training_health_metric.setText(label) - self._tip(self.training_health_metric, tip) - - @staticmethod - def _finite_metric(value: Any) -> Optional[float]: - """Return a finite metric value or ``None``. - - Args: - value: Raw metric value. - - Returns: - Finite float, or ``None`` when invalid. - """ - - if value is None: - return None - try: - numeric = float(value) - except (TypeError, ValueError): - return None - return numeric if math.isfinite(numeric) else None - - def _update_live_training_metrics( - self, - step: int, - event: dict[str, Any], - train_loss: Optional[float], - learning_rate: Optional[float], - grad_norm: Optional[float], - update_ratio: Optional[float], - tokens_per_second: Optional[float], - samples_per_second: Optional[float], - vram_allocated: Optional[float], - vram_reserved: Optional[float], - gpu_memory: Optional[float], - system_cpu: Optional[float], - system_ram: Optional[float], - data_workers: Optional[int], - ) -> None: - """Update live tracker widgets from one training progress event. - - Args: - step: Current optimizer step. - event: Progress event emitted by training. - train_loss: Latest training loss. - learning_rate: Current learning rate. - grad_norm: Current gradient norm. - update_ratio: Current parameter update ratio. - tokens_per_second: Current token throughput. - samples_per_second: Current sample throughput. - vram_allocated: Current CUDA allocated memory in GB. - vram_reserved: Current CUDA reserved memory in GB. - gpu_memory: Current GPU memory pressure percentage. - system_cpu: Current system CPU utilization percentage. - system_ram: Current system RAM utilization percentage. - data_workers: CPU data-loader worker count. - """ - - total_steps = event.get("total_steps") - if "epoch" in event and "total_epochs" in event: - self.live_epoch_metric.setText(f"Epoch: {event['epoch']}/{event['total_epochs']}") - if total_steps: - self.live_step_metric.setText(f"Step: {step:,}/{int(total_steps):,}") - data_percent = min(100.0, max(0.0, (step / max(1, int(total_steps))) * 100.0)) - self.live_data_metric.setText(f"Data: {data_percent:.1f}%") - self.live_progress.setValue(int(data_percent)) - else: - self.live_step_metric.setText(f"Step: {step:,}") - if tokens_per_second is not None: - self.live_tokens_metric.setText(f"Tokens/sec: {float(tokens_per_second):,.0f}") - if train_loss is not None: - self.live_loss_metric.setText(f"Loss: {float(train_loss):.4f}") - if learning_rate is not None: - self.live_lr_metric.setText(f"LR: {float(learning_rate):.2e}") - sample_text = str(event.get("sample_text") or "").strip() - if sample_text: - self.live_sample_text.setText(f"Training text: {self._compact_preview_text(sample_text, 220)}") - self.live_layer_status.setText(f"▣ Layers: {self.n_layer.value()}") - self.live_head_status.setText(f"◎ Heads: {self.n_head.value()}") - self.live_hidden_status.setText(f"▤ Hidden size: {self.n_embd.value()}") - self.live_batch_status.setText(f"▥ Batch size: {self.batch_size.value()}") - self.live_context_status.setText(f"▢ Context: {self.train_context_length.value()}") - self.live_device_status.setText(f"Device: {self.device.currentText()}") - self.live_worker_status.setText(f"CPU workers: {data_workers if data_workers is not None else self.data_loader_workers.value()}") - self._set_meter(self.live_cpu_bar, "CPU", system_cpu if system_cpu is not None else self._system_cpu_value()) - self._set_meter(self.live_gpu_bar, "GPU memory", gpu_memory) - if vram_allocated is not None or vram_reserved is not None: - allocated = float(vram_allocated or 0.0) - reserved = float(vram_reserved or 0.0) - reserved_percent = None - if self.device.currentText().startswith("cuda") and torch.cuda.is_available(): - try: - _, total_vram = torch.cuda.mem_get_info() - reserved_percent = min(100.0, 100.0 * reserved * (1024 ** 3) / max(total_vram, 1)) - except Exception: - reserved_percent = None - self._set_meter(self.live_vram_bar, "VRAM reserved", reserved_percent) - self.live_vram_label.setText(f"VRAM reserved: {reserved:.2f} GB ({allocated:.2f} GB active)") - self._set_meter(self.live_ram_bar, "System RAM", system_ram if system_ram is not None else self._system_ram_value()) - latest_loss = float(train_loss) if train_loss is not None else None - self.live_flow.set_state(self.n_layer.value(), self.n_head.value(), step, latest_loss) - self.live_prediction_chart.update_distribution(step, latest_loss) - self.live_attention_chart.update_heatmap(step, grad_norm) - self.live_activation_chart.update_histogram(step, tokens_per_second) - self.live_gradient_chart.update_flow(self.n_layer.value(), grad_norm, step) - - def _system_ram_value(self) -> Optional[float]: - """Read system RAM utilization for live telemetry. - - Returns: - System RAM percentage, or None when unavailable. - """ - - if psutil is None: - return None - return float(psutil.virtual_memory().percent) - - def _system_cpu_value(self) -> Optional[float]: - """Read system CPU utilization for live telemetry. - - Returns: - System CPU percentage, or None when unavailable. - """ - - if psutil is None: - return None - return float(psutil.cpu_percent(interval=None)) - - @staticmethod - def _format_duration(seconds: float) -> str: - """Format a duration for compact UI display. - - Args: - seconds: Duration in seconds. - - Returns: - Human-readable compact duration. - """ - - seconds = max(0, int(seconds)) - hours, remainder = divmod(seconds, 3600) - minutes, secs = divmod(remainder, 60) - if hours: - return f"{hours}h {minutes:02d}m" - if minutes: - return f"{minutes}m {secs:02d}s" - return f"{secs}s" - - def _apply_chat_delta(self, event: dict[str, Any]) -> None: - """Apply one streamed chat chunk to the rendered conversation. - - Args: - event: Chat stream progress event. - """ - - self.chat_stream_reply += str(event.get("content", "")) - should_follow = self._is_chat_near_bottom() - self._render_chat_markdown(self.chat_stream_reply) - if should_follow: - self.chat_scroll.verticalScrollBar().setValue(self.chat_scroll.verticalScrollBar().maximum()) - self._set_chat_stats( - float(event.get("elapsed_seconds", 0.0)), - int(event.get("token_count", 0)), - float(event.get("tokens_per_second", 0.0)), - ) - - def _drain_progress_queue(self) -> None: - """Drain queued worker progress events on the UI thread.""" - - if self.progress_queue is None or self.active_log is None or self.active_progress_bar is None: - return - drained = 0 - last_percent = None - while drained < 12: - try: - event = self.progress_queue.get_nowait() - except Empty: - break - notification_event = event - if isinstance(event, dict) and event.get("percent") is not None: - last_percent = event.get("percent") - event = {**event, "percent": None} - if ( - isinstance(notification_event, dict) - and self.active_task_kind == "dataset_download" - and notification_event.get("button_text") - and self.active_button is not None - ): - self.active_button_text = str(notification_event["button_text"]) - self.active_button.setText(self.active_button_text) - self._handle_progress(event, self.active_log, self.active_progress_bar) - if isinstance(notification_event, dict): - self._notify_progress(notification_event) - drained += 1 - if last_percent is not None: - self.active_progress_bar.setValue(max(0, min(100, int(last_percent)))) - - def _thread_finished(self) -> None: - """Clean up thread bookkeeping after a worker finishes.""" - - LOGGER.info("Background task thread finished") - self._drain_progress_queue() - if self.progress_timer.isActive(): - self.progress_timer.stop() - self.thread = None - self.worker = None - self.stop_event = None - self.progress_queue = None - self.active_log = None - self.active_progress_bar = None - if self.active_stop_button is not None: - self.active_stop_button.setEnabled(False) - self.active_stop_button = None - if self.active_button is not None: - self._clear_button_busy() - self.active_task_kind = "" - - def _task_failed(self, message: str, log: QTextEdit, progress_bar: QProgressBar) -> None: - """Handle background task failure. - - Args: - message: Error message. - log: Log widget to append to. - progress_bar: Progress bar to reset. - """ - - stopped_by_user = "stopped by user" in message.lower() - if stopped_by_user: - LOGGER.info("Background task stopped by user: %s", message) - else: - LOGGER.error("Background task error: %s", message) - log.append(f"Stopped: {message}" if stopped_by_user else f"Error: {message}") - self._notify_failure("Task stopped" if stopped_by_user else "Task failed", message) - progress_bar.setRange(0, 100) - progress_bar.setValue(0) - if stopped_by_user: - self.project_state.setText("Stopped") - self._clear_button_busy() - - def _set_button_busy(self, button: QPushButton, text: str) -> None: - """Disable a button and start its spinner text. - - Args: - button: Button to mark busy. - text: Busy label. - """ - - self.active_button = button - self.active_button_text = text - self.active_button_restore_text = button.text() - self.spinner_index = 0 - button.setEnabled(False) - button.setText(f"| {text}") - self.spinner_timer.start(150) - - def _clear_button_busy(self, final_text: Optional[str] = None) -> None: - """Restore the active busy button. - - Args: - final_text: Optional final button text. - """ - - if self.spinner_timer.isActive(): - self.spinner_timer.stop() - if self.active_button: - self.active_button.setEnabled(True) - self.active_button.setText(final_text or self.active_button_restore_text) - if self.active_stop_button: - self.active_stop_button.setEnabled(False) - self.active_button = None - self.active_button_text = "" - self.active_button_restore_text = "" - - def _tick_spinner(self) -> None: - """Advance the active button spinner frame.""" - - if not self.active_button: - return - frames = "|/-\\" - self.spinner_index = (self.spinner_index + 1) % len(frames) - self.active_button.setText(f"{frames[self.spinner_index]} {self.active_button_text}") - - def _dataset_config_from_ui(self) -> DatasetConfig: - """Collect dataset options from the current UI controls. - - Returns: - Dataset preparation configuration. - """ - - conversation_paths: list[Path] = [] - instruction_paths: list[Path] = [] - dataset_stage = self._dataset_stage_value() - return DatasetConfig( - input_dir=Path(self.input_dir.text()), - output_dir=Path(self.dataset_dir.text()), - vocab_size=None if self.auto_vocab.isChecked() else self.manual_vocab_size.value(), - conversation_datasets=self._selected_conversation_datasets(), - conversation_sample_limit=self.conversation_sample_limit.value(), - conversation_dataset_path=conversation_paths[0] if conversation_paths else None, - instruction_dataset_path=instruction_paths[0] if instruction_paths else None, - conversation_dataset_paths=conversation_paths, - instruction_dataset_paths=instruction_paths, - default_data_paths=self._selected_default_data_paths_for_stage(dataset_stage), - mixture_weights=self._mixture_weights_from_ui(), - min_frequency=self.min_frequency.value(), - context_length=self.context_length.value(), - validation_split=self.validation_split.value(), - lowercase=False, - max_workers=self.max_workers.value(), - code_training_mode=self.code_training_mode.isChecked(), - include_prose=self.include_prose.isChecked(), - include_source_code=self.include_source_code.isChecked(), - extract_code_blocks=self.extract_code_blocks.isChecked(), - preserve_indentation=self.preserve_indentation.isChecked(), - generate_instruction_samples=self.instruction_samples.isChecked(), - reasoning_sample_mode=self._reasoning_sample_mode_value(), - prepare_mode=self._prepare_mode_value(), - tokenizer_strategy=self._tokenizer_strategy_value(), - tokenizer_path=Path(self.tokenizer_path.text()) if self.tokenizer_path.text().strip() else None, - dataset_stage=dataset_stage, - tokenizer_training_max_gb=self.tokenizer_training_max_gb.value(), - ) - - def _selected_default_data_paths_for_stage(self, stage: str) -> list[Path]: - """Return selected bundled files that match the dataset purpose. - - Args: - stage: Dataset preparation stage. - - Returns: - Selected paths suitable for the requested stage. - """ - - # Folder selection is the workflow configuration. Do not apply a - # second hardcoded stage filter here; the Dataset Sources tree already - # contains exactly the files selected by the user. - return self._selected_default_data_paths() - - @staticmethod - def _split_path_list(text: str) -> list[Path]: - """Split a semicolon-delimited path field. - - Args: - text: Raw path field text. - - Returns: - Parsed paths. - """ - - return [Path(item.strip().strip('"')) for item in text.split(";") if item.strip()] - - def check_project_health(self) -> None: - """Run a project health check in the background.""" - - self.dataset_log.clear() - self.dataset_progress.setValue(0) - self.dataset_log.append("Checking project health...") - self.project_state.setText("Checking health") - self._run_task( - check_project_health, - ( - Path(self.input_dir.text()), - Path(self.dataset_dir.text()), - Path(self.model_dir.text()), - Path(self.export_dir.text()), - Path(self.gguf_path.text()) if self.gguf_path.text().strip() else None, - Path(self.llama_cpp_dir.text()) if self.llama_cpp_dir.text().strip() else None, - self.device.currentText(), - ), - self._health_check_finished, - self.dataset_log, - self.dataset_progress, - with_progress=True, - button=self.health_check_button, - stop_button=self.stop_dataset_button, - busy_text="Checking Health", - ) - - @Slot(object) - def _health_check_finished(self, result: Any) -> None: - """Display project health check results. - - Args: - result: Project health result. - """ - - self.dataset_progress.setValue(100) - self.dataset_log.append("") - self.dataset_log.append(f"Project health: {result.status.upper()} ({result.summary})") - for check in result.checks: - marker = {"ok": "OK", "warning": "WARN", "error": "ERROR"}.get(check.get("status"), "INFO") - self.dataset_log.append(f"[{marker}] {check.get('name')}: {check.get('detail')}") - self.project_state.setText("Health checked") - self._clear_button_busy("Check Health") - - def preview_dataset(self) -> None: - """Run a dataset preview and quality scan in the background.""" - - self.dataset_log.clear() - self.dataset_progress.setValue(0) - self.dataset_log.append("Previewing dataset...") - self.project_state.setText("Previewing dataset") - self._run_task( - scan_dataset_preview, - (self._dataset_config_from_ui(),), - self._dataset_preview_finished, - self.dataset_log, - self.dataset_progress, - with_progress=True, - button=self.preview_dataset_button, - stop_button=self.stop_dataset_button, - busy_text="Previewing Dataset", - ) - - @Slot(object) - def _dataset_preview_finished(self, result: Any) -> None: - """Display dataset preview and quality scan results. - - Args: - result: Dataset preview result. - """ - - self.dataset_progress.setValue(100) - suffix_text = ", ".join(f"{suffix}: {count}" for suffix, count in - result.suffix_counts.items()) or "none" - self.dataset_log.append("") - self.dataset_log.append( - f"Source files: {result.source_file_count:,}; size: {result.total_bytes / (1024 * 1024):.2f} MB") - self.dataset_log.append(f"File types: {suffix_text}") - self.dataset_log.append( - f"Prepared dataset artifacts: {'found' if result.prepared else 'not complete'}") - self.dataset_log.append( - f"Duplicate scan: {result.duplicate_count:,} file entries in {len(result.duplicate_groups):,} likely group(s).") - self.dataset_log.append( - f"Bad extraction scan: {result.bad_extraction_count:,} suspicious file(s).") - self.dataset_log.append( - f"Code/prose balance: {result.balance_label} ({result.code_preview_count:,}/{result.prose_preview_count:,}).") - self.dataset_log.append( - f"Training readiness: {result.readiness_label} ({result.readiness_score}/100).") - for reason in result.readiness_reasons[:8]: - self.dataset_log.append(f"- {reason}") - self.dataset_quality_duplicates.setText( - f"Duplicates: {result.duplicate_count:,}") - self.dataset_quality_extraction.setText( - f"Extraction: {result.bad_extraction_count:,} flagged") - self.dataset_quality_balance.setText( - f"Balance: {result.balance_label}") - self.dataset_quality_readiness.setText( - f"Readiness: {result.readiness_label} {result.readiness_score}/100") - if result.summary: - self._update_dataset_quality_report(result.summary) - # dataset_quality_duplicates is intentionally left alone here: - # _update_dataset_quality_report() just set it to the block-level - # duplication percentage from the prepared corpus (the more useful, - # actionable metric). Re-setting it to result.duplicate_count (a - # raw duplicate *file* count from the earlier preview scan) would - # silently discard that and always show the old metric instead. - self.dataset_quality_extraction.setText( - f"Extraction: {result.bad_extraction_count:,} flagged") - self.dataset_quality_balance.setText( - f"Balance: {result.balance_label}") - self.dataset_quality_readiness.setText( - f"Readiness: {result.readiness_label} {result.readiness_score}/100") - tokens = int(result.summary.get("token_count", 0) or 0) - vocab = int(result.summary.get("tokenizer_vocab_size", 0) or 0) - self.dataset_log.append( - f"Prepared summary: {tokens:,} tokens, vocab {vocab:,}.") - else: - self.dataset_quality_samples.setText( - f"Preview: {len(result.sample_previews):,} shown") - self.dataset_quality_tokens.setText("Tokens: not prepared") - self.dataset_quality_windows.setText("Windows: not prepared") - self.dataset_quality_vocab.setText("Vocab: not prepared") - self.dataset_quality_code.setText( - f"Code/prose: {result.code_preview_count:,}/{result.prose_preview_count:,}") - self.dataset_quality_cache.setText( - f"Files: {result.source_file_count:,} source") - if result.duplicate_groups: - self.dataset_log.append("") - self.dataset_log.append("Likely duplicates:") - for group in result.duplicate_groups[:8]: - self.dataset_log.append( - f"- {group.get('type')}: {group.get('count')} file(s)") - for path in group.get("files", [])[:4]: - self.dataset_log.append(f" {Path(path).name}") - if result.bad_extraction_files: - self.dataset_log.append("") - self.dataset_log.append("Suspicious extraction files:") - for item in result.bad_extraction_files[:12]: - self.dataset_log.append( - f"- {Path(item.get('path', '')).name}: {item.get('reasons')}") - suggestions: list[str] = [] - if result.duplicate_groups: - suggestions.append( - "Remove or move duplicate files before preparing the final dataset.") - if result.bad_extraction_files: - suggestions.append( - "Replace flagged PDFs with text/source versions, or remove files with bad extraction.") - if result.balance_label == "Prose heavy" and self.code_training_mode.isChecked(): - suggestions.append( - "Add real source-code folders or enable source-file inclusion for a stronger coding model.") - if result.balance_label == "Code heavy": - suggestions.append( - "Add README/tutorial/prose explanations if you want the model to explain code well.") - if result.readiness_label in {"Needs cleanup", "Not ready"}: - suggestions.append( - "Run Preview Dataset again after cleanup and only train once readiness improves.") - if hasattr(self, "dataset_advisor"): - if suggestions: - self.dataset_advisor.setPlainText( - "\n".join(f"- {suggestion}" for suggestion in suggestions)) - else: - self.dataset_advisor.setPlainText( - "No immediate cleanup suggestions. Dataset looks acceptable for the current preview.") - if suggestions: - self.dataset_log.append("") - self.dataset_log.append("Cleanup suggestions:") - for suggestion in suggestions: - self.dataset_log.append(f"- {suggestion}") - if result.issues: - self.dataset_quality_warning.setText( - f"Warnings: {len(result.issues)}") - self.dataset_log.append("") - self.dataset_log.append("Quality notes:") - for issue in result.issues[:12]: - self.dataset_log.append(f"- {issue}") - else: - self.dataset_quality_warning.setText("Warnings: none") - if result.sample_previews: - self.dataset_log.append("") - self.dataset_log.append("Preview samples:") - for index, sample in enumerate(result.sample_previews, start=1): - label = sample.get("language") or sample.get("kind") or "text" - self.dataset_log.append( - f"\n[{index}] {Path(sample.get('path', '')).name} ({label}, {sample.get('characters')} chars)") - self.dataset_log.append( - sample.get("preview", "").replace("\n", "\n ")[:1400]) - self.project_state.setText("Dataset previewed") - self._clear_button_busy("Preview Dataset") - - def prepare_dataset(self) -> None: - """Collect dataset options and start dataset preparation.""" - - config = self._dataset_config_from_ui() - self.dataset_log.clear() - self.dataset_progress.setValue(0) - self._reset_dataset_quality_report() - self.dataset_log.append("Preparing dataset...") - self.dataset_log.append(f"App log file: {self.log_file_path}") - self.dataset_log.append(f"Dataset purpose: {dataset_stage_label(config.dataset_stage)}") - if config.conversation_dataset_paths: - self.dataset_log.append(f"Local conversation JSON/JSONL: {len(config.conversation_dataset_paths)} path(s)") - LOGGER.info("Local conversation JSON/JSONL datasets: %s", "; ".join(str(path) for path in config.conversation_dataset_paths)) - if config.instruction_dataset_paths: - self.dataset_log.append(f"Local instruction JSON/JSONL: {len(config.instruction_dataset_paths)} path(s)") - LOGGER.info("Local instruction JSON/JSONL datasets: %s", "; ".join(str(path) for path in config.instruction_dataset_paths)) - if config.default_data_paths: - self.dataset_log.append(f"Bundled default data: {len(config.default_data_paths)} file(s)") - LOGGER.info("Bundled default data files: %s", "; ".join(str(path) for path in config.default_data_paths)) - if self.include_conversation_datasets.isChecked(): - selected_labels = [ - action.text() - for action in getattr(self, "conversation_dataset_actions", {}).values() - if action.isChecked() and action.isVisible() - ] - if selected_labels: - hf_cache = config.output_dir / "cache" / "huggingface" - self.dataset_log.append(f"Online training datasets: {', '.join(selected_labels)}") - self.dataset_log.append(f"Downloading/loading online data at: {hf_cache}") - LOGGER.info("Online training datasets: %s", ", ".join(selected_labels)) - LOGGER.info("Downloading/loading online data at: %s", hf_cache) - else: - self.dataset_log.append("Online training datasets are enabled, but no dataset is selected for this purpose.") - LOGGER.warning("Online training datasets enabled, but no dataset is selected") - else: - self.dataset_log.append("Online training datasets: off. Local source files only.") - LOGGER.info("Online training datasets: off. Local source files only.") - checked_count = sum( - 1 - for action in getattr(self, "conversation_dataset_actions", {}).values() - if action.isChecked() - ) - if checked_count: - self.dataset_log.append("Checked online dataset choices are ignored until the master checkbox is enabled.") - LOGGER.info("Checked online dataset choices are ignored until the master checkbox is enabled") - LOGGER.info( - "Preparing dataset: input=%s output=%s stage=%s online_datasets=%s conversation_json=%s instruction_json=%s", - config.input_dir, - config.output_dir, - config.dataset_stage, - ",".join(config.conversation_datasets) or "off", - ";".join(str(path) for path in config.conversation_dataset_paths) or "off", - ";".join(str(path) for path in config.instruction_dataset_paths) or "off", - ) - self.project_state.setText("Preparing dataset") - self.dataset_status.setText("Dataset: preparing") - self.auto_vocab_label.setText("Calculating...") - self._run_task( - build_dataset, - (config,), - self._dataset_finished, - self.dataset_log, - self.dataset_progress, - with_progress=True, - button=self.prepare_button, - stop_button=self.stop_dataset_button, - busy_text="Preparing Dataset", - task_kind="dataset", - isolate_process=True, - ) - - @Slot(object) - def _dataset_finished(self, result: Any) -> None: - """Update UI after dataset preparation finishes. - - Args: - result: Dataset build result. - """ - - self.dataset_progress.setValue(100) - self.auto_vocab_label.setText(f"{result.vocab_size:,}") - - LOGGER.info( - "Dataset prepared: documents=%s tokens=%s vocab=%s code=%s prose=%s conversation=%s output=%s", - result.document_count, - result.token_count, - result.vocab_size, - result.code_sample_count, - result.prose_sample_count, - getattr(result, "conversation_sample_count", 0), - result.output_dir, - ) - - self.dataset_log.append( - f"Prepared {result.document_count} documents, " - f"{result.character_count:,} characters, " - f"{result.token_count:,} tokens, " - f"vocab {result.vocab_size:,}." - ) - - if getattr(result, "train_window_count", 0) or getattr(result, - "val_window_count", - 0): - self.dataset_log.append( - f"Training windows: {result.train_window_count:,}; " - f"validation windows: {result.val_window_count:,}." - ) - - self.dataset_log.append( - f"Cache summary: reused {result.cached_file_count:,} file(s), " - f"processed {result.processed_file_count:,} file(s)." - ) - - if getattr(result, "dataset_version_id", ""): - self.dataset_log.append( - f"Dataset version: {result.dataset_version_id}" - ) - - if result.warning: - self.dataset_log.append(f"Recommendation: {result.warning}") - - self._update_dataset_quality_report( - { - "document_count": result.document_count, - "token_count": result.token_count, - "train_window_count": getattr(result, "train_window_count", 0), - "val_window_count": getattr(result, "val_window_count", 0), - "character_count": result.character_count, - "tokenizer_vocab_size": result.vocab_size, - "code_sample_count": result.code_sample_count, - "prose_sample_count": result.prose_sample_count, - "conversation_sample_count": getattr(result, - "conversation_sample_count", - 0), - "cached_file_count": result.cached_file_count, - "processed_file_count": result.processed_file_count, - "skipped_file_count": result.skipped_file_count, - "failed_file_count": result.failed_file_count, - "warning": result.warning, - "sequence_token_stats": getattr(result, "sequence_token_stats", - {}), - "duplicate_block_count": getattr(result, - "duplicate_block_count", 0), - "unique_block_count": getattr(result, "unique_block_count", 0), - "corpus_block_count": getattr(result, "corpus_block_count", 0), - "duplicate_block_ratio": getattr(result, - "duplicate_block_ratio", 0.0), - "unique_block_ratio": getattr(result, "unique_block_ratio", - 1.0), - } - ) - - self.train_data_dir.setText(str(result.output_dir)) - self.project_state.setText("Dataset ready") - - self.dataset_status.setText( - f"Dataset: {result.document_count} files, {result.token_count:,} tokens" - ) - - if result.code_sample_count: - self.dataset_status.setText( - f"Dataset: {result.code_sample_count:,} code, " - f"{result.prose_sample_count:,} prose, " - f"{result.token_count:,} tokens" - ) - - self.refresh_model_estimate() - self.refresh_fine_tune_workflow() - - self._notify_complete( - "dataset", - "Dataset preparation complete", - [ - f"Output: {result.output_dir}", - f"Documents: {result.document_count:,}", - f"Characters: {result.character_count:,}", - f"Tokens: {result.token_count:,}", - f"Vocabulary: {result.vocab_size:,}", - ( - "Windows: " - f"{getattr(result, 'train_window_count', 0):,} training, " - f"{getattr(result, 'val_window_count', 0):,} validation" - ), - ( - "Content mix: " - f"{result.code_sample_count:,} code, " - f"{result.prose_sample_count:,} prose, " - f"{getattr(result, 'conversation_sample_count', 0):,} conversation" - ), - ( - "Files: " - f"{result.processed_file_count:,} processed, " - f"{result.cached_file_count:,} cached, " - f"{result.skipped_file_count:,} skipped, " - f"{result.failed_file_count:,} failed" - ), - f"Dataset version: {getattr(result, 'dataset_version_id', '') or '-'}", - f"Health: {'warning - ' + result.warning if result.warning else 'ready'}", - ], - ) - - self._clear_button_busy("DataSet Prepared") - - def _prepare_mode_value(self) -> str: - """Return the selected dataset preparation mode. - - Returns: - Internal mode value. - """ - - label = self.prepare_mode.currentText() - if label == "Full rebuild": - return "full_rebuild" - if label == "Force reprocess": - return "force_reprocess" - return "incremental" - - def _tokenizer_strategy_value(self) -> str: - """Return the selected tokenizer strategy. - - Returns: - Internal tokenizer strategy value. - """ - - label = self.tokenizer_strategy.currentText() - if label == "Train new tokenizer": - return "train_new" - if label == "Reuse dataset tokenizer": - return "reuse_dataset" - if label == "Import tokenizer.json": - return "import_tokenizer" - return "auto" - - def _reasoning_sample_mode_value(self) -> str: - """Return the selected reasoning sample mode. - - Returns: - Internal reasoning sample mode. - """ - - label = self.reasoning_sample_mode.currentText() - if label == "Detailed code reasoning": - return "detailed" - if label == "No reasoning wrapper": - return "none" - return "scaffold" - - def _dataset_stage_value(self) -> str: - """Return the selected dataset preparation stage. - - Returns: - Dataset stage identifier. - """ - - return self.dataset_stage.currentText().strip().lower().replace(" ", "_") or "base" - - def _set_dataset_stage(self, stage: str) -> None: - """Set the dataset stage combo from an internal stage value. - - Args: - stage: Dataset stage identifier. - """ - - index = self.dataset_stage.findText(stage, Qt.MatchFixedString) - if index < 0: - self.dataset_stage.addItem(stage) - index = self.dataset_stage.count() - 1 - self.dataset_stage.setCurrentIndex(index) - self._update_online_dataset_stage_controls() - - def _update_online_dataset_stage_controls(self) -> None: - """Show and enable online datasets for the selected training stage.""" - - if not hasattr(self, "dataset_stage"): - return - stage = self._dataset_stage_value() - allowed = set(CONVERSATION_DATASET_PRESETS) - include_online = self.include_conversation_datasets.isChecked() - for dataset_id, action in getattr(self, "conversation_dataset_actions", {}).items(): - visible = dataset_id in allowed - action.setVisible(visible) - action.setEnabled(include_online and visible) - if not visible: - action.setChecked(False) - if hasattr(self, "conversation_dataset_button"): - self.conversation_dataset_button.setEnabled(include_online) - self.conversation_sample_limit.setEnabled(include_online) - self._update_conversation_dataset_button_text() - self.conversation_datasets_status.setText( - f"{self.dataset_stage.currentText()}: choose optional online datasets." - if include_online else "Choose optional online datasets, or use local folders only." - ) - - def _selected_conversation_datasets(self) -> list[str]: - """Return selected built-in conversation dataset IDs. - - Returns: - Selected dataset identifiers. - """ - - allowed = set(CONVERSATION_DATASET_PRESETS) - selected = [ - dataset_id - for dataset_id, action in getattr(self, "conversation_dataset_actions", {}).items() - if dataset_id in allowed and action.isChecked() and self.include_conversation_datasets.isChecked() - ] - custom = self.custom_huggingface_dataset.text().strip() - if custom and self.include_conversation_datasets.isChecked(): - selected.append(f"hf_custom:{custom}") - return selected - - def _download_custom_huggingface_dataset(self) -> None: - """Enable the entered Hugging Face dataset for the next preparation run.""" - value = self.custom_huggingface_dataset.text().strip() - if not value: - self.conversation_datasets_status.setText("Enter a Hugging Face dataset ID or URL first.") - return - self.include_conversation_datasets.setChecked(True) - self.conversation_datasets_status.setText( - f"Custom dataset queued: {value}. It will download during dataset preparation." - ) - self._update_conversation_dataset_button_text() - - def _set_selected_conversation_datasets(self, dataset_ids: list[str]) -> None: - """Restore selected conversation dataset actions. - - Args: - dataset_ids: Dataset IDs to select. - """ - - selected = set(dataset_ids) - allowed = set(dataset_ids_for_stage(self._dataset_stage_value())) - for dataset_id, action in getattr(self, "conversation_dataset_actions", {}).items(): - action.setChecked(dataset_id in selected and dataset_id in allowed) - action.setEnabled(self.include_conversation_datasets.isChecked() and dataset_id in allowed) - if hasattr(self, "custom_huggingface_dataset"): - self.custom_huggingface_dataset.setText( - next((value[10:] for value in dataset_ids if value.startswith("hf_custom:")), "") - ) - if hasattr(self, "conversation_sample_limit"): - self.conversation_sample_limit.setEnabled(self.include_conversation_datasets.isChecked()) - self._update_conversation_dataset_button_text() - if hasattr(self, "conversation_datasets_status"): - self._update_online_dataset_stage_controls() - - def _update_conversation_dataset_button_text(self) -> None: - """Refresh the compact online dataset selector label.""" - - if not hasattr(self, "conversation_dataset_button"): - return - allowed = set(dataset_ids_for_stage(self._dataset_stage_value())) if hasattr(self, "dataset_stage") else set() - selected_labels = [ - action.text() - for dataset_id, action in getattr(self, "conversation_dataset_actions", {}).items() - if dataset_id in allowed and action.isChecked() - ] - if not self.include_conversation_datasets.isChecked(): - self.conversation_dataset_button.setText("Online datasets off") - elif not selected_labels: - self.conversation_dataset_button.setText("Choose online datasets") - elif len(selected_labels) == 1: - self.conversation_dataset_button.setText(selected_labels[0]) - else: - self.conversation_dataset_button.setText(f"{len(selected_labels)} online datasets selected") - - def configure_fine_tune_dataset_builder(self) -> None: - """Configure the Ingest tab for the selected fine-tune dataset type.""" - - stage_label = self.fine_tune_dataset_builder_stage.currentText() - stage = { - "Instruction fine-tune": "instruction", - "Conversation fine-tune": "conversation", - "Code fine-tune": "code", - }.get(stage_label, "instruction") - starter_datasets = { - "instruction": ["alpaca_52k"], - "conversation": ["dailydialog"], - "code": ["codealpaca_20k"], - } - self._set_dataset_stage(stage) - self.include_conversation_datasets.setChecked(True) - self._set_selected_conversation_datasets(starter_datasets.get(stage, [])) - if stage == "code": - self.code_training_mode.setChecked(True) - self.include_source_code.setChecked(True) - self.extract_code_blocks.setChecked(True) - self.preserve_indentation.setChecked(True) - self._set_mixture_weights({}) - self._switch_page(0) - self.dataset_log.append(f"Configured Ingest for {dataset_stage_label(stage)}. Import the base tokenizer before preparing.") - self.project_state.setText(f"Configured {dataset_stage_label(stage)} data") - - def _dataset_plan_from_ui(self) -> dict[str, float]: - """Return dataset blueprint state. - - Returns: - Empty mapping because category percentages are disabled. - """ - - return {} - - def _selected_default_data_paths(self) -> list[Path]: - """Return bundled default data files selected in the Dataset Blueprint. - - Returns: - Selected bundled data paths. - """ - - if not hasattr(self, "default_data_actions"): - return [path for path, _category in iter_default_data_files()] - return [ - Path(path) - for path, item in self.default_data_actions.items() - if item.checkState(0) == Qt.Checked - ] - - def _set_selected_default_data_paths(self, paths: Optional[list[Any]]) -> None: - """Restore bundled default data checkbox selections. - - Args: - paths: Saved bundled data file paths. ``None`` means no - preference was ever saved (a brand-new project), and every - file is selected by default. An explicit empty list means - the user deliberately deselected everything, and that - choice is restored as-is rather than falling back to - "select everything" -- previously the two cases were - indistinguishable, so saving a project with nothing - selected silently reset to everything selected on reload. - """ - - if not hasattr(self, "default_data_actions"): - return - if paths is None: - selected = set(self.default_data_actions) - else: - selected = {str(Path(path)) for path in paths} - self.default_data_tree_updating = True - try: - for path, item in self.default_data_actions.items(): - item.setCheckState(0, Qt.Checked if path in selected else Qt.Unchecked) - self._refresh_default_data_category_states() - finally: - self.default_data_tree_updating = False - - def _set_dataset_blueprint_refresh_busy(self, busy: bool) -> None: - """Toggle refresh busy state indicators for the Dataset Sources page.""" - - if hasattr(self, "dataset_plan_refresh_button"): - self.dataset_plan_refresh_button.setEnabled(not busy) - self.dataset_plan_refresh_button.setText("Refreshing..." if busy else "Refresh") - if hasattr(self, "dataset_plan_progress"): - if busy: - self.dataset_plan_progress.setRange(0, 0) - self.dataset_plan_progress.setVisible(True) - else: - self.dataset_plan_progress.setRange(0, 100) - self.dataset_plan_progress.setValue(0) - self.dataset_plan_progress.setVisible(False) - - def refresh_dataset_blueprint_files(self) -> None: - """Reload the Dataset Blueprint file tree from disk.""" - - root = getattr(self, "blueprint_data_root", default_data_root()) - self._refresh_external_dataset_status() - selected_paths = [str(path) for path in self._selected_default_data_paths()] - self._set_dataset_blueprint_refresh_busy(True) - QApplication.processEvents() - try: - self._refresh_dataset_blueprint_source( - Path(root), - saved_paths=selected_paths, - saved_plan=self._dataset_plan_from_ui(), - preset="Custom", - ) - self.project_state.setText("Blueprint refreshed") - LOGGER.info("Dataset blueprint tree refreshed from %s", root) - finally: - self._set_dataset_blueprint_refresh_busy(False) - - def _handle_default_data_tree_changed(self, item: Any, column: int) -> None: - """Handle category and file toggles in the bundled data tree. - - Args: - item: Changed tree item. - column: Changed column index. - """ - - if column != 0 or getattr(self, "default_data_tree_updating", False): - return - data = item.data(0, Qt.UserRole) or {} - if data.get("kind") != "category": - self.default_data_tree_updating = True - try: - self._refresh_default_data_category_states() - finally: - self.default_data_tree_updating = False - if hasattr(self, "_mixture_weights_state"): - delattr(self, "_mixture_weights_state") - return - state = item.checkState(0) - if state == Qt.PartiallyChecked: - return - self.default_data_tree_updating = True - try: - for index in range(item.childCount()): - item.child(index).setCheckState(0, state) - finally: - self.default_data_tree_updating = False - if hasattr(self, "_mixture_weights_state"): - delattr(self, "_mixture_weights_state") - - def _refresh_default_data_category_states(self) -> None: - """Refresh category checkbox states from child file selections.""" - - if not hasattr(self, "default_data_category_items"): - return - for category_item in self.default_data_category_items.values(): - checked = 0 - partial = False - for index in range(category_item.childCount()): - state = category_item.child(index).checkState(0) - if state == Qt.Checked: - checked += 1 - elif state == Qt.PartiallyChecked: - partial = True - if partial or 0 < checked < category_item.childCount(): - category_item.setCheckState(0, Qt.PartiallyChecked) - elif checked == category_item.childCount() and category_item.childCount() > 0: - category_item.setCheckState(0, Qt.Checked) - else: - category_item.setCheckState(0, Qt.Unchecked) - - def _set_dataset_plan(self, plan: dict[str, Any], preset: str = "Custom") -> None: - """Restore high-level dataset blueprint controls. - - Args: - plan: Saved dataset domain percentages. - preset: Saved preset label. - """ - - if not hasattr(self, "dataset_plan_spins"): - return - self._restoring_dataset_plan = True - try: - values = {**dataset_plan_defaults(), **(plan or {})} - for key, widget in self.dataset_plan_spins.items(): - widget.blockSignals(True) - try: - widget.setValue(float(values.get(key, 0.0))) - except (TypeError, ValueError): - widget.setValue(0.0) - widget.blockSignals(False) - self.dataset_plan_preset.blockSignals(True) - if preset == "Custom": - self.dataset_plan_preset.setCurrentText(preset) - else: - self.dataset_plan_preset.setCurrentText("Custom") - self.dataset_plan_preset.blockSignals(False) - finally: - self._restoring_dataset_plan = False - self._update_dataset_plan_total() - - def _dataset_plan_mark_custom(self, *_args: Any) -> None: - """Mark the dataset blueprint as custom after manual edits.""" - - if getattr(self, "_restoring_dataset_plan", False): - return - if hasattr(self, "_mixture_weights_state"): - delattr(self, "_mixture_weights_state") - if hasattr(self, "dataset_plan_preset") and self.dataset_plan_preset.currentText() != "Custom": - self.dataset_plan_preset.blockSignals(True) - self.dataset_plan_preset.setCurrentText("Custom") - self.dataset_plan_preset.blockSignals(False) - - def _update_dataset_plan_total(self) -> None: - """No-op retained for compatibility after blueprint percentage removal.""" - - return - - def normalize_dataset_plan(self) -> None: - """No-op retained for compatibility after blueprint percentage removal.""" - - return - - def apply_dataset_plan_preset(self, preset: str) -> None: - """No-op retained for compatibility after blueprint percentage removal. - - Args: - preset: Preset label from the Dataset Blueprint combo box. - """ - - return - - def apply_dataset_plan_to_ingestion(self) -> None: - """Clear ingestion mixture overrides (category percentages are disabled).""" - - self._set_mixture_weights({}) - if hasattr(self, "dataset_log"): - self.dataset_log.append("Dataset blueprint applied: category percentages are disabled.") - self.project_state.setText("Blueprint applied") - LOGGER.info("Dataset blueprint applied with category percentages disabled") - - def _mixture_weights_from_ui(self) -> dict[str, float]: - """Return dataset mixture weights from the Ingest tab. - - Returns: - Empty mapping because category percentages are disabled. - """ - - if not hasattr(self, "_mixture_weights_state"): - self._mixture_weights_state = {} - return {} - - def _set_mixture_weights(self, weights: dict[str, Any]) -> None: - """Restore dataset mixture weights. - - Args: - weights: Saved mixture weights by source family. - """ - - self._mixture_weights_state = {} - - def _update_mixture_total(self) -> None: - """No-op retained for compatibility after mixture percentage removal.""" - - return - - def _normalize_mixture_weights(self) -> None: - """No-op retained for compatibility after mixture percentage removal.""" - - return - - def _training_launch_target_value(self) -> str: - """Return whether training should launch locally or remotely. - - Returns: - ``local`` or ``remote``. - """ - - if self.training_launch_target.currentText() == "RunPod cloud": - return "runpod" - return "remote" if self.training_launch_target.currentText() == "Remote workers" else "local" - - def _fine_tune_launch_target_value(self) -> str: - """Return whether fine-tuning should launch locally or remotely. - - Returns: - ``local`` or ``remote``. - """ - - if not hasattr(self, "fine_tune_launch_target"): - return "local" - if self.fine_tune_launch_target.currentText() == "RunPod cloud": - return "runpod" - return "remote" if self.fine_tune_launch_target.currentText() == "Remote workers" else "local" - - def _architecture_style_config(self) -> dict[str, Any]: - """Return ModelConfig keyword arguments for the selected block style. - - Returns: - Architecture style settings. - """ - - if self.architecture_style.currentText() == "Llama-like": - return { - "norm_type": "rmsnorm", - "position_encoding": "rope", - "mlp_type": "swiglu", - "rope_theta": self.rope_theta.value(), - } - return { - "norm_type": "layernorm", - "position_encoding": "learned", - "mlp_type": "gelu", - "rope_theta": self.rope_theta.value(), - } - - def _optimizer_value(self) -> str: - """Return the selected optimizer identifier. - - Returns: - Stable optimizer name used by the trainer. - """ - - return { - "AdamW": "adamw", - "Adam": "adam", - "Lion": "lion", - "Adafactor": "adafactor", - }.get(self.optimizer_name.currentText(), "adamw") - - def _scheduler_value(self) -> str: - """Return the selected scheduler identifier. - - Returns: - Stable scheduler name used by the trainer. - """ - - return { - "Warmup linear": "warmup_linear", - "Cosine decay": "cosine", - "Polynomial decay": "polynomial", - "One-cycle": "one_cycle", - "Constant": "constant", - }.get(self.scheduler_name.currentText(), "warmup_linear") - - def _precision_value(self) -> str: - """Return the selected numeric precision identifier. - - Returns: - Stable precision name used by the trainer. - """ - - return { - "FP16": "fp16", - "BF16": "bf16", - "FP32": "fp32", - }.get(self.precision.currentText(), "fp16") - - def _fine_tune_output_path(self) -> Path: - """Return the selected fine-tune output folder. - - Returns: - Folder where fine-tuned artifacts should be written. - """ - - text = self.fine_tune_output_dir.text().strip() if hasattr(self, "fine_tune_output_dir") else "" - if text: - path = Path(text) - elif self.current_project_file is not None: - path = self.current_project_file.parent / "fine_tunes" / "latest" - else: - path = Path(self.model_dir.text()) / "fine_tuned" - try: - if path.resolve() == Path(self.model_dir.text()).resolve(): - path = Path(self.model_dir.text()) / "fine_tuned" - except OSError: - pass - if hasattr(self, "fine_tune_output_dir"): - self.fine_tune_output_dir.setText(str(path)) - return path - - def _refresh_fine_tune_default_output(self, *_args: Any) -> None: - """Keep the fine-tune output folder stage-specific unless a custom folder was chosen.""" - - if not hasattr(self, "fine_tune_output_dir") or self.current_project_file is None: - return - project_dir = self.current_project_file.parent - fine_tunes_dir = project_dir / "fine_tunes" - stage = self._training_stage_value() - stage_folder = { - "instruction": "instruction_latest", - "conversation": "conversation_latest", - "code": "code_latest", - "domain": "domain_latest", - }.get(stage, "fine_tune_latest") - desired = fine_tunes_dir / stage_folder - current_text = self.fine_tune_output_dir.text().strip() - if not current_text: - self.fine_tune_output_dir.setText(str(desired)) - return - try: - current = Path(current_text) - current_resolved = current.resolve() - fine_tunes_resolved = fine_tunes_dir.resolve() - except OSError: - return - managed_names = { - "latest", - "fine_tune", - "fine_tuned", - "instruction", - "conversation", - "code", - "domain", - "instruction_latest", - "conversation_latest", - "code_latest", - "domain_latest", - "fine_tune_latest", - } - if current_resolved.parent == fine_tunes_resolved and current.name in managed_names: - self.fine_tune_output_dir.setText(str(desired)) - - def _prepare_fine_tune_run_folder(self, training_config: TrainingConfig) -> None: - """Create fine-tune folders and snapshot the base checkpoint. - - Args: - training_config: Fine-tune training configuration. - """ - - output_dir = Path(training_config.output_dir) - output_dir.mkdir(parents=True, exist_ok=True) - checkpoints_dir = output_dir / "checkpoints" - checkpoints_dir.mkdir(parents=True, exist_ok=True) - base_checkpoint = training_config.fine_tune_from_checkpoint - if base_checkpoint is None: - return - base_checkpoint = Path(base_checkpoint) - if not base_checkpoint.exists(): - return - try: - base_resolved = base_checkpoint.resolve() - output_resolved = output_dir.resolve() - if base_resolved == (output_resolved / base_checkpoint.name) or output_resolved in base_resolved.parents: - raise ValueError( - "Fine-tune base checkpoint must be outside the selected fine-tune output folder. " - "Choose the original pretrained model checkpoint instead." - ) - except RuntimeError as exc: - raise ValueError(f"Could not validate fine-tune base checkpoint path: {exc}") from exc - snapshot_dir = output_dir / "base_model" - snapshot_dir.mkdir(parents=True, exist_ok=True) - copied_checkpoint = snapshot_dir / base_checkpoint.name - if not copied_checkpoint.exists() or copied_checkpoint.stat().st_size != base_checkpoint.stat().st_size: - shutil.copy2(base_checkpoint, copied_checkpoint) - base_parent = base_checkpoint.parent - for file_name in ("tokenizer.json", "training_summary.json", "model_lineage.json"): - source = base_parent / file_name - if source.exists(): - target = snapshot_dir / file_name - if not target.exists() or target.stat().st_size != source.stat().st_size: - shutil.copy2(source, target) - manifest = { - "base_checkpoint": str(base_checkpoint), - "copied_checkpoint": str(copied_checkpoint), - "fine_tune_output": str(output_dir), - "created_at": datetime.now().isoformat(timespec="seconds"), - } - (snapshot_dir / "base_model_manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8") - self.fine_tune_log.append(f"Base model snapshot: {copied_checkpoint}") - - def _training_output_dir_for_mode(self, training_mode: Optional[str]) -> Path: - """Return the output folder for a training mode. - - Args: - training_mode: Training mode override. - - Returns: - Base model or fine-tune output folder. - """ - - return self._fine_tune_output_path() if training_mode == "fine_tune" else Path(self.model_dir.text()) - - def _training_mode_value(self) -> str: - """Return the selected training mode identifier. - - Returns: - Stable training mode used by the trainer. - """ - - return { - "Pretrain from scratch": "pretrain", - "Fine-tune checkpoint": "fine_tune", - "Instruction fine-tune": "fine_tune", - "Conversation fine-tune": "fine_tune", - "Code fine-tune": "fine_tune", - }.get(self.training_mode.currentText(), "pretrain") - - def _training_stage_value(self) -> str: - """Return the higher-level training stage selected in the UI. - - Returns: - Training stage identifier. - """ - - return { - "Pretrain from scratch": "base", - "Fine-tune checkpoint": "domain", - "Instruction fine-tune": "instruction", - "Conversation fine-tune": "conversation", - "Code fine-tune": "code", - }.get(self.training_mode.currentText(), "base") - - def _peft_method_value(self) -> str: - """Return the selected PEFT method identifier. - - Returns: - Stable PEFT method used by the trainer. - """ - - return { - "Full fine-tune": "none", - "LoRA adapters": "lora", - }.get(self.peft_method.currentText(), "none") - - def _lora_target_value(self) -> str: - """Return selected LoRA target groups. - - Returns: - Comma-separated target group string. - """ - - return { - "Attention projections": "attention", - "MLP projections": "mlp", - "Attention + MLP": "attention,mlp", - }.get(self.lora_targets.currentText(), "attention") - - def _update_training_mode_controls(self) -> None: - """Enable fine-tune controls only when fine-tuning is selected.""" - - enabled = self._training_mode_value() == "fine_tune" - lora_enabled = enabled and self._peft_method_value() == "lora" - self.fine_tune_checkpoint.setEnabled(enabled) - self.peft_method.setEnabled(enabled) - self.fine_tune_check_button.setEnabled(enabled) - self.lora_rank.setEnabled(lora_enabled) - self.lora_alpha.setEnabled(lora_enabled) - self.lora_dropout.setEnabled(lora_enabled) - self.lora_targets.setEnabled(lora_enabled) - self.refresh_fine_tune_workflow() - - def _current_dataset_summary(self) -> dict[str, Any]: - """Read the active prepared dataset summary. - - Returns: - Dataset summary dictionary, or an empty dictionary. - """ - - summary_path = Path(self.train_data_dir.text()) / "dataset_summary.json" - if not summary_path.exists(): - summary_path = Path(self.dataset_dir.text()) / "dataset_summary.json" - if not summary_path.exists(): - return {} - try: - data = json.loads(summary_path.read_text(encoding="utf-8")) - return data if isinstance(data, dict) else {} - except Exception as exc: - LOGGER.warning("Could not read dataset summary %s: %s", summary_path, exc) - return {} - - def _fine_tune_dataset_stage_status(self) -> tuple[bool, str]: - """Check whether the prepared dataset matches the fine-tune type. - - Returns: - Tuple containing whether the workflow may proceed and a user-facing message. - """ - - expected_stage = self._training_stage_value() - summary = self._current_dataset_summary() - if not summary: - return False, "Dataset: not prepared. Prepare the fine-tune dataset first." - dataset_stage = str(summary.get("dataset_stage") or self._dataset_stage_value()) - tokens = int(summary.get("token_count", 0) or 0) - vocab = int(summary.get("tokenizer_vocab_size", 0) or 0) - stage_name = dataset_stage_label(dataset_stage) if dataset_stage in {"base", "instruction", "conversation", "code"} else dataset_stage - details = f"{stage_name}, {tokens:,} tokens, vocab {vocab:,}" - if expected_stage == "instruction" and dataset_stage != "instruction": - return False, f"Dataset mismatch: selected Instruction fine-tune, but prepared dataset is {details}." - if expected_stage == "conversation" and dataset_stage != "conversation": - return False, f"Dataset mismatch: selected Conversation fine-tune, but prepared dataset is {details}." - if expected_stage == "code" and dataset_stage != "code": - return False, f"Dataset mismatch: selected Code fine-tune, but prepared dataset is {details}." - if expected_stage == "domain" and dataset_stage == "base": - return True, f"Dataset warning: {details}. Base datasets usually belong to pretraining; continue only for domain adaptation." - return True, f"Dataset ready: {details}." - - def refresh_fine_tune_workflow(self) -> None: - """Refresh fine-tune workflow guidance in the Fine-Tuning tab.""" - - if not hasattr(self, "fine_tune_dataset_status"): - return - self._refresh_fine_tune_default_output() - ok, message = self._fine_tune_dataset_stage_status() - self.fine_tune_dataset_status.setText(message) - self.fine_tune_dataset_status.setProperty("state", "ok" if ok else "warning") - self.fine_tune_dataset_status.style().unpolish(self.fine_tune_dataset_status) - self.fine_tune_dataset_status.style().polish(self.fine_tune_dataset_status) - - def apply_recommended_fine_tune_settings(self) -> None: - """Apply conservative fine-tuning defaults for the selected workflow.""" - - stage = self._training_stage_value() - synced = self._sync_architecture_from_fine_tune_base() - self._set_combo_text(self.peft_method, "LoRA adapters") - self.lora_dropout.setValue(0.05) - self._set_combo_text(self.lora_targets, "Attention projections") - self.max_grad_norm.setValue(0.5) - self.weight_decay.setValue(0.05) - self._set_combo_by_data(self.scheduler_name, "cosine", { - "warmup_linear": "Warmup linear", - "cosine": "Cosine decay", - "polynomial": "Polynomial decay", - "one_cycle": "One-cycle", - "constant": "Constant", - }) - if stage == "conversation": - self.lora_rank.setValue(16) - self.lora_alpha.setValue(32.0) - self.learning_rate.setValue(0.00003) - self.epochs.setValue(max(1, min(self.epochs.value(), 2))) - elif stage == "code": - self.lora_rank.setValue(8) - self.lora_alpha.setValue(16.0) - self.lora_dropout.setValue(0.05) - self.learning_rate.setValue(0.00005) - self.max_grad_norm.setValue(0.5) - self.epochs.setValue(max(1, min(self.epochs.value(), 3))) - elif stage == "instruction": - self.lora_rank.setValue(8) - self.lora_alpha.setValue(16.0) - self.learning_rate.setValue(0.00005) - self.epochs.setValue(max(1, min(self.epochs.value(), 3))) - else: - self.lora_rank.setValue(8) - self.lora_alpha.setValue(16.0) - self.learning_rate.setValue(0.00005) - self._update_training_mode_controls() - message = "Recommended LoRA settings applied." - if synced: - message += "\nArchitecture was synced from the selected base checkpoint." - message += "\nUse Check Fine-tune before starting so checkpoint and tokenizer compatibility are verified." - self.fine_tune_preview.setText(message) - - def _sync_architecture_from_fine_tune_base(self) -> bool: - """Sync architecture controls from the selected fine-tune base checkpoint. - - Returns: - True when a checkpoint was read and architecture controls were updated. - """ - - if not hasattr(self, "fine_tune_checkpoint"): - return False - checkpoint_text = self.fine_tune_checkpoint.text().strip() - if not checkpoint_text: - return False - checkpoint_path = Path(checkpoint_text) - if not checkpoint_path.exists(): - return False - try: - checkpoint = torch.load(checkpoint_path, map_location="cpu") - except Exception as exc: - LOGGER.warning("Could not read fine-tune base checkpoint %s: %s", checkpoint_path, exc) - return False - model_config = checkpoint.get("model_config", {}) if isinstance(checkpoint, dict) else {} - if not isinstance(model_config, dict): - return False - mappings = { - "embedding_size": self.n_embd, - "head_count": self.n_head, - "layer_count": self.n_layer, - # NOT self.context_length -- that is the Dataset tab's tokenizer - # window-size setting (DatasetConfig.context_length), an - # unrelated dataset-preparation parameter. _current_model_config() - # reads self.train_context_length for ModelConfig.context_length, - # which is the field resume-compatibility actually checks. - "context_length": self.train_context_length, - } - for key, widget in mappings.items(): - if key in model_config: - try: - widget.setValue(int(model_config[key])) - except (TypeError, ValueError): - LOGGER.warning("Invalid %s in checkpoint %s: %r", key, checkpoint_path, model_config[key]) - if "dropout" in model_config: - try: - self.dropout.setValue(float(model_config["dropout"])) - except (TypeError, ValueError): - LOGGER.warning("Invalid dropout in checkpoint %s: %r", checkpoint_path, model_config["dropout"]) - if "rope_theta" in model_config: - try: - self.rope_theta.setValue(float(model_config["rope_theta"])) - except (TypeError, ValueError): - LOGGER.warning("Invalid rope_theta in checkpoint %s: %r", checkpoint_path, model_config["rope_theta"]) - if "bias" in model_config: - self.use_bias.setChecked(bool(model_config["bias"])) - norm_type = str(model_config.get("norm_type", "layernorm")).lower() - position_encoding = str(model_config.get("position_encoding", "learned")).lower() - mlp_type = str(model_config.get("mlp_type", "gelu")).lower() - if norm_type == "rmsnorm" or position_encoding == "rope" or mlp_type == "swiglu": - # Must match training_tab.py's actual combo item text exactly - # ("Llama-like") -- _set_combo_text() silently no-ops on a - # non-editable combo when the text doesn't match any item, so a - # wrong string here does not raise or log anything. It used to - # say "Modern LLM", which does not exist as an option: this - # left architecture_style un-synced while every other field - # (n_embd, n_head, n_layer, ...) synced correctly, guaranteeing - # a resume-compatibility mismatch on norm_type/position_encoding - # /mlp_type with no indication of why. - self._set_combo_text(self.architecture_style, "Llama-like") - else: - self._set_combo_text(self.architecture_style, "Classic GPT") - attention_type = str(model_config.get("attention_type", "mha")).lower() - self._set_combo_by_data( - self.attention_type, - attention_type, - { - "mha": "Multi-head", - "mqa": "Multi-query", - "gqa": "Grouped-query", - }, - ) - if "kv_head_count" in model_config: - try: - self.kv_head_count.setValue(int(model_config["kv_head_count"])) - except (TypeError, ValueError): - LOGGER.warning("Invalid kv_head_count in checkpoint %s: %r", checkpoint_path, model_config["kv_head_count"]) - backend = str(model_config.get("attention_backend", "sdpa")).lower() - self._set_combo_by_data( - self.attention_backend, - backend, - { - "sdpa": "SDPA / Flash when available", - "eager": "PyTorch eager", - }, - ) - if "attention_window" in model_config: - try: - self.attention_window.setValue(int(model_config["attention_window"])) - except (TypeError, ValueError): - LOGGER.warning("Invalid attention_window in checkpoint %s: %r", checkpoint_path, model_config["attention_window"]) - LOGGER.info("Fine-tune architecture synced from base checkpoint: %s", checkpoint_path) - return True - - def _attention_type_value(self) -> str: - """Return the selected attention layout identifier. - - Returns: - Stable attention type used by the model. - """ - - return { - "Multi-head": "mha", - "Grouped-query": "gqa", - "Multi-query": "mqa", - }.get(self.attention_type.currentText(), "mha") - - def _attention_backend_value(self) -> str: - """Return the selected attention backend identifier. - - Returns: - Stable attention backend used by the model. - """ - - return { - "SDPA / Flash when available": "sdpa", - "Manual": "manual", - }.get(self.attention_backend.currentText(), "sdpa") - - def apply_training_profile(self) -> None: - """Apply the selected optimizer/scheduler/regularization profile. - - Each branch below explicitly sets every field it conceptually owns - (optimizer, scheduler, LR/regularization, precision/memory knobs, - batch shape, and early-stopping patience), even fields that happen - to match the previous profile's value. This is deliberate: profiles - must be idempotent when switched between, or a field set by a - previously applied profile (e.g. activation_checkpointing=True from - Low-memory) can silently survive into a later profile that never - mentions it, producing a configuration no single profile actually - intended. - - Two categories of fields are deliberately NOT touched here: - - attention_type / kv_head_count: an architecture choice, not a - training-strategy choice. Low-memory sets these to - Grouped-query because that specific profile is about reducing - memory end-to-end; the other profiles leave whatever the user - has selected alone rather than silently reverting it. - - training_mode / peft_method / lora_* (Code fine-tune only): - these belong to the fine-tuning tab's widgets, not this tab's. - """ - - profile = self.training_profile.currentText() - if profile == "Low-memory": - self._set_combo_text(self.optimizer_name, "Adafactor") - self._set_combo_text(self.scheduler_name, "Cosine decay") - self.learning_rate.setValue(0.0002) - self.weight_decay.setValue(0.05) - self.min_lr_ratio.setValue(0.05) - self.polynomial_power.setValue(1.0) - self.max_grad_norm.setValue(1.0) - self._set_combo_text(self.precision, "BF16" if torch.cuda.is_available() else "FP32") - self.use_amp.setChecked(True) - self._set_combo_text(self.attention_type, "Grouped-query") - self.kv_head_count.setValue(max(1, self.n_head.value() // 2)) - self.activation_checkpointing.setChecked(True) - # The two knobs that most directly control peak memory: shrink - # the batch and make up the lost effective batch size with - # gradient accumulation, and avoid extra data-loader worker - # processes competing for memory. - self.batch_size.setValue(4) - self.gradient_accumulation.setValue(4) - self.data_loader_workers.setValue(0) - self.warmup_steps.setValue(100) - self.dropout.setValue(0.1) - self.early_stopping_patience.setValue(3) - elif profile == "Code fine-tune": - self._set_combo_text(self.optimizer_name, "AdamW") - self._set_combo_text(self.scheduler_name, "Cosine decay") - self.learning_rate.setValue(0.00005) - self.weight_decay.setValue(0.05) - self.min_lr_ratio.setValue(0.1) - self.polynomial_power.setValue(1.0) - self.max_grad_norm.setValue(0.5) - self._set_combo_text(self.precision, "FP16") - self.use_amp.setChecked(True) - self.activation_checkpointing.setChecked(False) - self.batch_size.setValue(16) - self.gradient_accumulation.setValue(1) - self.data_loader_workers.setValue(0) - self.warmup_steps.setValue(50) - self.dropout.setValue(0.05) - # Fine-tuning generally needs less patience than a full - # pretraining run before validation loss plateaus meaningfully. - self.early_stopping_patience.setValue(2) - self._set_combo_text(self.training_mode, "Fine-tune checkpoint") - self._set_combo_text(self.peft_method, "LoRA adapters") - self.lora_rank.setValue(8) - self.lora_alpha.setValue(16.0) - self.lora_dropout.setValue(0.05) - self._set_combo_text(self.lora_targets, "Attention projections") - elif profile == "Experimental Lion": - self._set_combo_text(self.optimizer_name, "Lion") - self._set_combo_text(self.scheduler_name, "One-cycle") - self.learning_rate.setValue(0.0001) - self.weight_decay.setValue(0.1) - self.min_lr_ratio.setValue(0.01) - self.polynomial_power.setValue(1.0) - self.max_grad_norm.setValue(1.0) - # Lion is reported to be more sensitive to fp16 under/overflow - # than AdamW; prefer bf16 where available, fp32 otherwise. - self._set_combo_text(self.precision, "BF16" if torch.cuda.is_available() else "FP32") - self.use_amp.setChecked(True) - self.activation_checkpointing.setChecked(False) - self.batch_size.setValue(16) - self.gradient_accumulation.setValue(1) - self.data_loader_workers.setValue(0) - self.warmup_steps.setValue(100) - self.dropout.setValue(0.1) - self.early_stopping_patience.setValue(3) - else: - self._set_combo_text(self.optimizer_name, "AdamW") - self._set_combo_text(self.scheduler_name, "Cosine decay") - self.learning_rate.setValue(0.0003) - self.weight_decay.setValue(0.1) - self.min_lr_ratio.setValue(0.1) - self.polynomial_power.setValue(1.0) - self.max_grad_norm.setValue(1.0) - self._set_combo_text(self.precision, "FP16") - self.use_amp.setChecked(True) - self.activation_checkpointing.setChecked(False) - self.batch_size.setValue(16) - self.gradient_accumulation.setValue(1) - self.data_loader_workers.setValue(0) - self.warmup_steps.setValue(100) - self.dropout.setValue(0.1) - self.early_stopping_patience.setValue(3) - self._update_training_mode_controls() - self.refresh_model_estimate() - self.training_log.append(f"Applied training profile: {profile}") - - def _tokenizer_strategy_reuses(self) -> bool: - """Return whether current tokenizer strategy ignores vocabulary controls. - - Returns: - True when an existing tokenizer is selected directly. - """ - - return self.tokenizer_strategy.currentText() in {"Reuse dataset tokenizer", "Import tokenizer.json"} - - def _update_tokenizer_strategy_controls(self) -> None: - """Enable only the tokenizer inputs relevant to the selected strategy.""" - - imports_tokenizer = self.tokenizer_strategy.currentText() == "Import tokenizer.json" - reuses_tokenizer = self._tokenizer_strategy_reuses() - if hasattr(self, "tokenizer_path_row"): - self.tokenizer_path_row.setEnabled(imports_tokenizer) - self.tokenizer_path.setEnabled(imports_tokenizer) - self.auto_vocab.setEnabled(not reuses_tokenizer) - self.manual_vocab_size.setEnabled(not reuses_tokenizer and not self.auto_vocab.isChecked()) - self.min_frequency.setEnabled(not reuses_tokenizer) - - def _update_model_estimate_chips( - self, - estimate: dict[str, Any], - model_config: Optional[ModelConfig] = None, - training_config: Optional[TrainingConfig] = None, - train_tokens: int = 0, - ) -> None: - """Update model and VRAM estimate chips. - - Args: - estimate: Estimate dictionary from the training planning service. - model_config: Model architecture used for the estimate. - training_config: Training options used for the estimate. - train_tokens: Number of available training tokens. - """ - - params = int(estimate.get("parameters", 0)) - checkpoint_bytes = float(estimate.get("checkpoint_bytes", 0)) - vram_bytes = float(estimate.get("vram_bytes", 0)) - self.model_size_metric.setText(f"Model: {params / 1_000_000:.2f}M, ckpt {format_bytes(checkpoint_bytes)}") - self.vram_estimate_metric.setText(f"VRAM est: {format_bytes(vram_bytes)}") - parameter_breakdown = estimate.get("parameter_breakdown", {}) or {} - memory_breakdown = estimate.get("memory_breakdown", {}) or {} - embedding_params = int(parameter_breakdown.get("token_embedding", 0)) + int( - parameter_breakdown.get("position_embedding", 0) - ) - attention_params = int(parameter_breakdown.get("attention", 0)) - mlp_params = int(parameter_breakdown.get("mlp", 0)) - norm_params = int(parameter_breakdown.get("norms", 0)) - self.parameter_breakdown_metric.setText( - "Params: " - f"emb {self._compact_number(embedding_params)}, " - f"attn {self._compact_number(attention_params)}, " - f"mlp {self._compact_number(mlp_params)}" - ) - self._tip( - self.parameter_breakdown_metric, - ( - f"Embedding: {embedding_params:,}\n" - f"Attention: {attention_params:,}\n" - f"MLP: {mlp_params:,}\n" - f"Norms/output: {norm_params:,}\n" - f"Total: {params:,}" - ), - ) - weights = float(memory_breakdown.get("weights", 0)) - optimizer = float(memory_breakdown.get("optimizer", 0)) - activations = float(memory_breakdown.get("activations", 0)) - kv_cache = float(memory_breakdown.get("kv_cache", 0)) - self.memory_breakdown_metric.setText( - f"Memory: w {format_bytes(weights)}, opt {format_bytes(optimizer)}, act {format_bytes(activations)}" - ) - self._tip( - self.memory_breakdown_metric, - ( - f"Weights: {format_bytes(weights)}\n" - f"Optimizer state: {format_bytes(optimizer)}\n" - f"Activations: {format_bytes(activations)}\n" - f"KV cache estimate: {format_bytes(kv_cache)}\n" - f"Total training estimate: {format_bytes(vram_bytes)}" - ), - ) - self._update_architecture_advisor(estimate, model_config, training_config, train_tokens) - - def _update_architecture_advisor( - self, - estimate: dict[str, Any], - model_config: Optional[ModelConfig], - training_config: Optional[TrainingConfig], - train_tokens: int, - ) -> None: - """Update the compact architecture advisor chip. - - Args: - estimate: Estimate dictionary from the training planning service. - model_config: Model architecture used for the estimate. - training_config: Training options used for the estimate. - train_tokens: Number of available training tokens. - """ - - params = max(int(estimate.get("parameters", 0) or 0), 1) - tokens_per_param = float(train_tokens) / float(params) if train_tokens > 0 else 0.0 - vram_bytes = float(estimate.get("vram_bytes", 0) or 0) - notes: list[str] = [] - if tokens_per_param <= 0: - label = "Advisor: prepare data" - notes.append("Prepare a dataset to compare token budget against model size.") - elif tokens_per_param < 20: - label = "Advisor: data-light" - notes.append( - f"Token budget is about {tokens_per_param:.1f} tokens per parameter. More data or fewer epochs may reduce overfitting." - ) - elif tokens_per_param > 150: - label = "Advisor: data-rich" - notes.append( - f"Token budget is about {tokens_per_param:.1f} tokens per parameter. The model may be small for this much data." - ) - else: - label = "Advisor: balanced" - notes.append(f"Token budget is about {tokens_per_param:.1f} tokens per parameter.") - if model_config is not None: - if model_config.context_length >= 2048 and model_config.embedding_size <= 256: - notes.append("Long context with a small embedding can be memory-heavy without adding much capacity.") - if model_config.attention_type in {"grouped_query", "multi_query"}: - notes.append("Grouped/multi-query attention reduces KV memory and is useful for longer contexts.") - if model_config.mlp_type == "swiglu" and model_config.norm_type == "rmsnorm": - notes.append("Llama-like blocks improve modern compatibility but must match checkpoints when resuming.") - if training_config is not None and training_config.device == "cuda" and vram_bytes > 3.5 * 1024**3: - notes.append("Estimated VRAM is high for 4 GB GPUs. Try lower batch, context, embedding, or layers.") - if label == "Advisor: balanced": - label = "Advisor: memory check" - self.architecture_advisor_metric.setText(label) - self._tip(self.architecture_advisor_metric, "\n".join(notes)) - - @staticmethod - def _compact_number(value: int) -> str: - """Format a large count for tight metric chips. - - Args: - value: Count to format. - - Returns: - Compact display string. - """ - - magnitude = abs(value) - if magnitude >= 1_000_000_000: - return f"{value / 1_000_000_000:.1f}B" - if magnitude >= 1_000_000: - return f"{value / 1_000_000:.1f}M" - if magnitude >= 1_000: - return f"{value / 1_000:.1f}K" - return str(value) - - def _current_model_config(self, vocab_size: int = 1) -> ModelConfig: - """Build a model config from the current AI tab settings. - - Args: - vocab_size: Tokenizer vocabulary size to use. - - Returns: - Current model configuration. - """ - - return ModelConfig( - vocab_size=vocab_size, - context_length=self.train_context_length.value(), - embedding_size=self.n_embd.value(), - head_count=self.n_head.value(), - layer_count=self.n_layer.value(), - dropout=self.dropout.value(), - bias=self.use_bias.isChecked(), - attention_type=self._attention_type_value(), - kv_head_count=self.kv_head_count.value(), - attention_backend=self._attention_backend_value(), - attention_window=self.attention_window.value(), - **self._architecture_style_config(), - ) - - def _current_training_config( - self, - resume_path: Optional[Path] = None, - training_mode: Optional[str] = None, - ) -> TrainingConfig: - """Build a training config from the current AI tab settings. - - Args: - resume_path: Optional specific checkpoint to resume from. - training_mode: Optional explicit training mode override. - - Returns: - Current training configuration. - """ - - return TrainingConfig( - output_dir=self._training_output_dir_for_mode(training_mode), - epochs=self.epochs.value(), - batch_size=self.batch_size.value(), - learning_rate=self.learning_rate.value(), - weight_decay=self.weight_decay.value(), - optimizer_name=self._optimizer_value(), - scheduler_name=self._scheduler_value(), - scheduler_min_lr_ratio=self.min_lr_ratio.value(), - polynomial_power=self.polynomial_power.value(), - gradient_accumulation=self.gradient_accumulation.value(), - sample_stride=self.sample_stride.value(), - warmup_steps=self.warmup_steps.value(), - eval_interval=self.eval_interval.value(), - max_eval_batches=self.max_eval_batches.value(), - save_interval=self.save_interval.value(), - data_loader_workers=self.data_loader_workers.value(), - max_grad_norm=self.max_grad_norm.value(), - activation_checkpointing=self.activation_checkpointing.isChecked(), - device=self.device.currentText(), - use_amp=self.use_amp.isChecked(), - precision=self._precision_value(), - seed=self.seed.value(), - training_mode=training_mode or self._training_mode_value(), - fine_tune_from_checkpoint=( - Path(self.fine_tune_checkpoint.text()) - if training_mode != "pretrain" and self.fine_tune_checkpoint.text().strip() - else None - ), - peft_method="none" if training_mode == "pretrain" else self._peft_method_value(), - lora_rank=self.lora_rank.value(), - lora_alpha=self.lora_alpha.value(), - lora_dropout=self.lora_dropout.value(), - lora_target_modules=self._lora_target_value(), - resume=self.resume_training.isChecked(), - resume_from_checkpoint=resume_path if self.resume_training.isChecked() else None, - require_compatible_resume=self.resume_safety.isChecked(), - early_stopping=self.early_stopping.isChecked(), - early_stopping_patience=self.early_stopping_patience.value(), - ) - - def _current_training_vocab_size(self, data_dir: Path) -> int: - """Return the tokenizer vocabulary size for the current training dataset. - - Args: - data_dir: Prepared dataset folder. - - Returns: - Vocabulary size, or zero if unavailable. - """ - - summary_path = data_dir / "dataset_summary.json" - if summary_path.exists(): - summary = json.loads(summary_path.read_text(encoding="utf-8")) - vocab_size = int(summary.get("tokenizer_vocab_size", 0) or 0) - if vocab_size > 0: - return vocab_size - tokenizer_path = data_dir / "tokenizer.json" - if tokenizer_path.exists(): - tokenizer_data = json.loads(tokenizer_path.read_text(encoding="utf-8")) - vocab = tokenizer_data.get("model", {}).get("vocab", {}) - if isinstance(vocab, dict): - return len(vocab) - return 0 - - def _checkpoint_vocab_size(self, checkpoint_path: Path) -> int: - """Return the tokenizer vocabulary size saved in a checkpoint. - - Args: - checkpoint_path: Checkpoint file to inspect. - - Returns: - Saved vocabulary size, or zero when unavailable. - """ - - try: - checkpoint = torch.load(checkpoint_path, map_location="cpu") - model_config = checkpoint.get("model_config", {}) - if isinstance(model_config, dict): - return int(model_config.get("vocab_size", 0) or 0) - except Exception as exc: - LOGGER.warning("Could not inspect checkpoint vocab size for %s: %s", checkpoint_path, exc) - return 0 - - @staticmethod - def _tokenizer_mismatch_help(checkpoint_vocab: int, dataset_vocab: int) -> str: - """Return user-facing help for tokenizer mismatch errors. - - Args: - checkpoint_vocab: Vocabulary size saved in the checkpoint. - dataset_vocab: Vocabulary size in the prepared dataset. - - Returns: - Help text. - """ - - return ( - f"Tokenizer mismatch: base checkpoint vocab is {checkpoint_vocab:,}, " - f"but prepared dataset vocab is {dataset_vocab:,}.\n" - "Fix: rebuild the fine-tune dataset using the exact tokenizer from the base model. " - "In Ingest, set Tokenizer policy to Import tokenizer.json and choose the tokenizer.json " - "beside the base checkpoint, then prepare the fine-tune dataset again." - ) - - def _training_run_artifacts(self, output_dir: Path) -> list[Path]: - """Return training-run artifacts in a model output folder. - - Args: - output_dir: Model output folder to inspect. - - Returns: - Existing training-run artifact paths. - """ - - candidates = [ - output_dir / "checkpoints", - output_dir / "final_model.pt", - output_dir / "final_adapter.pt", - output_dir / "training_summary.json", - output_dir / "training_history.json", - output_dir / "model_lineage.json", - ] - return [path for path in candidates if path.exists()] - - def _clear_training_run_artifacts(self, output_dir: Path) -> list[Path]: - """Delete resumable training artifacts from a model output folder. - - Args: - output_dir: Model output folder to clean. - - Returns: - Paths that were removed. - """ - - output_dir = output_dir.resolve() - removed: list[Path] = [] - candidates = self._training_run_artifacts(output_dir) - for path in candidates: - try: - resolved = path.resolve() - except FileNotFoundError: - resolved = path - if output_dir not in resolved.parents and resolved != output_dir: - LOGGER.warning("Skipped training cleanup outside model output folder: %s", path) - continue - if path.is_dir(): - shutil.rmtree(path) - removed.append(path) - elif path.exists(): - path.unlink() - removed.append(path) - return removed - - def _selected_resume_path(self) -> Optional[Path]: - """Return the selected or latest checkpoint path. - - Returns: - Checkpoint path, or ``None`` when no checkpoint exists. - """ - - if self.resume_checkpoint.text().strip(): - return Path(self.resume_checkpoint.text()) - return latest_checkpoint(Path(self.model_dir.text()) / "checkpoints") - - def preview_resume_compatibility(self) -> None: - """Preview whether the selected checkpoint can resume safely.""" - - if not self.resume_training.isChecked(): - self.resume_training_preview.setText("[INFO] Resume latest is off. Enable resume to continue from a checkpoint.") - return - resume_path = self._selected_resume_path() - if resume_path is None: - self.resume_training_preview.setText("[INFO] No checkpoint found in the current model folder.") - return - if not resume_path.exists(): - self.resume_training_preview.setText(f"[BLOCK] Checkpoint does not exist:\n{resume_path}") - return - try: - vocab_size = self._current_training_vocab_size(Path(self.train_data_dir.text())) - if vocab_size <= 0: - self.resume_training_preview.setText("[BLOCK] Could not determine current dataset tokenizer vocabulary size.") - return - model_config = self._current_model_config(vocab_size=vocab_size) - # Explicit override: this is the AI/Training tab's own "Check - # Resume" button, checking a pretrain checkpoint. Without this, - # training_mode falls back to reading the separate Fine-Tuning - # tab's mode combo (self.training_mode), which defaults to - # "Instruction fine-tune" on a fresh session -- resolving to - # "fine_tune" and making training_config.validate() below raise - # "fine_tune_from_checkpoint is required for fine_tune mode", - # a confusing error unrelated to what the user is checking. - training_config = self._current_training_config(resume_path, training_mode="pretrain") - model_config.validate() - training_config.validate() - report = check_resume_compatibility(resume_path, model_config, training_config) - errors = list(report.errors) - if training_config.require_compatible_resume: - if not report.can_load_optimizer_state: - errors.append("Safe resume requires matching optimizer state.") - if not report.can_load_scheduler_state: - errors.append("Safe resume requires matching scheduler state.") - if not report.can_load_scaler_state: - errors.append("Safe resume requires matching AMP scaler state.") - lines: list[str] = [] - if errors: - lines.append("[BLOCK] Resume is not safe with the current settings.") - elif report.warnings: - lines.append("[WARN] Resume is possible, but settings changed.") - else: - lines.append("[OK] Checkpoint can resume with the current settings.") - lines.extend(f"[OK] {line}" for line in report.info) - lines.extend(f"[WARN] {line}" for line in report.warnings) - lines.extend(f"[BLOCK] {line}" for line in errors) - if not training_config.require_compatible_resume and not errors: - lines.append("[INFO] Safe resume is off. Compatible weights will load; incompatible optimizer state may be skipped.") - self.resume_training_preview.setText("\n".join(lines)) - except Exception as exc: - self.resume_training_preview.setText(f"[BLOCK] Could not check resume compatibility:\n{exc}") - - def preview_fine_tune_compatibility(self) -> None: - """Preview whether the selected checkpoint can be used for fine-tuning.""" - - stage_ok, stage_message = self._fine_tune_dataset_stage_status() - if not stage_ok: - self.fine_tune_preview.setText(f"[BLOCK] {stage_message}") - return - base_path = Path(self.fine_tune_checkpoint.text()) if self.fine_tune_checkpoint.text().strip() else None - if base_path is None: - self.fine_tune_preview.setText("[BLOCK] Choose a base checkpoint for fine-tuning.") - return - if not base_path.exists(): - self.fine_tune_preview.setText(f"[BLOCK] Fine-tune base checkpoint does not exist:\n{base_path}") - return - try: - vocab_size = self._current_training_vocab_size(Path(self.train_data_dir.text())) - if vocab_size <= 0: - self.fine_tune_preview.setText("[BLOCK] Could not determine current dataset tokenizer vocabulary size.") - return - model_config = self._current_model_config(vocab_size=vocab_size) - training_config = self._current_training_config() - model_config.validate() - report = check_resume_compatibility(base_path, model_config, training_config) - lines: list[str] = [] - if report.errors: - lines.append("[BLOCK] Base checkpoint cannot be fine-tuned with the current model/dataset settings.") - else: - lines.append("[OK] Base checkpoint weights can be used for fine-tuning.") - lines.append(f"[OK] {stage_message}" if stage_ok else f"[BLOCK] {stage_message}") - lines.extend(self._fine_tune_lineage_advice(base_path)) - lines.extend(f"[OK] {line}" for line in report.info) - behavior_warnings = [ - warning for warning in report.warnings - if not warning.startswith("Optimizer changed:") and not warning.startswith("LR scheduler changed:") - ] - lines.extend(f"[WARN] {line}" for line in behavior_warnings) - lines.extend(f"[BLOCK] {line}" for line in report.errors) - checkpoint_vocab = self._checkpoint_vocab_size(base_path) - if checkpoint_vocab and checkpoint_vocab != vocab_size: - lines.append(f"[FIX] {self._tokenizer_mismatch_help(checkpoint_vocab, vocab_size)}") - if not report.errors: - lines.append("[INFO] Fine-tuning starts fresh optimizer, scheduler, and scaler state.") - self.fine_tune_preview.setText("\n".join(lines)) - except Exception as exc: - self.fine_tune_preview.setText(f"[BLOCK] Could not check fine-tune compatibility:\n{exc}") - - def _fine_tune_lineage_advice(self, base_path: Path) -> list[str]: - """Return guidance about the selected fine-tune base checkpoint. - - Args: - base_path: Selected checkpoint path. - - Returns: - Lines for the fine-tune compatibility report. - """ - - lines: list[str] = [] - try: - output_dir = self._fine_tune_output_path().resolve() - base_resolved = base_path.resolve() - if output_dir == base_resolved.parent or output_dir in base_resolved.parents: - return [ - "[BLOCK] Selected base checkpoint is inside the current fine-tune output folder.", - "[FIX] Choose the original pretrained model or a completed earlier fine-tune from another folder.", - ] - except OSError: - pass - lineage_path = base_path.parent / "model_lineage.json" - summary_path = base_path.parent / "training_summary.json" - lineage = read_json(lineage_path, default={}) or {} - summary = read_json(summary_path, default={}) or {} - training_mode = str(lineage.get("training_mode") or (summary.get("training_config") or {}).get("training_mode") or "") - stage = str((summary.get("model_lineage") or lineage).get("fine_tune_stage") or "") - if training_mode == "fine_tune": - stage_text = f" ({stage})" if stage else "" - lines.append(f"[INFO] Selected base is a previous fine-tuned checkpoint{stage_text}.") - lines.append("[INFO] This is correct for cumulative tuning, such as conversation -> instruction -> code.") - elif training_mode == "pretrain": - lines.append("[OK] Selected base is the pretrained model checkpoint.") - lines.append("[INFO] This is correct when starting a new independent fine-tune branch.") - else: - lines.append("[INFO] Could not read model lineage; compatibility check will still validate tensor shapes.") - project_base = self.current_project_file.parent / "models" / "final_model.pt" if self.current_project_file else None - if project_base and project_base.exists(): - try: - if base_path.resolve() != project_base.resolve() and training_mode != "fine_tune": - lines.append(f"[HINT] Project pretrained model is: {project_base}") - except OSError: - pass - return lines - - def _training_history_path(self) -> Path: - """Return the training history path for the selected model folder. - - Returns: - Path to ``training_history.json``. - """ - - output_dir = getattr(self, "active_training_output_dir", None) - if output_dir is None: - output_dir = Path(self.model_dir.text()) - return Path(output_dir) / "training_history.json" - - def _load_training_history(self) -> list[dict[str, Any]]: - """Load training run history. - - Returns: - List of training run entries. - """ - - path = self._training_history_path() - if not path.exists(): - return [] - try: - data = json.loads(path.read_text(encoding="utf-8")) - return data if isinstance(data, list) else [] - except Exception: - return [] - - def refresh_model_estimate(self) -> None: - """Refresh model size, rough VRAM, and run history widgets.""" - - model_config = self._current_model_config() - # Same reasoning as preview_resume_compatibility: this is the - # AI/Training tab's shared "Model Estimate" card, not the - # Fine-Tuning tab's; pass an explicit override rather than - # inheriting the Fine-Tuning tab's mode combo by fallback. - training_config = self._current_training_config(training_mode="pretrain") - data_dir = Path(self.train_data_dir.text()) - train_tokens = max(model_config.context_length * training_config.batch_size, 1) - try: - summary_path = data_dir / "dataset_summary.json" - if summary_path.exists(): - summary = json.loads(summary_path.read_text(encoding="utf-8")) - vocab_size = int(summary.get("tokenizer_vocab_size", 0) or 0) - train_tokens = int(summary.get("train_token_count", summary.get("token_count", train_tokens)) or train_tokens) - if vocab_size > 0: - model_config.vocab_size = vocab_size - elif (data_dir / "tokenizer.json").exists(): - tokenizer_data = json.loads((data_dir / "tokenizer.json").read_text(encoding="utf-8")) - vocab = tokenizer_data.get("model", {}).get("vocab", {}) - if vocab: - model_config.vocab_size = len(vocab) - except Exception as exc: - self.training_log.append(f"[WARN] Could not refresh dataset-based estimate: {exc}") - estimate = estimate_training_resources(model_config, training_config, train_tokens) - self.last_training_estimate = estimate - self._update_model_estimate_chips(estimate, model_config, training_config, train_tokens) - self.history_metric.setText(f"Runs: {len(self._load_training_history())}") - self.training_log.append( - "Model estimate refreshed: " - f"{int(estimate['parameters']):,} params, " - f"checkpoint {format_bytes(float(estimate['checkpoint_bytes']))}, " - f"VRAM {format_bytes(float(estimate['vram_bytes']))}." - ) - - def _append_training_history(self, result: Any) -> None: - """Persist a training run entry to ``training_history.json``. - - Args: - result: Training result object. - """ - - history_path = self._training_history_path() - history_path.parent.mkdir(parents=True, exist_ok=True) - history = self._load_training_history() - summary = {} - try: - if Path(result.summary_path).exists(): - summary = json.loads(Path(result.summary_path).read_text(encoding="utf-8")) - except Exception: - summary = {} - estimate = getattr(self, "last_training_estimate", {}) or {} - entry = { - "completed_at": datetime.now().isoformat(timespec="seconds"), - "checkpoint_path": str(result.checkpoint_path), - "summary_path": str(result.summary_path), - "stopped": bool(getattr(result, "stopped", False)), - "final_train_loss": result.final_train_loss, - "final_val_loss": result.final_val_loss, - "best_val_loss": summary.get("best_val_loss"), - "recommended_checkpoint_path": summary.get("recommended_checkpoint_path"), - "best_checkpoint_path": summary.get("best_checkpoint_path"), - "dataset_dir": self.train_data_dir.text(), - "dataset_version": (summary.get("model_lineage") or {}).get("dataset_version"), - "training_run_id": summary.get("training_run_id"), - "parameters": estimate.get("parameters") or summary.get("parameters"), - "model_config": summary.get("model_config"), - "training_config": summary.get("training_config"), - } - history.append(entry) - history_path.write_text(json.dumps(history[-200:], indent=2), encoding="utf-8") - self.history_metric.setText(f"Runs: {len(history[-200:])}") - (self.active_training_log or self.training_log).append(f"Training history updated: {history_path}") - - def _run_training_preflight(self, model_config: ModelConfig, training_config: TrainingConfig) -> bool: - """Run pre-training checklist and disk-space guard. - - Args: - model_config: Selected model architecture. - training_config: Selected training settings. - - Returns: - True when training may continue. - """ - - log = self.active_training_log or self.training_log - data_dir = Path(self.train_data_dir.text()) - output_dir = training_config.output_dir - errors: list[str] = [] - warnings: list[str] = [] - info: list[str] = [] - resettable_errors: list[str] = [] - missing: list[str] = [] - if not (data_dir / "tokenizer.json").exists(): - missing.append("tokenizer.json") - has_npy_tokens = (data_dir / "train_tokens.npy").exists() and (data_dir / "val_tokens.npy").exists() - has_json_tokens = (data_dir / "train_tokens.json").exists() and (data_dir / "val_tokens.json").exists() - if not has_npy_tokens and not has_json_tokens: - missing.append("train_tokens.(npy/json), val_tokens.(npy/json)") - if not data_dir.exists(): - errors.append(f"Dataset folder does not exist: {data_dir}") - elif missing: - errors.append(f"Dataset is not prepared. Missing: {', '.join(missing)}") - else: - info.append("Dataset artifacts found.") - - vocab_size = 0 - train_tokens = 0 - val_tokens = 0 - summary = {} - try: - summary_path = data_dir / "dataset_summary.json" - if summary_path.exists(): - summary = json.loads(summary_path.read_text(encoding="utf-8")) - vocab_size = int(summary.get("tokenizer_vocab_size", 0) or 0) - if vocab_size <= 0: - tokenizer_data = json.loads((data_dir / "tokenizer.json").read_text(encoding="utf-8")) - vocab_size = len(tokenizer_data.get("model", {}).get("vocab", {})) - train_tokens = int(summary.get("train_token_count", 0) or 0) - val_tokens = int(summary.get("val_token_count", 0) or 0) - if train_tokens <= 0 and (data_dir / "train_tokens.json").exists(): - train_tokens = len(json.loads((data_dir / "train_tokens.json").read_text(encoding="utf-8"))) - if val_tokens <= 0 and (data_dir / "val_tokens.json").exists(): - val_tokens = len(json.loads((data_dir / "val_tokens.json").read_text(encoding="utf-8"))) - except Exception as exc: - warnings.append(f"Could not fully inspect dataset metadata: {exc}") - - if vocab_size > 0: - model_config.vocab_size = vocab_size - info.append(f"Tokenizer vocab: {vocab_size:,}.") - elif not missing: - errors.append("Could not determine tokenizer vocabulary size.") - if train_tokens and train_tokens <= model_config.context_length: - errors.append("Training token count must be larger than context length.") - elif train_tokens: - info.append(f"Training tokens: {train_tokens:,}; validation tokens: {val_tokens:,}.") - if train_tokens < 50_000: - warnings.append("Training token count is very small; expect smoke-test quality.") - - try: - model_config.validate() - except Exception as exc: - errors.append(f"Model architecture is invalid: {exc}") - try: - training_config.validate() - except Exception as exc: - errors.append(f"Training options are invalid: {exc}") - if model_config.attention_backend == "sdpa": - if hasattr(torch.nn.functional, "scaled_dot_product_attention"): - if training_config.device == "cuda" and torch.cuda.is_available(): - flash_enabled = bool(getattr(torch.backends.cuda, "flash_sdp_enabled", lambda: False)()) - info.append("Attention backend: SDPA selected; Flash Attention may be used by PyTorch." if flash_enabled else "Attention backend: SDPA selected; CUDA flash kernel is not enabled.") - else: - info.append("Attention backend: SDPA selected; CPU/backend fallback will be used if needed.") - else: - warnings.append("SDPA attention selected, but this PyTorch build does not expose scaled_dot_product_attention.") - else: - warnings.append("Manual attention backend selected. This is useful for debugging but can be slower.") - if training_config.peft_method == "lora": - info.append( - "PEFT: LoRA adapters enabled. Intermediate checkpoints will save adapter weights; final_model.pt will be merged." - ) - - if training_config.device == "cuda" and not torch.cuda.is_available(): - errors.append("CUDA is selected, but PyTorch cannot use CUDA on this machine.") - elif training_config.device == "cuda": - info.append(f"CUDA ready: {torch.cuda.get_device_name(0)}.") - if training_config.data_loader_workers > 0: - info.append(f"CPU-assisted batch loading enabled with {training_config.data_loader_workers} worker(s).") - else: - warnings.append("CPU training is selected. This can be very slow.") - if sys.platform.startswith("win") and training_config.data_loader_workers > 4: - warnings.append("High CPU worker counts can duplicate dataset memory on Windows. Start with 2-4 workers and increase carefully.") - - active_resume_path: Optional[Path] = None - resume_path = training_config.resume_from_checkpoint if training_config.resume else None - if resume_path and not Path(resume_path).exists(): - errors.append(f"Selected resume checkpoint does not exist: {resume_path}") - elif training_config.resume: - if resume_path is None: - resume_path = latest_checkpoint(output_dir / "checkpoints") - if resume_path is None: - info.append("Resume latest is enabled, but no checkpoint exists yet.") - else: - active_resume_path = Path(resume_path) - try: - compatibility = check_resume_compatibility(active_resume_path, model_config, training_config) - info.extend(compatibility.info) - warnings.extend(compatibility.warnings) - errors.extend(compatibility.errors) - resettable_errors.extend(compatibility.errors) - if training_config.require_compatible_resume: - if not compatibility.can_load_optimizer_state: - message = "Safe resume requires matching optimizer state." - errors.append(message) - resettable_errors.append(message) - if not compatibility.can_load_scheduler_state: - message = "Safe resume requires matching scheduler state." - errors.append(message) - resettable_errors.append(message) - if not compatibility.can_load_scaler_state: - message = "Safe resume requires matching AMP scaler state." - errors.append(message) - resettable_errors.append(message) - except Exception as exc: - errors.append(f"Could not inspect resume checkpoint: {exc}") - if training_config.training_mode == "fine_tune" and active_resume_path is None: - base_path = training_config.fine_tune_from_checkpoint - if base_path is None: - errors.append("Fine-tune mode requires a base checkpoint.") - elif not Path(base_path).exists(): - errors.append(f"Fine-tune base checkpoint does not exist: {base_path}") - else: - try: - compatibility = check_resume_compatibility(Path(base_path), model_config, training_config) - info.append(f"Fine-tune base checkpoint: {Path(base_path).name}.") - warnings.extend( - warning for warning in compatibility.warnings - if not warning.startswith("Optimizer changed:") and not warning.startswith("LR scheduler changed:") - ) - errors.extend(compatibility.errors) - checkpoint_vocab = self._checkpoint_vocab_size(Path(base_path)) - if checkpoint_vocab and checkpoint_vocab != vocab_size: - errors.append(self._tokenizer_mismatch_help(checkpoint_vocab, vocab_size)) - if not compatibility.errors: - info.append("Fine-tune base weights are compatible. Optimizer state will start fresh.") - except Exception as exc: - errors.append(f"Could not inspect fine-tune base checkpoint: {exc}") - elif training_config.training_mode == "pretrain" and active_resume_path is None: - info.append("A fresh pretraining run will start from random weights.") - elif training_config.training_mode == "fine_tune" and active_resume_path is not None: - info.append("Existing run checkpoint found; training will resume that run instead of reloading the fine-tune base.") - - output_dir.mkdir(parents=True, exist_ok=True) - estimate = estimate_training_resources(model_config, training_config, train_tokens) - self.last_training_estimate = estimate - self._update_model_estimate_chips(estimate, model_config, training_config, train_tokens) - params = int(estimate["parameters"]) - checkpoint_bytes = float(estimate["checkpoint_bytes"]) - checkpoint_count = int(estimate["checkpoint_count"]) - estimated_storage = float(estimate["estimated_storage"]) - estimated_vram = float(estimate["vram_bytes"]) - free_bytes = shutil.disk_usage(output_dir).free - info.append(f"Estimated parameters: {params:,}.") - info.append(f"Estimated checkpoint size: {format_bytes(checkpoint_bytes)}.") - info.append(f"Estimated training VRAM: {format_bytes(estimated_vram)}.") - info.append(f"Estimated training storage need: {format_bytes(estimated_storage)}.") - info.append(f"Free space on model drive: {format_bytes(free_bytes)}.") - if training_config.device == "cuda" and torch.cuda.is_available(): - free_vram, total_vram = torch.cuda.mem_get_info() - info.append(f"GPU free/total VRAM: {format_bytes(free_vram)} / {format_bytes(total_vram)}.") - if estimated_vram > free_vram * 0.9: - warnings.append("Estimated VRAM is close to or above currently free GPU memory.") - if free_bytes < estimated_storage * 1.25: - errors.append("Not enough free disk space for estimated checkpoints and final model.") - elif free_bytes < estimated_storage * 2: - warnings.append("Free disk space is close to the estimated training storage need.") - if checkpoint_count > 50: - warnings.append("Save interval may create many checkpoints. Increase Save every or clean old checkpoints.") - - log.clear() - log.append("Training checklist") - for line in info: - log.append(f"[OK] {line}") - for line in warnings: - log.append(f"[WARN] {line}") - for line in errors: - log.append(f"[ERROR] {line}") - - if errors: - hard_errors = [error for error in errors if error not in resettable_errors] - if active_resume_path is not None and resettable_errors and not hard_errors: - message = ( - "The existing checkpoint was created with different model settings, so it cannot be resumed.\n\n" - "This is expected if you intentionally changed architecture, block style, tokenizer, " - "context length, attention layout, or other checkpoint-shaped settings.\n\n" - "You can start a fresh training run with the current settings. This will delete old " - "checkpoints and training summaries in the selected model output folder.\n\n" - f"Model folder:\n{output_dir}\n\n" - "Continue and start from scratch?" - ) - choice = QMessageBox.question( - self, - "Start From Scratch?", - message, - QMessageBox.Yes | QMessageBox.No, - QMessageBox.No, - ) - if choice == QMessageBox.Yes: - removed = self._clear_training_run_artifacts(output_dir) - training_config.resume = False - training_config.resume_from_checkpoint = None - self.resume_training.setChecked(False) - self.resume_checkpoint.clear() - log.append("") - log.append("Starting fresh run with current settings.") - for path in removed: - log.append(f"Removed old training artifact: {path}") - LOGGER.warning( - "User chose to discard incompatible resume checkpoint %s and start fresh in %s", - active_resume_path, - output_dir, - ) - if training_config.training_mode == "fine_tune": - base_path = training_config.fine_tune_from_checkpoint - if base_path is None or not Path(base_path).exists(): - log.append("[ERROR] Fine-tune mode requires an existing base checkpoint after reset.") - QMessageBox.warning(self, "Training blocked", "Fine-tune mode still needs a valid base checkpoint.") - return False - compatibility = check_resume_compatibility(Path(base_path), model_config, training_config) - if compatibility.errors: - for line in compatibility.errors: - log.append(f"[ERROR] {line}") - QMessageBox.warning(self, "Training blocked", "The base checkpoint is still incompatible with current settings.") - return False - log.append("[OK] Old run cleared; fine-tune base checkpoint is compatible.") - else: - log.append("[OK] Old run cleared; pretraining will start from random weights.") - self.project_state.setText("Training reset") - self.train_status.setText("Training: starting fresh") - return True - LOGGER.error("Training blocked by preflight checklist.") - for line in info: - LOGGER.info("Training preflight OK: %s", line) - for line in warnings: - LOGGER.warning("Training preflight warning: %s", line) - for line in errors: - LOGGER.error("Training preflight error: %s", line) - self.project_state.setText("Training blocked") - self.train_status.setText("Training: blocked") - QMessageBox.warning(self, "Training blocked", "Fix the checklist errors before starting training.") - return False - existing_artifacts = self._training_run_artifacts(output_dir) - if ( - training_config.training_mode == "pretrain" - and active_resume_path is None - and existing_artifacts - ): - artifact_text = "\n".join(f"- {path.name}" for path in existing_artifacts) - message = ( - "This model folder already contains training artifacts from a previous run.\n\n" - "If you changed architecture or low-memory settings and want a clean start, " - "the old checkpoints should be removed first.\n\n" - f"Model folder:\n{output_dir}\n\n" - f"Artifacts found:\n{artifact_text}\n\n" - "Delete these artifacts and start from scratch?" - ) - choice = QMessageBox.question( - self, - "Clean Previous Run?", - message, - QMessageBox.Yes | QMessageBox.No, - QMessageBox.No, - ) - if choice != QMessageBox.Yes: - self.project_state.setText("Training cancelled") - self.train_status.setText("Training: idle") - log.append("Training cancelled. Previous run artifacts were kept.") - return False - removed = self._clear_training_run_artifacts(output_dir) - training_config.resume = False - training_config.resume_from_checkpoint = None - self.resume_training.setChecked(False) - self.resume_checkpoint.clear() - log.append("") - log.append("Previous run artifacts removed. Training will start from scratch with current settings.") - for path in removed: - log.append(f"Removed old training artifact: {path}") - LOGGER.warning("User cleaned previous training artifacts in %s before starting from scratch.", output_dir) - if warnings: - message = "Training checklist has warnings. Continue anyway?\n\n" + "\n".join(f"- {warning}" for warning in warnings[:8]) - choice = QMessageBox.question(self, "Training warnings", message, QMessageBox.Yes | QMessageBox.No, QMessageBox.No) - if choice != QMessageBox.Yes: - self.project_state.setText("Training cancelled") - self.train_status.setText("Training: idle") - return False - return True - - def start_training(self) -> None: - """Collect training options and start model training.""" - - launch_target = self._training_launch_target_value() - if launch_target == "runpod": - self.launch_runpod_worker_for_current_training() - return - if launch_target == "remote": - self.publish_remote_training_job() - return - self.active_training_log = self.training_log - self.active_training_progress = self.training_progress - self.active_training_final_button_text = "Start Training" - resume_path = Path(self.resume_checkpoint.text()) if self.resume_checkpoint.text().strip() else None - dataset_dir = Path(self.train_data_dir.text()) - vocab_size = self._current_training_vocab_size(dataset_dir) - if vocab_size <= 0: - QMessageBox.warning(self, "Training blocked", "Could not determine tokenizer vocabulary size. Prepare the dataset first.") - return - model_config = self._current_model_config(vocab_size=vocab_size) - training_config = self._current_training_config(resume_path, training_mode="pretrain") - if not self._run_training_preflight(model_config, training_config): - return - self.active_training_output_dir = training_config.output_dir - self._init_telemetry_store(training_config.output_dir) - self.training_log.append("") - self.training_progress.setValue(0) - self.training_epoch_metric.setText("Epoch: -") - self.training_step_metric.setText("Step: -") - self.training_loss_metric.setText("Train loss: -") - self.training_val_metric.setText("Val loss: -") - self.training_health_metric.setText("Health: -") - self.training_health_points = [] - self.training_lr_metric.setText("LR: -") - self.training_speed_metric.setText("Speed: -") - self.training_grad_metric.setText("Grad: -") - self.training_vram_metric.setText("VRAM: -") - self.training_eta_metric.setText("ETA: -") - self.loss_chart.clear() - self.optimization_chart.clear() - self.stability_chart.clear() - self.throughput_chart.clear() - self.memory_chart.clear() - self.live_prediction_chart.update_distribution(0, None) - self.live_attention_chart.update_heatmap(0, None) - self.live_activation_chart.update_histogram(0, None) - self.live_gradient_chart.update_flow(self.n_layer.value(), None, 0) - self.live_progress.setValue(0) - self.live_epoch_metric.setText("Epoch: -") - self.live_step_metric.setText("Step: -") - self.live_tokens_metric.setText("Tokens/sec: -") - self.live_loss_metric.setText("Loss: -") - self.live_lr_metric.setText("LR: -") - self.live_data_metric.setText("Data: -") - self.live_sample_text.setText("Training text: -") - self.live_flow.set_state(self.n_layer.value(), self.n_head.value(), 0, None) - self._set_meter(self.live_cpu_bar, "CPU", self._system_cpu_value()) - self._set_meter(self.live_gpu_bar, "GPU memory", None) - self._set_meter(self.live_vram_bar, "VRAM reserved", None) - self._set_meter(self.live_ram_bar, "System RAM", self._system_ram_value()) - self.live_worker_status.setText(f"CPU workers: {self.data_loader_workers.value()}") - self.training_log.append("Training started...") - self.project_state.setText("Training") - self.train_status.setText("Training: running") - self._run_task( - run_training_job, - (dataset_dir, model_config, training_config), - self._training_finished, - self.training_log, - self.training_progress, - with_progress=True, - button=self.train_button, - stop_button=self.stop_training_button, - busy_text="Training", - task_kind="training", - ) - - def start_fine_tuning(self) -> None: - """Collect fine-tuning options and start adaptation training.""" - - fine_tune_launch = self._fine_tune_launch_target_value() - if fine_tune_launch in {"remote", "runpod"}: - stage_ok, stage_message = self._fine_tune_dataset_stage_status() - self.refresh_fine_tune_workflow() - if not stage_ok: - self.fine_tune_log.append(stage_message) - QMessageBox.warning(self, "Fine-tune blocked", stage_message) - return - if fine_tune_launch == "runpod": - self.launch_runpod_worker_for_current_training(training_mode="fine_tune", stage=self._training_stage_value()) - self.fine_tune_log.append("RunPod fine-tune job launched. Watch Job Manager for worker assignment and progress.") - else: - self.publish_remote_training_job(training_mode="fine_tune", stage=self._training_stage_value()) - self.fine_tune_log.append("Remote fine-tune job queued. Watch Job Manager for worker assignment and progress.") - return - self.active_training_log = self.fine_tune_log - self.active_training_progress = self.fine_tune_progress - self.active_training_final_button_text = "Start Fine-Tune" - stage_ok, stage_message = self._fine_tune_dataset_stage_status() - self.refresh_fine_tune_workflow() - if not stage_ok: - self.fine_tune_log.append(stage_message) - QMessageBox.warning(self, "Fine-tune blocked", stage_message) - return - resume_path = Path(self.resume_checkpoint.text()) if self.resume_checkpoint.text().strip() else None - dataset_dir = Path(self.train_data_dir.text()) - vocab_size = self._current_training_vocab_size(dataset_dir) - if vocab_size <= 0: - QMessageBox.warning(self, "Fine-tune blocked", "Could not determine tokenizer vocabulary size. Prepare the fine-tuning dataset first.") - return - model_config = self._current_model_config(vocab_size=vocab_size) - training_config = self._current_training_config(resume_path, training_mode="fine_tune") - if not self._run_training_preflight(model_config, training_config): - return - self.active_training_output_dir = training_config.output_dir - self._prepare_fine_tune_run_folder(training_config) - self._init_telemetry_store(training_config.output_dir) - self.fine_tune_log.append("") - self.fine_tune_progress.setValue(0) - self.training_progress.setValue(0) - self.fine_tune_eta_metric.setText("ETA: -") - self.fine_tune_epoch_metric.setText("Epoch: -") - self.fine_tune_step_metric.setText("Step: -") - self.fine_tune_loss_metric.setText("Train loss: -") - self.fine_tune_val_metric.setText("Val loss: -") - self.fine_tune_lr_metric.setText("LR: -") - self.fine_tune_speed_metric.setText("Speed: -") - self.fine_tune_grad_metric.setText("Grad: -") - self.training_epoch_metric.setText("Epoch: -") - self.training_step_metric.setText("Step: -") - self.training_loss_metric.setText("Train loss: -") - self.training_val_metric.setText("Val loss: -") - self.training_health_metric.setText("Health: -") - self.training_health_points = [] - self.training_lr_metric.setText("LR: -") - self.training_speed_metric.setText("Speed: -") - self.training_grad_metric.setText("Grad: -") - self.training_vram_metric.setText("VRAM: -") - self.training_eta_metric.setText("ETA: -") - self.loss_chart.clear() - self.optimization_chart.clear() - self.stability_chart.clear() - self.throughput_chart.clear() - self.memory_chart.clear() - self.live_progress.setValue(0) - self.live_sample_text.setText("Training text: -") - self.live_flow.set_state(self.n_layer.value(), self.n_head.value(), 0, None) - self.fine_tune_log.append("Fine-tuning started...") - self.project_state.setText("Fine-tuning") - self.train_status.setText("Training: fine-tuning") - self._run_task( - run_fine_tuning_job, - (dataset_dir, model_config, training_config, self._training_stage_value()), - self._training_finished, - self.fine_tune_log, - self.fine_tune_progress, - with_progress=True, - button=self.fine_tune_button, - stop_button=self.stop_fine_tune_button, - busy_text="Fine-tuning", - task_kind="fine_tune", - ) - - @Slot(object) - def _training_finished(self, result: Any) -> None: - """Update UI after training finishes. - - Args: - result: Training result. - """ - - log = self.active_training_log or self.training_log - progress = self.active_training_progress or self.training_progress - progress.setValue(100) - if progress is not self.training_progress: - self.training_progress.setValue(100) - if hasattr(self, "live_progress"): - self.live_progress.setValue(100) - log.append(f"Saved model: {result.checkpoint_path}") - log.append(f"Final train loss: {result.final_train_loss:.4f}") - if result.final_val_loss is not None: - log.append(f"Final validation loss: {result.final_val_loss:.4f}") - training_summary: dict[str, Any] = {} - try: - training_summary = json.loads(Path(result.summary_path).read_text(encoding="utf-8")) - except Exception: - training_summary = {} - best_checkpoint = str(training_summary.get("recommended_checkpoint_path") or "") - best_val_loss = training_summary.get("best_val_loss") - if best_checkpoint: - if best_val_loss is not None: - log.append(f"Recommended checkpoint: {best_checkpoint} (best validation loss {float(best_val_loss):.4f})") - else: - log.append(f"Recommended checkpoint: {best_checkpoint}") - output_dir = self.active_training_output_dir or Path(result.checkpoint_path).parent - stage_key = self.active_task_kind if self.active_task_kind in {"training", "fine_tune"} else "training" - self.export_model_dir.setText(str(output_dir)) - try: - if stage_key != "fine_tune" and Path(output_dir).resolve() == Path(self.model_dir.text()).resolve(): - self.fine_tune_checkpoint.setText(str(result.checkpoint_path)) - except OSError: - pass - if getattr(result, "stopped", False): - self.project_state.setText("Training stopped") - self.train_status.setText("Training: stopped, checkpoint saved") - log.append("Training stopped safely. Resume from this checkpoint or the latest checkpoint.") - else: - self.project_state.setText("Training complete") - self.train_status.setText(f"Training: loss {result.final_train_loss:.4f}") - title = "Fine-tuning complete" if stage_key == "fine_tune" else "Model training complete" - if getattr(result, "stopped", False): - title = "Fine-tuning stopped" if stage_key == "fine_tune" else "Model training stopped" - completion_lines = [ - f"Checkpoint: {result.checkpoint_path}", - f"Summary: {result.summary_path}", - f"Final train loss: {result.final_train_loss:.4f}", - ] - if result.final_val_loss is not None: - completion_lines.append(f"Final validation loss: {result.final_val_loss:.4f}") - if best_checkpoint: - completion_lines.append(f"Recommended checkpoint: {best_checkpoint}") - if best_val_loss is not None: - completion_lines.append(f"Best validation loss: {float(best_val_loss):.4f}") - completion_lines.append(f"Output: {output_dir}") - self._notify_complete(stage_key, title, completion_lines) - self._append_training_history(result) - self._clear_button_busy(self.active_training_final_button_text) - self.active_training_log = None - self.active_training_progress = None - self.active_training_output_dir = None - - def run_benchmark(self) -> None: - """Run benchmark prompts against the current trained model.""" - - prompts = normalize_prompts(self.benchmark_prompts.toPlainText()) - self.benchmark_log.append(f"Running benchmark with {len(prompts)} prompt(s)...") - self.benchmark_progress.setValue(0) - self.project_state.setText("Benchmarking") - self._run_task( - evaluate_checkpoint, - ( - Path(self.model_dir.text()), - prompts, - None, - self.benchmark_tokens.value(), - self.benchmark_temperature.value(), - 50, - self.device.currentText(), - self.benchmark_kv_cache.isChecked(), - ), - self._benchmark_finished, - self.benchmark_log, - self.benchmark_progress, - with_progress=True, - button=self.run_benchmark_button, - stop_button=self.stop_benchmark_button, - busy_text="Benchmarking", - ) - - @Slot(object) - def _benchmark_finished(self, result: Any) -> None: - """Update UI after benchmark prompts finish. - - Args: - result: Benchmark result object. - """ - - self.benchmark_progress.setRange(0, 100) - self.benchmark_progress.setValue(100) - self.benchmark_log.append( - f"Benchmark complete: {result.prompt_count} prompt(s), {result.total_seconds:.2f}s, " - f"{result.total_generated_tokens} generated token(s), {result.tokens_per_second:.2f} tok/s." - ) - self.benchmark_log.append(f"Benchmark saved: {result.output_path}") - self.project_state.setText("Benchmark complete") - self._clear_button_busy("Run Benchmark") - - def toggle_llm_model(self) -> None: - """Load or unload the selected chat model depending on current state.""" - - if self.chat_session is not None: - self.unload_llm_model() - return - self.load_llm_model() - - def load_llm_model(self) -> None: - """Load a selected model backend for chat testing.""" - - backend = self._chat_backend_value() - path_text = self.microgpt_chat_path.text().strip() if backend == "microgpt" else self.gguf_path.text().strip() - if not path_text: - required = "MicroGPT model folder or checkpoint" if backend == "microgpt" else "GGUF model file" - QMessageBox.information(self, "Model required", f"Choose a {required} first.") - return - model_path = Path(path_text) - self.chat_progress.setValue(0) - self._render_chat_markdown("**Loading model...**") - self.chat_stats.setText("Loading model...") - self.project_state.setText("Loading chat model") - self.chat_status.setText("Chat: loading model") - loader = load_microgpt_chat_session if backend == "microgpt" else load_llama_chat_session - args = ( - (model_path, self.device.currentText()) - if backend == "microgpt" - else (model_path, self.llama_context.value(), self.llama_threads.value(), self.llama_gpu_layers.value()) - ) - self._run_task( - loader, - args, - self._llm_loaded, - self.chat_event_log, - self.chat_progress, - button=self.load_llm_button, - busy_text="Loading Model", - task_kind="chat", - ) - - @Slot(object) - def _llm_loaded(self, session: Any) -> None: - """Store a loaded GGUF chat session. - - Args: - session: Loaded ``LlamaChatSession``. - """ - - self.chat_session = session - self._clear_chat_messages() - self.chat_markdown = "" - self._add_chat_message( - "assistant", - f"Loaded model: `{session.model_path.name}`\n\n{session.runtime_summary}\n\nSend a message to begin.", - ) - self.chat_progress.setValue(100) - self.chat_stats.setText(session.runtime_summary) - self.project_state.setText("Chat model loaded") - self.chat_status.setText(f"Chat: {session.runtime_summary}") - self._clear_button_busy("Unload") - self._tip(self.load_llm_button, "Unload the currently loaded model from memory.") - - def unload_llm_model(self) -> None: - """Unload the active chat model and clear chat state.""" - - if self.thread is not None: - QMessageBox.information(self, "Task running", "Please wait for the current task to finish.") - return - if self.chat_session is not None and hasattr(self.chat_session, "reset"): - self.chat_session.reset() - self.chat_session = None - self._clear_chat_messages() - self.chat_markdown = "" - self._add_chat_message("assistant", "Model unloaded.\n\nLoad a model to start testing.") - self.chat_progress.setRange(0, 100) - self.chat_progress.setValue(0) - self.chat_stats.setText("Idle") - self.project_state.setText("Ready") - self.chat_status.setText("Chat: no model loaded") - self.load_llm_button.setText("Load Model") - self._update_chat_backend_controls() - - def send_chat_message(self) -> None: - """Send a prompt to the loaded chat model.""" - - if self.chat_session is None: - QMessageBox.information(self, "Load model", "Load a model before sending a message.") - return - prompt = self.chat_input.toPlainText().strip() - if not prompt: - return - self.pending_user_message = prompt - self.chat_input.clear() - self._add_chat_message("user", prompt, resend_prompt=prompt) - self.chat_stream_reply = "" - self._add_chat_message("assistant", "_Thinking..._", resend_prompt=prompt) - self.chat_progress.setRange(0, 0) - self.chat_stats.setText("Thinking...") - self.project_state.setText("Generating") - self.chat_status.setText("Chat: generating reply") - streamer = stream_microgpt_chat_reply if self._chat_backend_value() == "microgpt" else stream_chat_reply - self._run_task( - streamer, - ( - self.chat_session, - prompt, - self.system_prompt.toPlainText(), - self.chat_max_tokens.value(), - self.chat_temperature.value(), - self.chat_top_p.value(), - self.chat_repeat_penalty.value(), - self.reasoning_effort.currentText(), - self.thinking_enabled.isChecked(), - ), - self._chat_reply_finished, - self.chat_event_log, - self.chat_progress, - with_progress=True, - button=self.send_chat_button, - stop_button=self.stop_chat_button, - busy_text="Thinking", - ) - - @Slot(object) - def _chat_reply_finished(self, reply: Any) -> None: - """Render the model reply. - - Args: - reply: Assistant reply text and metrics. - """ - - result = reply if isinstance(reply, dict) else {"reply": str(reply)} - text = str(result.get("reply", "")).strip() - if text: - self.chat_stream_reply = text - else: - self.chat_stream_reply = self.chat_stream_reply or "_No reply returned._" - self._render_chat_markdown(self.chat_stream_reply) - self.chat_progress.setRange(0, 100) - self.chat_progress.setValue(100) - self._set_chat_stats( - float(result.get("elapsed_seconds", 0.0)), - int(result.get("token_count", 0)), - float(result.get("tokens_per_second", 0.0)), - ) - self.project_state.setText("Ready") - self.chat_status.setText("Chat: ready") - self._clear_button_busy("Send") - - def reset_chat(self) -> None: - """Clear the chat transcript and model conversation memory.""" - - if self.chat_session is not None: - self.chat_session.reset() - self._clear_chat_messages() - self.chat_markdown = "" - self.chat_stream_prefix = "" - self.chat_stream_reply = "" - self._add_chat_message("assistant", "Chat reset.") - self.chat_stats.setText("Idle") - self.chat_status.setText("Chat: ready") - - def _append_chat_markdown(self, role: str, content: str) -> None: - """Append one rendered chat message. - - Args: - role: Display role heading. - content: Markdown content. - """ - - block = f"### {role}\n{content.strip()}\n" - self.chat_markdown = f"{self.chat_markdown.rstrip()}\n\n{block}" if self.chat_markdown else block - self._add_chat_message("user" if role.lower() in {"you", "user"} else "assistant", content) - - def create_bundle(self) -> None: - """Create a portable model export bundle.""" - - self.export_log.append("Creating model bundle...") - self.export_progress.setValue(15) - try: - output = export_project_bundle(Path(self.export_model_dir.text()), Path(self.export_dir.text())) - except Exception as exc: - self.export_log.append(f"Error: {exc}") - self.export_progress.setValue(0) - return - self.export_progress.setValue(100) - self.export_log.append(f"Bundle created: {output}") - self.export_status.setText("Export: bundle created") - - def quantize_model(self) -> None: - """Create a quantized FP16 checkpoint when selected.""" - - mode = self.quant_mode.currentText() - if not mode.startswith("FP16"): - self.export_log.append("This GGUF quantization target is planned. FP16 checkpoint quantization is available now.") - return - checkpoint = Path(self.export_model_dir.text()) / "final_model.pt" - output = Path(self.export_dir.text()) / "final_model_fp16.pt" - self.export_log.append("Creating FP16 checkpoint...") - self.export_progress.setValue(20) - try: - result = quantize_checkpoint(checkpoint, output, mode="fp16") - except Exception as exc: - self.export_log.append(f"Error: {exc}") - self.export_progress.setValue(0) - return - self.export_progress.setValue(100) - self.export_log.append(f"Quantized checkpoint created: {result}") - self.export_status.setText("Export: FP16 checkpoint ready") - - def export_hf_package(self) -> None: - """Create an HF-style MicroGPT package.""" - - self.export_log.append("Creating HF-style MicroGPT package...") - self.export_progress.setValue(20) - try: - result = export_hf_microgpt_package(Path(self.export_model_dir.text())) - except Exception as exc: - self.export_log.append(f"Error: {exc}") - self.export_progress.setValue(0) - return - self.export_progress.setValue(100) - self.export_log.append(f"HF package created: {result}") - self.export_log.append("Note: this package is MicroGPT model_type, not a llama.cpp-supported Llama model.") - self.export_status.setText("Export: HF package ready") - - def export_llama_adapter(self) -> None: - """Create a directly loadable Llama-family package when compatible.""" - - self.export_log.append("Creating Llama-compatible adapter package...") - self.export_progress.setValue(20) - try: - result = export_llama_adapter_package(Path(self.export_model_dir.text())) - except Exception as exc: - self.export_log.append(f"Error: {exc}") - self.export_progress.setValue(0) - return - self.export_progress.setValue(100) - self.export_log.append(f"Llama adapter package created: {result}") - self.export_status.setText("Export: Llama adapter ready") - - def convert_hf_to_gguf(self) -> None: - """Convert an HF-compatible model folder to GGUF through llama.cpp.""" - - model_dir_text = self.export_model_dir.text().strip() - llama_dir_text = self.llama_cpp_dir.text().strip() - output_text = self.gguf_output_path.text().strip() - if not model_dir_text: - QMessageBox.warning(self, "GGUF blocked", "Choose the model core folder first.") - return - if not (Path(model_dir_text) / "hf_model").exists(): - QMessageBox.warning( - self, - "GGUF blocked", - "GGUF conversion needs an HF model package first. Use Export HF Package, then convert a llama.cpp-supported model.", - ) - return - if not llama_dir_text: - QMessageBox.warning(self, "GGUF blocked", "Choose your local llama.cpp folder containing convert_hf_to_gguf.py.") - return - if not output_text: - QMessageBox.warning(self, "GGUF blocked", "Choose a GGUF output file path.") - return - self.export_log.append("Starting llama.cpp GGUF conversion...") - self.export_progress.setValue(0) - self._run_task( - export_gguf_with_llama_cpp, - ( - Path(model_dir_text), - Path(llama_dir_text), - Path(output_text), - self.gguf_outtype.currentText(), - ), - self._gguf_conversion_finished, - self.export_log, - self.export_progress, - button=self.gguf_convert_button, - busy_text="Converting GGUF", - ) - - @Slot(object) - def _gguf_conversion_finished(self, result: Any) -> None: - """Update UI after GGUF conversion finishes. - - Args: - result: GGUF output path. - """ - - self.export_progress.setValue(100) - self.export_log.append(f"GGUF created: {result}") - self.gguf_path.setText(str(result)) - self.export_status.setText("Export: GGUF ready") - self._clear_button_busy("Convert HF to GGUF") - - def _apply_preset(self, preset: str) -> None: - """Apply architecture values for a preset. - - Args: - preset: Selected preset name. - """ - - if preset == "Tiny": - self.n_embd.setValue(128) - self.n_head.setValue(4) - self.n_layer.setValue(4) - elif preset == "Small": - self.n_embd.setValue(512) - self.n_head.setValue(8) - self.n_layer.setValue(8) - - -def _ensure_valid_license(splash: "StartupValidationSplash") -> bool: - """Block app launch until a valid license is confirmed. - - Checks the currently stored license key (if any). On failure, shows - :class:`LicenseActivationDialog` in a loop -- unlike the general - startup-validation flow elsewhere in ``main()``, there is deliberately - no "continue anyway" option here: an unlicensed launch is not a - degraded-but-usable state, it's the one thing this app must not do. - - Args: - splash: Startup splash screen, used to show progress. - - Returns: - True if the app is licensed to proceed, False if the user cancelled - activation and the app should exit. - """ - - splash.append_log("Checking license...") - QApplication.processEvents() - - stored_key = load_stored_license_key() - if stored_key: - result = run_license_check_responsively(APP_VERSION, LICENSE_SERVER_URL) - if result.valid: - splash.append_log( - "✓ License valid" - + (" (offline grace period)" if result.used_offline_grace else "") - ) - return True - initial_message = result.reason - else: - initial_message = "No license activated on this machine yet." - - # A QSplashScreen-style window is designed to stay on top of other - # windows during startup -- which means it can end up covering a newly - # created dialog instead of the other way around. Hide it while the - # dialog is up rather than fight window-stacking order; it isn't doing - # anything useful to look at during activation anyway. - splash.hide() - try: - while True: - dialog = LicenseActivationDialog(APP_VERSION, LICENSE_SERVER_URL, initial_message) - dialog.setWindowIcon(MainWindow._static_app_icon()) - dialog.show() - dialog.raise_() - dialog.activateWindow() - if dialog.exec() != QDialog.Accepted: - LOGGER.info("License activation cancelled by user; exiting.") - QApplication.instance().setProperty("startup_aborted", True) - return False - splash.append_log("✓ License activated") - return True - finally: - splash.show() - splash.raise_() - - -def main(app: Optional[QApplication] = None, splash: Optional[StartupSplash] = None) -> None: - """Launch the PySide6 desktop application.""" - - owns_app = app is None - log_file = setup_logging() - qInstallMessageHandler(qt_message_handler) - LOGGER.info("Starting %s. Log file: %s", APP_NAME, log_file) - if sys.platform == "win32": - try: - ctypes.windll.shell32.SetCurrentProcessExplicitAppUserModelID(WINDOWS_APP_ID) - except Exception: - LOGGER.exception("Could not set Windows app user model ID") - app = app or QApplication(sys.argv) - app.setFont(QFont("Arial", 10)) - app.setWindowIcon(MainWindow._static_app_icon()) - splash = splash or StartupValidationSplash() - splash.setWindowIcon(MainWindow._static_app_icon()) - splash.show() - QTimer.singleShot(0, lambda: _apply_windows_taskbar_icon(splash)) - QApplication.processEvents() - if not _ensure_valid_license(splash): - splash.close() - if not owns_app: - app.quit() - app.setProperty("startup_aborted", True) - return - try: - _run_startup_validations(splash) - except Exception as exc: - LOGGER.exception("Startup validation failed") - splash.append_log(f"✗ Startup blocked: {exc}") - splash.close() - proceed = QMessageBox.question( - None, - "Startup validation failed", - "One or more startup checks failed.\n\n" - f"{exc}\n\n" - "Do you want to continue anyway?", - QMessageBox.Yes | QMessageBox.No, - QMessageBox.No, - ) - if proceed != QMessageBox.Yes: - if not owns_app: - app.quit() - app.setProperty("startup_aborted", True) - return - LOGGER.warning("User chose to continue after failed startup validation.") - splash.close() - while True: - chooser = ProjectChoiceDialog() - chooser.setWindowIcon(MainWindow._static_app_icon()) - QTimer.singleShot(0, lambda dialog=chooser: _apply_windows_taskbar_icon(dialog)) - if chooser.exec() != QDialog.Accepted: - LOGGER.info("Startup closed at project selection screen") - if not owns_app: - app.quit() - app.setProperty("startup_aborted", True) - return - window = MainWindow() - try: - if chooser.choice == "new": - base_dir = QFileDialog.getExistingDirectory( - None, - "Choose folder where the new project will be created", - str(DEFAULT_PROJECTS_DIR), - ) - if not base_dir: - window.deleteLater() - continue - project_name, ok = QInputDialog.getText(None, "Project name", "Enter project name:", text="MicroLLMProject") - if not ok: - window.deleteLater() - continue - project_name = project_name.strip() or "MicroLLMProject" - window._create_project_at(project_name, Path(base_dir)) - elif chooser.choice == "open": - project_file, _ = QFileDialog.getOpenFileName( - None, - "Open Micro LLM project", - str(DEFAULT_PROJECTS_DIR), - "Micro LLM project (project.json *.json);;All files (*)", - ) - if not project_file: - window.deleteLater() - continue - window._open_project_file(Path(project_file)) - elif chooser.choice == "recent": - if chooser.selected_project_file is None: - window.deleteLater() - continue - window._open_project_file(chooser.selected_project_file) - elif chooser.choice == "test_local_llm": - window.show_chat_only_mode() - except Exception as exc: - LOGGER.exception("Project setup failed during startup") - QMessageBox.critical(None, "Project setup failed", f"Could not complete project setup.\n\n{exc}") - window.deleteLater() - continue - break - window.show() - QTimer.singleShot(0, window.apply_windows_taskbar_icon) - interrupt_timer = QTimer() - interrupt_timer.timeout.connect(lambda: None) - interrupt_timer.start(200) - window.interrupt_timer = interrupt_timer - signal.signal(signal.SIGINT, lambda *_: QTimer.singleShot(0, window.request_shutdown_from_signal)) - if owns_app: - sys.exit(app.exec()) - - -if __name__ == "__main__": - main() diff --git a/llm_trainer/wiki_download.py b/llm_trainer/wiki_download.py deleted file mode 100644 index 81e6efc..0000000 --- a/llm_trainer/wiki_download.py +++ /dev/null @@ -1,1443 +0,0 @@ -#!/usr/bin/env python3 -""" -Wikipedia Dataset Downloader GUI - PySide6 -Fixed color theme for better visibility -""" - -import sys -import os -import json -import time -import re -import threading -from pathlib import Path -from datetime import datetime -from typing import List, Dict, Optional, Tuple -import requests -from PySide6.QtWidgets import * -from PySide6.QtCore import * -from PySide6.QtGui import * - - - -# ============================================================ -# Cleanup Configuration -# ============================================================ - -WORDS_PER_CHUNK = 2000 - -REMOVE_SECTIONS = [ - "References", - "External links", - "Bibliography", - "Further reading", - "See also", - "Notes", - "Sources", - "Citations", -] - -# ============================================================ - -css_patterns = [ - r"\.mw-parser-output.*?(?=(?:[A-Z][a-z].*?\n)|$)", - r"@media.*?(?=\n[A-Z]|\Z)", -] - -line_patterns = [ - r"^\s*v\s*t\s*e\s*$", - r"^\s*Main article:.*$", - r"^\s*Further information:.*$", - r"^\s*See also:.*$", - r"^\s*Coordinates:.*$", - r"^\s*This article is about.*$", -] - -reference_pattern = re.compile(r"\[[^\]]+\]") -whitespace_pattern = re.compile(r"\s+") - - -# ============================================================================ -# Backend: Wikipedia Downloader -# ============================================================================ - -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 -# ============================================================================ - -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 -# ============================================================================ - -class WikipediaDownloaderGUI(QMainWindow): - """Main GUI window for Wikipedia Downloader""" - - def __init__(self): - super().__init__() - self.downloader = WikipediaDownloaderBackend() - self.current_pages = [] - self.worker = None - self.output_dir = str(Path.home() / "wikipedia_dataset") - - self.init_ui() - self.setup_connections() - - 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() - - 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; - } - """) - - 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 -# ============================================================================ - -def main(): - app = QApplication(sys.argv) - - # ---------------------------------------------------------------------- - # Force Fusion style instead of Windows native style - # This prevents Windows Dark Mode from overriding widget colors. - # ---------------------------------------------------------------------- - app.setStyle("Fusion") - - # ---------------------------------------------------------------------- - # Light application palette - # ---------------------------------------------------------------------- - palette = QPalette() - - palette.setColor(QPalette.Window, QColor("#f0f2f5")) - palette.setColor(QPalette.WindowText, QColor("#1a1a1a")) - - palette.setColor(QPalette.Base, QColor("#ffffff")) - palette.setColor(QPalette.AlternateBase, QColor("#f8f9fa")) - - palette.setColor(QPalette.Text, QColor("#1a1a1a")) - - palette.setColor(QPalette.Button, QColor("#2ea44f")) - palette.setColor(QPalette.ButtonText, QColor("#ffffff")) - - palette.setColor(QPalette.BrightText, QColor("#ffffff")) - - palette.setColor(QPalette.Highlight, QColor("#4CAF50")) - palette.setColor(QPalette.HighlightedText, QColor("#ffffff")) - - palette.setColor(QPalette.ToolTipBase, QColor("#ffffff")) - palette.setColor(QPalette.ToolTipText, QColor("#1a1a1a")) - - palette.setColor(QPalette.PlaceholderText, QColor("#777777")) - - app.setPalette(palette) - - app.setApplicationName("Wikipedia Dataset Downloader") - app.setOrganizationName("TinyLLM") - - window = WikipediaDownloaderGUI() - window.show() - - sys.exit(app.exec()) - - - - -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() \ No newline at end of file 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 760f6cf..0000000 --- a/llm_trainer/worker/client.py +++ /dev/null @@ -1,530 +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("/")) - - -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/packager.py b/packager.py index ee59e1a..5d314a8 100644 --- a/packager.py +++ b/packager.py @@ -130,12 +130,13 @@ def build(*, clean: bool, runtime_dir: Path | None = None, gpu: bool = False) -> shutil.copytree(runtime_dir, bundle / "runtime", dirs_exist_ok=True) shutil.copy2(ROOT / "run_app.py", bundle / "run_app.py") shutil.copy2(ROOT / "runtime_setup.py", bundle / "runtime_setup.py") - shutil.copytree( - ROOT / "llm_trainer", - bundle / "llm_trainer", - ignore=shutil.ignore_patterns("default_data", "__pycache__", "*.pyc"), - dirs_exist_ok=True, - ) + for package_name in ("engine", "interface"): + shutil.copytree( + ROOT / package_name, + bundle / package_name, + ignore=shutil.ignore_patterns("default_data", "__pycache__", "*.pyc"), + dirs_exist_ok=True, + ) if target == "windows": installer = OUTPUT_ROOT / f"{APP_NAME}-{architecture}-Setup.exe" iscc = _find_inno_compiler() diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..dc423d4 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,52 @@ +[build-system] +requires = ["setuptools>=68"] +build-backend = "setuptools.build_meta" + +[tool.ruff] +target-version = "py312" +line-length = 100 +extend-exclude = [ + "packaging", +] + +[tool.ruff.lint] +select = [ + "E", + "F", + "W", + "I", + "D", + "UP", + "ANN", + "B", + "C4", + "SIM", +] +ignore = [ + "D203", + "D213", +] + +[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" + +[tool.mypy] +python_version = "3.12" +check_untyped_defs = true +disallow_any_generics = true +disallow_incomplete_defs = true +disallow_untyped_defs = true +no_implicit_optional = true +warn_redundant_casts = true +warn_return_any = true +warn_unused_ignores = true + +[tool.pytest.ini_options] +testpaths = ["tests"] +python_files = ["test_*.py"] +addopts = "-ra" diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..d6f2c4b --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,4 @@ +# Development quality tooling +mypy>=1.10 +pytest>=8.0 +ruff>=0.6 diff --git a/run_app.py b/run_app.py index 34ffc5b..63025e5 100644 --- a/run_app.py +++ b/run_app.py @@ -7,13 +7,13 @@ from PySide6.QtWidgets import QApplication -from llm_trainer.ui.startup_splash import StartupSplash +from interface.startup_splash import StartupSplash def _load_app(result: dict[str, object]) -> None: """Load the application module on the Qt thread for startup tests.""" try: - result["module"] = importlib.import_module("llm_trainer.ui.app") + result["module"] = importlib.import_module("interface.app") except Exception as exc: result["error"] = exc diff --git a/tests/test_dataset_mixture.py b/tests/test_dataset_mixture.py index 91d9a12..0330b50 100644 --- a/tests/test_dataset_mixture.py +++ b/tests/test_dataset_mixture.py @@ -3,8 +3,8 @@ import unittest from pathlib import Path -from llm_trainer.data import Document -from llm_trainer.dataset_mixture import ( +from engine.data import Document +from engine.dataset_mixture import ( _apply_dataset_mixture, _deduplicate_documents, _filter_repetitive_documents, diff --git a/tests/test_dataset_preview_artifacts.py b/tests/test_dataset_preview_artifacts.py index ca37ebb..36d269a 100644 --- a/tests/test_dataset_preview_artifacts.py +++ b/tests/test_dataset_preview_artifacts.py @@ -4,7 +4,7 @@ import unittest from pathlib import Path -from llm_trainer.dataset_preview import _has_prepared_token_artifacts +from engine.dataset_preview import _has_prepared_token_artifacts class DatasetPreviewArtifactsTests(unittest.TestCase): diff --git a/tests/test_dataset_preview_fast_scan.py b/tests/test_dataset_preview_fast_scan.py index e327cc7..53742af 100644 --- a/tests/test_dataset_preview_fast_scan.py +++ b/tests/test_dataset_preview_fast_scan.py @@ -4,8 +4,8 @@ import unittest from pathlib import Path -from llm_trainer.config import DatasetConfig -from llm_trainer.dataset_preview import scan_dataset_preview +from engine.config import DatasetConfig +from engine.dataset_preview import scan_dataset_preview class DatasetPreviewFastScanTests(unittest.TestCase): diff --git a/tests/test_external_dataset.py b/tests/test_external_dataset.py index 779080d..1446645 100644 --- a/tests/test_external_dataset.py +++ b/tests/test_external_dataset.py @@ -7,7 +7,7 @@ from pathlib import Path from unittest.mock import patch -from llm_trainer.external_dataset import ( +from engine.external_dataset import ( DatasetManifest, install_categories, is_newer_version, @@ -58,12 +58,38 @@ def copy_archive(_url: str, destination: Path, _progress: object) -> None: destination.write_bytes(archive.read_bytes()) destination = root / "installed" - with patch("llm_trainer.external_dataset._download", side_effect=copy_archive): + with patch("engine.external_dataset._download", side_effect=copy_archive): install_categories(manifest, destination, manifest_url="https://example.test/release/manifest.json") self.assertEqual((destination / "base_training" / "example.txt").read_text(), "example") self.assertEqual((destination / "version.txt").read_text(), "1.0.0\n") + def test_install_skips_existing_category_at_same_version(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + destination = root / "installed" + category_dir = destination / "base_training" + category_dir.mkdir(parents=True) + (category_dir / "existing.txt").write_text("keep", encoding="utf-8") + (destination / "version.txt").write_text("1.0.0\n", encoding="utf-8") + manifest = DatasetManifest.from_json({ + "dataset_id": "owner/repo", + "version": "1.0.0", + "categories": [{ + "name": "base_training", + "archive": "base.zip", + "size_bytes": 1, + "file_count": 1, + "sha256": "0" * 64, + }], + }) + + with patch("engine.external_dataset._download") as download: + install_categories(manifest, destination, categories=["base_training"]) + + download.assert_not_called() + self.assertEqual((category_dir / "existing.txt").read_text(encoding="utf-8"), "keep") + if __name__ == "__main__": unittest.main() diff --git a/tests/test_generation_kv_cache.py b/tests/test_generation_kv_cache.py index efe8643..d1d006b 100644 --- a/tests/test_generation_kv_cache.py +++ b/tests/test_generation_kv_cache.py @@ -4,8 +4,8 @@ import torch -from llm_trainer.config import ModelConfig -from llm_trainer.model import MicroGPT +from engine.config import ModelConfig +from engine.model import MicroGPT def _build_model(position_encoding: str) -> MicroGPT: diff --git a/tests/test_llama_chat_stream_metrics.py b/tests/test_llama_chat_stream_metrics.py index 99169d5..8296315 100644 --- a/tests/test_llama_chat_stream_metrics.py +++ b/tests/test_llama_chat_stream_metrics.py @@ -4,7 +4,7 @@ from threading import Lock from typing import Any, Iterator -from llm_trainer.llama_chat import LlamaChatSession +from engine.llama_chat import LlamaChatSession class _FakeLlama: diff --git a/tests/test_resume_checks.py b/tests/test_resume_checks.py index 4828feb..60774f9 100644 --- a/tests/test_resume_checks.py +++ b/tests/test_resume_checks.py @@ -7,8 +7,8 @@ import torch -from llm_trainer.config import ModelConfig, TrainingConfig -from llm_trainer.resume_checks import _resume_checkpoint_for, _validate_resume_compatibility +from engine.config import ModelConfig, TrainingConfig +from engine.resume_checks import _resume_checkpoint_for, _validate_resume_compatibility class ResumeChecksTests(unittest.TestCase): diff --git a/tests/test_startup_bootstrap.py b/tests/test_startup_bootstrap.py index d06a689..471d025 100644 --- a/tests/test_startup_bootstrap.py +++ b/tests/test_startup_bootstrap.py @@ -11,7 +11,7 @@ def test_app_module_is_loaded_lazily(self) -> None: result: dict[str, object] = {} with patch("run_app.importlib.import_module", return_value="loaded") as importer: run_app._load_app(result) - importer.assert_called_once_with("llm_trainer.ui.app") + importer.assert_called_once_with("interface.app") self.assertEqual(result["module"], "loaded") def test_import_errors_are_reported(self) -> None: diff --git a/tests/test_startup_tests.py b/tests/test_startup_tests.py index 32546ee..b2489f9 100644 --- a/tests/test_startup_tests.py +++ b/tests/test_startup_tests.py @@ -5,10 +5,20 @@ from pathlib import Path from unittest.mock import MagicMock, patch -from llm_trainer.ui.app import _run_startup_tests +from PySide6.QtWidgets import QApplication + +from interface.app import _run_startup_tests +from interface.startup import ProjectChoiceDialog class StartupTestsReportingTests(unittest.TestCase): + def test_project_choice_dialog_constructs_after_ui_module_split(self) -> None: + app = QApplication.instance() or QApplication([]) + dialog = ProjectChoiceDialog() + self.assertIsNotNone(dialog) + dialog.close() + app.processEvents() + def test_verbose_test_names_are_reported_to_splash_callback(self) -> None: with tempfile.TemporaryDirectory() as temp_dir: root = Path(temp_dir) @@ -25,7 +35,7 @@ def test_verbose_test_names_are_reported_to_splash_callback(self) -> None: process.wait.return_value = 0 reported: list[str] = [] - with patch("llm_trainer.ui.app.subprocess.Popen", return_value=process) as popen: + with patch("interface.app.subprocess.Popen", return_value=process) as popen: _run_startup_tests(root, tests_root, reported.append) command = popen.call_args.args[0] diff --git a/tests/test_token_dataset_stride.py b/tests/test_token_dataset_stride.py index 4b26287..116141c 100644 --- a/tests/test_token_dataset_stride.py +++ b/tests/test_token_dataset_stride.py @@ -2,7 +2,7 @@ import unittest -from llm_trainer.training import TokenDataset +from engine.training import TokenDataset class TokenDatasetStrideTests(unittest.TestCase): diff --git a/tests/test_training_gradient_accumulation.py b/tests/test_training_gradient_accumulation.py index 52e5c24..b680f6c 100644 --- a/tests/test_training_gradient_accumulation.py +++ b/tests/test_training_gradient_accumulation.py @@ -6,9 +6,9 @@ import torch from pathlib import Path -from llm_trainer.config import ModelConfig, TrainingConfig -from llm_trainer.training import train_model -from llm_trainer.model import MicroGPT +from engine.config import ModelConfig, TrainingConfig +from engine.training import train_model +from engine.model import MicroGPT def _run_training(train_tokens: list[int], gradient_accumulation: int, batch_size: int) -> int: diff --git a/tests/test_training_orchestrator_token_loading.py b/tests/test_training_orchestrator_token_loading.py index 04fc46a..5fa01e9 100644 --- a/tests/test_training_orchestrator_token_loading.py +++ b/tests/test_training_orchestrator_token_loading.py @@ -7,8 +7,8 @@ import numpy as np -from llm_trainer.training_orchestrator import _load_tokens_for_training -from llm_trainer.tokenizer import BOS_TOKEN, EOS_TOKEN, PAD_TOKEN, UNK_TOKEN, save_tokenizer_package +from engine.training_orchestrator import _load_tokens_for_training +from engine.tokenizer import BOS_TOKEN, EOS_TOKEN, PAD_TOKEN, UNK_TOKEN, save_tokenizer_package from tokenizers import Tokenizer from tokenizers.models import BPE diff --git a/tools/check_code_standards.py b/tools/check_code_standards.py new file mode 100644 index 0000000..54ad471 --- /dev/null +++ b/tools/check_code_standards.py @@ -0,0 +1,117 @@ +"""Check repository-wide Python source constraints. + +This dependency-free check is intentionally small so it can run before the +optional formatting, linting, and type-checking tools are installed. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path +import sys + + +DEFAULT_IGNORED_DIRECTORIES = { + ".git", + ".mypy_cache", + ".pytest_cache", + ".ruff_cache", + "__pycache__", + "build", + "dist", + "packaging", +} + + +def iter_python_files(root: Path) -> list[Path]: + """Return Python source files below a repository root. + + Args: + root: Repository directory to scan. + + Returns: + Python files in deterministic path order. + """ + files: list[Path] = [] + for path in root.rglob("*.py"): + if any(part in DEFAULT_IGNORED_DIRECTORIES for part in path.parts): + continue + files.append(path) + return sorted(files) + + +def count_lines(path: Path) -> int: + """Count lines in a UTF-8 Python source file. + + Args: + path: Python file to read. + + Returns: + Number of lines in the file. + """ + return len(path.read_text(encoding="utf-8-sig").splitlines()) + + +def find_oversized_files(root: Path, maximum_lines: int) -> list[tuple[Path, int]]: + """Find Python files exceeding the configured line limit. + + Args: + root: Repository directory to scan. + maximum_lines: Maximum permitted number of lines per file. + + Returns: + Paths and line counts for files over the limit. + """ + oversized: list[tuple[Path, int]] = [] + for path in iter_python_files(root): + line_count = count_lines(path) + if line_count > maximum_lines: + oversized.append((path, line_count)) + return oversized + + +def build_parser() -> argparse.ArgumentParser: + """Build the command-line argument parser. + + Returns: + Configured argument parser. + """ + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--root", + type=Path, + default=Path(__file__).resolve().parents[1], + help="Repository root to scan.", + ) + parser.add_argument( + "--max-lines", + type=int, + default=500, + help="Maximum permitted lines in a Python file.", + ) + return parser + + +def main() -> int: + """Run the repository source constraint check. + + Returns: + Zero when all files satisfy the configured constraints; otherwise one. + """ + args = build_parser().parse_args() + if args.max_lines < 1: + raise ValueError("--max-lines must be greater than zero") + + oversized = find_oversized_files(args.root.resolve(), args.max_lines) + if not oversized: + print(f"All Python files are at or below {args.max_lines} lines.") + return 0 + + print(f"Python files exceeding {args.max_lines} lines:") + for path, line_count in oversized: + print(f" {path.relative_to(args.root.resolve())}: {line_count} lines") + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/check_dependency_boundaries.py b/tools/check_dependency_boundaries.py new file mode 100644 index 0000000..944fb66 --- /dev/null +++ b/tools/check_dependency_boundaries.py @@ -0,0 +1,92 @@ +"""Check the dependency direction between the engine and desktop interface.""" + +from __future__ import annotations + +import argparse +import ast +import sys +from dataclasses import dataclass +from pathlib import Path + + +@dataclass(frozen=True) +class BoundaryViolation: + """A forbidden import found in a package.""" + + path: Path + line: int + imported: str + + +def _python_files(package_root: Path) -> list[Path]: + """Return Python files below a package in deterministic order.""" + + return sorted( + path + for path in package_root.rglob("*.py") + if "__pycache__" not in path.parts + ) + + +def _imports(path: Path) -> list[tuple[int, str]]: + """Return absolute top-level imports from a Python file.""" + + tree = ast.parse(path.read_text(encoding="utf-8-sig"), filename=str(path)) + imports: list[tuple[int, str]] = [] + for node in ast.walk(tree): + if isinstance(node, ast.Import): + imports.extend((node.lineno, alias.name) for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.level == 0 and node.module: + imports.append((node.lineno, node.module)) + return imports + + +def find_violations(root: Path) -> list[BoundaryViolation]: + """Find imports that violate the engine/interface dependency boundary.""" + + rules = (("engine", "interface"),) + violations: list[BoundaryViolation] = [] + for package_name, forbidden_root in rules: + package_root = root / package_name + if not package_root.is_dir(): + continue + for path in _python_files(package_root): + for line, imported in _imports(path): + if imported == forbidden_root or imported.startswith(f"{forbidden_root}."): + violations.append( + BoundaryViolation(path.relative_to(root), line, imported) + ) + return violations + + +def build_parser() -> argparse.ArgumentParser: + """Build the boundary-check command-line parser.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--root", + type=Path, + default=Path(__file__).resolve().parents[1], + help="Repository root to scan.", + ) + return parser + + +def main() -> int: + """Run the boundary check and return a process status.""" + + root = build_parser().parse_args().root.resolve() + violations = find_violations(root) + if violations: + print("Forbidden package-boundary imports found:") + for violation in violations: + print(f" {violation.path}:{violation.line}: {violation.imported}") + return 1 + print( + "Dependency boundaries are clean: engine cannot import interface." + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/curriculum_code.py b/tools/curriculum_code.py new file mode 100644 index 0000000..e77e3b8 --- /dev/null +++ b/tools/curriculum_code.py @@ -0,0 +1,305 @@ +from __future__ import annotations + +from pathlib import Path +try: + from .curriculum_shared import * +except ImportError: + from curriculum_shared import * + +def code_training_block(language: str, index: int) -> str: + """Create one base-training code explanation block. + + Args: + language: Key from CODE_LANGUAGE_SPECS. + index: Unique deterministic block number. + + Returns: + Plain text code teaching block. + """ + + spec = CODE_LANGUAGE_SPECS[language] + task_i, pattern_i, type_i, container_i, error_i, name_i = mixed_radix_pick( + index, + len(CODE_TASKS), + len(CODE_PATTERNS), + len(spec["types"]), + len(spec["containers"]), + len(spec["errors"]), + 997, + ) + label = spec["label"] + comment = spec["comment"] + task = CODE_TASKS[task_i] + pattern = CODE_PATTERNS[pattern_i] + type_name = spec["types"][type_i] + container = spec["containers"][container_i] + error = spec["errors"][error_i] + unique = f"{language}_{name_i}_{index}" + if language == "python": + snippet = ( + f"def process_{unique}(items: list[int]) -> int:\n" + f" total = 0\n" + f" for value in items:\n" + f" if value >= 0:\n" + f" total += value\n" + f" return total\n\n" + f"assert process_{unique}([1, -2, 3]) == 4" + ) + elif language in {"javascript", "typescript"}: + annotation = ": number[]" if language == "typescript" else "" + return_type = ": number" if language == "typescript" else "" + snippet = ( + f"function process_{unique}(items{annotation}){return_type} {{\n" + f" let total = 0;\n" + f" for (const value of items) {{\n" + f" if (value >= 0) total += value;\n" + f" }}\n" + f" return total;\n" + f"}}\n\n" + f"console.assert(process_{unique}([1, -2, 3]) === 4);" + ) + elif language == "java": + snippet = ( + f"static int process{unique.title().replace('_', '')}(java.util.List items) {{\n" + f" int total = 0;\n" + f" for (int value : items) {{\n" + f" if (value >= 0) total += value;\n" + f" }}\n" + f" return total;\n" + f"}}" + ) + elif language == "csharp": + snippet = ( + f"static int Process{unique.title().replace('_', '')}(IEnumerable items) {{\n" + f" var total = 0;\n" + f" foreach (var value in items) {{\n" + f" if (value >= 0) total += value;\n" + f" }}\n" + f" return total;\n" + f"}}" + ) + elif language == "cpp": + snippet = ( + f"int process_{unique}(const std::vector& items) {{\n" + f" int total = 0;\n" + f" for (int value : items) {{\n" + f" if (value >= 0) total += value;\n" + f" }}\n" + f" return total;\n" + f"}}" + ) + elif language == "rust": + snippet = ( + f"fn process_{unique}(items: &[i32]) -> i32 {{\n" + f" let mut total = 0;\n" + f" for value in items {{\n" + f" if *value >= 0 {{ total += *value; }}\n" + f" }}\n" + f" total\n" + f"}}" + ) + elif language == "go": + snippet = ( + f"func process{unique.title().replace('_', '')}(items []int) int {{\n" + f" total := 0\n" + f" for _, value := range items {{\n" + f" if value >= 0 {{ total += value }}\n" + f" }}\n" + f" return total\n" + f"}}" + ) + elif language == "sql": + snippet = ( + f"SELECT user_id, SUM(amount) AS total_{name_i}\n" + f"FROM payments\n" + f"WHERE amount >= 0\n" + f"GROUP BY user_id\n" + f"ORDER BY total_{name_i} DESC;" + ) + else: + snippet = ( + f"process_{unique}() {{\n" + f" local total=0\n" + f" for value in \"$@\"; do\n" + f" if [ \"$value\" -ge 0 ]; then total=$((total + value)); fi\n" + f" done\n" + f" printf '%s\\n' \"$total\"\n" + f"}}" + ) + return ( + f"{label} example {index}.\n" + f"Goal: teach how to {task} with a {pattern}.\n" + f"The example uses a {container} and a {type_name} value.\n" + f"```{spec['ext']}\n{snippet}\n```\n" + f"{comment} Read the code from top to bottom.\n" + f"The function receives data, skips invalid values, and returns one clear result.\n" + f"A common mistake in this topic is {error}.\n" + f"Check the empty input case before trusting the code.\n" + f"Keep names descriptive, keep steps small, and test one behavior at a time.\n" + ) + + +def code_fine_tune_block(language: str, index: int) -> str: + """Create one code fine-tuning instruction block. + + Args: + language: Key from CODE_LANGUAGE_SPECS. + index: Unique deterministic block number. + + Returns: + Instruction-style code fine-tuning block. + """ + + spec = CODE_LANGUAGE_SPECS[language] + task_i, pattern_i, type_i, container_i, error_i, variant = mixed_radix_pick( + index, + len(CODE_TASKS), + len(CODE_PATTERNS), + len(spec["types"]), + len(spec["containers"]), + len(spec["errors"]), + 2003, + ) + label = spec["label"] + task = CODE_TASKS[task_i] + pattern = CODE_PATTERNS[pattern_i] + type_name = spec["types"][type_i] + container = spec["containers"][container_i] + error = spec["errors"][error_i] + unique = f"{language}_{variant}_{index}" + if language == "sql": + response_code = ( + f"SELECT category, COUNT(*) AS count_{variant}\n" + f"FROM events\n" + f"WHERE created_at >= CURRENT_DATE - INTERVAL '7 days'\n" + f"GROUP BY category\n" + f"ORDER BY count_{variant} DESC;" + ) + bug_code = "SELECT category, COUNT(*) FROM events WHERE created_at >= CURRENT_DATE - INTERVAL '7 days';" + elif language == "bash": + response_code = ( + f"count_{unique}() {{\n" + f" local path=\"$1\"\n" + f" if [ ! -f \"$path\" ]; then return 1; fi\n" + f" grep -c \"ERROR\" \"$path\"\n" + f"}}" + ) + bug_code = "grep -c ERROR $path" + elif language == "python": + response_code = ( + f"def solve_{unique}(items: list[int], limit: int) -> list[int]:\n" + f" result: list[int] = []\n" + f" for value in items:\n" + f" if 0 <= value <= limit:\n" + f" result.append(value)\n" + f" return result\n\n" + f"assert solve_{unique}([1, -1, 5], 3) == [1]" + ) + bug_code = f"def solve_{unique}(items, limit):\n return [x for x in items if x <= limit]" + elif language in {"javascript", "typescript"}: + annotation = ": number[]" if language == "typescript" else "" + limit_annotation = ": number" if language == "typescript" else "" + return_annotation = ": number[]" if language == "typescript" else "" + response_code = ( + f"function solve_{unique}(items{annotation}, limit{limit_annotation}){return_annotation} {{\n" + f" const result = [];\n" + f" for (const value of items) {{\n" + f" if (value >= 0 && value <= limit) result.push(value);\n" + f" }}\n" + f" return result;\n" + f"}}\n\n" + f"console.assert(JSON.stringify(solve_{unique}([1, -1, 5], 3)) === JSON.stringify([1]));" + ) + bug_code = f"function solve_{unique}(items, limit) {{ return items.filter(x => x <= limit); }}" + elif language == "java": + method = f"solve{unique.title().replace('_', '')}" + response_code = ( + f"static java.util.List {method}(java.util.List items, int limit) {{\n" + f" java.util.List result = new java.util.ArrayList<>();\n" + f" for (int value : items) {{\n" + f" if (value >= 0 && value <= limit) result.add(value);\n" + f" }}\n" + f" return result;\n" + f"}}" + ) + bug_code = f"static java.util.List {method}(java.util.List items, int limit) {{ return null; }}" + elif language == "csharp": + method = f"Solve{unique.title().replace('_', '')}" + response_code = ( + f"static List {method}(IEnumerable items, int limit) {{\n" + f" var result = new List();\n" + f" foreach (var value in items) {{\n" + f" if (value >= 0 && value <= limit) result.Add(value);\n" + f" }}\n" + f" return result;\n" + f"}}" + ) + bug_code = f"static List {method}(IEnumerable items, int limit) => items.Where(x => x <= limit).ToList();" + elif language == "cpp": + response_code = ( + f"std::vector solve_{unique}(const std::vector& items, int limit) {{\n" + f" std::vector result;\n" + f" for (int value : items) {{\n" + f" if (value >= 0 && value <= limit) result.push_back(value);\n" + f" }}\n" + f" return result;\n" + f"}}" + ) + bug_code = f"std::vector solve_{unique}(std::vector& items, int limit) {{ return items; }}" + elif language == "rust": + response_code = ( + f"fn solve_{unique}(items: &[i32], limit: i32) -> Vec {{\n" + f" items.iter()\n" + f" .copied()\n" + f" .filter(|value| *value >= 0 && *value <= limit)\n" + f" .collect()\n" + f"}}" + ) + bug_code = f"fn solve_{unique}(items: Vec, limit: i32) -> Vec {{ items }}" + elif language == "go": + method = f"solve{unique.title().replace('_', '')}" + response_code = ( + f"func {method}(items []int, limit int) []int {{\n" + f" result := make([]int, 0, len(items))\n" + f" for _, value := range items {{\n" + f" if value >= 0 && value <= limit {{ result = append(result, value) }}\n" + f" }}\n" + f" return result\n" + f"}}" + ) + bug_code = f"func {method}(items []int, limit int) []int {{ return items }}" + return ( + f"Instruction: Write {label} code to {task}.\n" + f"User context: Use a {pattern}. The input involves a {container}. The important type is {type_name}.\n" + f"Response:\n" + f"```{spec['ext']}\n{response_code}\n```\n" + f"Explanation: The solution separates input handling from the core operation.\n" + f"It names the result clearly and keeps each step small.\n" + f"Edge case: empty input should return a safe default or a clear error.\n" + f"Debugging example:\n" + f"```{spec['ext']}\n{bug_code}\n```\n" + f"The likely issue is {error}.\n" + f"Fix: validate inputs, check boundaries, and test the smallest failing case first.\n" + f"Final answer: use the shown pattern, then add tests for normal, empty, and invalid inputs.\n" + ) + + +def write_target_bytes(path: Path, block_factory, target_bytes: int) -> None: + """Write generated corpus blocks until a file reaches the target size. + + Args: + path: Output file path. + block_factory: Callable accepting a block index and returning text. + target_bytes: Minimum UTF-8 byte size to write. + """ + + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8", newline="\n") as handle: + index = 0 + while handle.tell() < target_bytes: + handle.write(block_factory(index)) + handle.write("\n\n") + index += 1 + + + diff --git a/tools/curriculum_finetune.py b/tools/curriculum_finetune.py new file mode 100644 index 0000000..26cf23d --- /dev/null +++ b/tools/curriculum_finetune.py @@ -0,0 +1,462 @@ +from __future__ import annotations + +try: + from .curriculum_shared import * +except ImportError: + from curriculum_shared import * + +def programming_deep_blocks(count: int, topic: str) -> list[str]: + """Create programming-focused corpus blocks. + + Args: + count: Number of blocks to generate. + topic: Programming topic name. + + Returns: + Generated corpus blocks. + """ + + python_examples = [ + ( + "Python list filtering", + "numbers = [1, 2, 3, 4, 5]\n" + "even = []\n" + "for number in numbers:\n" + " if number % 2 == 0:\n" + " even.append(number)\n" + "print(even)", + "The list stores numbers in order.\nThe loop checks each number.\nThe percent operator gives the remainder.\nA remainder of zero means the number is even.", + ), + ( + "Python function", + "def area(width, height):\n" + " return width * height\n\n" + "result = area(6, 4)\n" + "print(result)", + "The function receives width and height.\nThe return statement sends back the answer.\nThe area is twenty four.", + ), + ( + "Python dictionary", + "scores = {'Mina': 8, 'Ravi': 9}\n" + "scores['Lena'] = 7\n" + "for name, score in scores.items():\n" + " print(name, score)", + "A dictionary maps keys to values.\nThe key is a name.\nThe value is a score.\nThe items method gives key and value pairs.", + ), + ] + javascript_examples = [ + ( + "JavaScript array map", + "const prices = [10, 20, 30];\n" + "const doubled = prices.map(price => price * 2);\n" + "console.log(doubled);", + "An array keeps values in order.\nThe map method creates a new array.\nThe arrow function runs once for each value.", + ), + ( + "JavaScript async function", + "async function loadUser(id) {\n" + " const response = await fetch(`/users/${id}`);\n" + " return response.json();\n" + "}", + "The async keyword allows await.\nAwait pauses until the promise settles.\nThis is useful for network requests.", + ), + ] + java_examples = [ + ( + "Java class", + "class Counter {\n" + " private int value = 0;\n" + " void increment() {\n" + " value++;\n" + " }\n" + " int getValue() {\n" + " return value;\n" + " }\n" + "}", + "A class groups data and behavior.\nThe field stores the count.\nThe method changes the count.\nPrivate data is hidden from outside code.", + ), + ] + cpp_examples = [ + ( + "C++ vector loop", + "#include \n" + "#include \n\n" + "int main() {\n" + " std::vector values{1, 2, 3};\n" + " int total = 0;\n" + " for (int value : values) {\n" + " total += value;\n" + " }\n" + " std::cout << total << '\\n';\n" + "}", + "A vector stores many values.\nThe range loop visits each value.\nThe total variable accumulates the sum.", + ), + ( + "C pointer safety", + "#include \n\n" + "int main(void) {\n" + " int value = 5;\n" + " int *ptr = &value;\n" + " printf(\"%d\\n\", *ptr);\n" + " return 0;\n" + "}", + "A pointer stores an address.\nThe address operator gets the address.\nThe star operator reads the value at the address.", + ), + ] + rust_go_examples = [ + ( + "Rust ownership", + "fn main() {\n" + " let name = String::from(\"Mina\");\n" + " print_name(&name);\n" + " println!(\"{}\", name);\n" + "}\n\n" + "fn print_name(value: &String) {\n" + " println!(\"{}\", value);\n" + "}", + "The ampersand borrows the string.\nBorrowing lets a function read without taking ownership.\nThe original value can still be used later.", + ), + ( + "Go error handling", + "file, err := os.Open(\"data.txt\")\n" + "if err != nil {\n" + " return err\n" + "}\n" + "defer file.Close()", + "Go returns errors as values.\nThe code checks the error immediately.\nThe defer statement closes the file later.", + ), + ] + sql_shell_examples = [ + ( + "SQL selection", + "SELECT name, age\n" + "FROM users\n" + "WHERE age >= 18\n" + "ORDER BY name;", + "The SELECT clause chooses columns.\nThe FROM clause chooses a table.\nThe WHERE clause filters rows.\nThe ORDER BY clause sorts the result.", + ), + ( + "Bash pipeline", + "cat access.log | grep ERROR | sort | uniq -c", + "A pipeline sends output to the next command.\nGrep filters matching lines.\nSort groups similar lines.\nUniq counts repeated lines.", + ), + ( + "PowerShell pipeline", + "Get-ChildItem -File | Where-Object { $_.Length -gt 1MB } | Select-Object Name, Length", + "PowerShell passes objects through the pipeline.\nWhere-Object filters objects.\nSelect-Object chooses properties to display.", + ), + ] + web_examples = [ + ( + "HTML form", + "
    \n" + " \n" + " \n" + " \n" + "
    ", + "A form collects input.\nA label tells the user what to enter.\nA button submits the form.", + ), + ( + "CSS button", + ".button {\n" + " background: #222;\n" + " color: white;\n" + " padding: 8px 12px;\n" + "}\n" + ".button:hover {\n" + " background: #444;\n" + "}", + "CSS changes how elements look.\nThe hover rule runs when the pointer is over the button.", + ), + ] + algorithm_examples = [ + ( + "Binary search", + "def binary_search(values, target):\n" + " low = 0\n" + " high = len(values) - 1\n" + " while low <= high:\n" + " mid = (low + high) // 2\n" + " if values[mid] == target:\n" + " return mid\n" + " if values[mid] < target:\n" + " low = mid + 1\n" + " else:\n" + " high = mid - 1\n" + " return -1", + "Binary search works on sorted data.\nEach step removes half of the remaining choices.\nThis makes it faster than checking every item.", + ), + ( + "Queue with list", + "from collections import deque\n\n" + "queue = deque()\n" + "queue.append('first')\n" + "queue.append('second')\n" + "item = queue.popleft()\n" + "print(item)", + "A queue is first in, first out.\nAppend adds to the back.\nPopleft removes from the front.", + ), + ] + debugging_examples = [ + ( + "Read the traceback", + "Traceback says the error line.\nStart at the last line.\nFind the exception name.\nThen inspect the code near that line.\nA NameError often means a variable name is missing or misspelled.", + "Debugging starts with evidence.\nDo not guess first.\nRead the error.\nReproduce the bug.\nChange one thing.\nRun the test again.", + ), + ( + "Off by one error", + "for index in range(len(items)):\n" + " print(items[index])", + "Indexes start at zero in many languages.\nThe last index is length minus one.\nAn off by one error reads before the start or after the end.", + ), + ] + sets = { + "python": python_examples, + "javascript_web": javascript_examples + web_examples, + "java_csharp": java_examples, + "c_cpp_systems": cpp_examples, + "rust_go": rust_go_examples, + "sql_shell": sql_shell_examples, + "algorithms": algorithm_examples, + "debugging": debugging_examples, + "full_stack": web_examples + sql_shell_examples + javascript_examples, + "data_structures": algorithm_examples + python_examples, + "software_engineering": debugging_examples + java_examples + rust_go_examples, + "mixed_language": python_examples + javascript_examples + cpp_examples + rust_go_examples + sql_shell_examples, + } + examples = sets[topic] + blocks = [] + for index in range(count): + title, code, explanation = examples[index % len(examples)] + scenario = index % 11 + blocks.append( + f"{title}.\n" + f"Example number {index + 1}.\n" + f"{code}\n" + f"{explanation}\n" + f"The programmer should name variables clearly.\n" + f"The program should handle expected input.\n" + f"The program should fail clearly when input is wrong.\n" + f"A small test should check the normal case.\n" + f"A second test should check an edge case.\n" + f"If scenario {scenario} changes, update the test first.\n" + f"Good code is readable, correct, and easy to change." + ) + return blocks + + +def conversation_fine_tune_blocks(count: int, topic: str) -> list[str]: + """Create conversation fine-tuning corpus blocks. + + Args: + count: Number of blocks to generate. + topic: Conversation scenario group. + + Returns: + Conversation training blocks. + """ + + scenarios = { + "daily_help": [ + ("planning a day", "I have school and chores today.", "Make a short list. Do the urgent task first."), + ("finding a lost key", "I lost my key.", "Check the last place you used it. Then check pockets and bags."), + ("choosing food", "I want a simple dinner.", "Rice, vegetables, and eggs can make a quick meal."), + ("packing for a trip", "I don't know what to pack.", "List essentials first: documents, clothes, chargers."), + ("running late", "I'm running late for an appointment.", "Call ahead if you can, then focus on getting there safely."), + ("organizing a closet", "My closet is a mess.", "Sort into keep, donate, and discard piles first."), + ("choosing a gift", "I need a gift for a friend.", "Think about what they enjoy doing in their free time."), + ("scheduling chores", "I keep forgetting chores.", "Write them on a weekly checklist you can see each morning."), + ("saving time", "Mornings feel rushed.", "Prepare clothes and bags the night before."), + ("handling a spill", "I spilled water on my desk.", "Unplug electronics first, then blot the liquid with a cloth."), + ("choosing a route", "Traffic is bad on my usual route.", "Check a map app for a lighter route before leaving."), + ("staying hydrated", "I forget to drink water.", "Keep a bottle within sight as a visual reminder."), + ], + "learning_tutor": [ + ("asking about fractions", "What is one half?", "One half means one part out of two equal parts."), + ("grammar help", "When do I use went?", "Use went for the past. I went to school yesterday."), + ("science question", "Why does ice melt?", "Ice melts when it gains heat. It becomes liquid water."), + ("history question", "Why do maps matter in history?", "Maps show where events happened and how borders changed."), + ("vocabulary help", "What does 'concise' mean?", "Concise means expressing something clearly in few words."), + ("multiplication help", "What is 7 times 8?", "7 times 8 is 56."), + ("reading comprehension", "How do I find the main idea?", "Look at the first and last sentences of the paragraph."), + ("study habits", "How should I study for a test?", "Break the material into small chunks and review daily."), + ("essay structure", "How do I start an essay?", "Open with a clear sentence stating your main point."), + ("geometry help", "What is a right angle?", "A right angle measures exactly ninety degrees."), + ("spelling help", "How do I remember 'necessary'?", "Remember: one collar, two sleeves -- one C, two S's."), + ("note taking", "How should I take notes?", "Write short phrases, not full sentences, and use bullet points."), + ], + "coding_mentor": [ + ("debugging Python", "My loop prints too many lines.", "Check the range. The stop value may be too large."), + ("understanding variables", "What is a variable?", "A variable is a name that stores a value."), + ("reading an error", "What does NameError mean?", "It usually means Python cannot find that variable name."), + ("understanding functions", "Why use functions?", "Functions let you reuse steps without repeating code."), + ("list vs tuple", "What's the difference between a list and a tuple?", "Lists can change after creation; tuples cannot."), + ("indexing errors", "I got an IndexError.", "You likely tried to access a position beyond the list's length."), + ("understanding loops", "When do I use a while loop?", "Use a while loop when you don't know how many times to repeat in advance."), + ("git basics", "What does git commit do?", "It saves a snapshot of your staged changes to the project history."), + ("testing code", "Why write tests?", "Tests catch mistakes early and confirm code still works after changes."), + ("naming variables", "How should I name variables?", "Use short, descriptive names that explain what the value represents."), + ("recursion confusion", "I don't understand recursion.", "A recursive function solves a small case, then calls itself on a smaller one."), + ("reading stack traces", "The stack trace is confusing.", "Start from the bottom-most line pointing to your own code."), + ], + "empathy_support": [ + ("nervous before exam", "I feel nervous about my exam.", "That is normal. Study small parts and take breaks."), + ("friend conflict", "My friend ignored me.", "Ask calmly what happened. Listen before deciding."), + ("mistake at work", "I made a mistake.", "Own it, fix what you can, and learn the cause."), + ("feeling overwhelmed", "I have too much to do.", "Pick one task, finish it, then move to the next."), + ("disappointment", "I didn't get the result I wanted.", "It's okay to feel disappointed. Consider what to try differently."), + ("homesickness", "I miss home.", "That feeling is common. Reach out to family when you can."), + ("public speaking fear", "I'm scared to speak in front of others.", "Practice out loud a few times; familiarity reduces nerves."), + ("comparison worry", "I feel behind compared to others.", "Everyone moves at a different pace. Focus on your own progress."), + ("difficult feedback", "I got harsh feedback.", "Take a breath, look for the useful part, and set the rest aside."), + ("change anxiety", "Things are changing and I feel unsettled.", "Focus on what stays the same and what you can control."), + ("apologizing", "I need to apologize but don't know how.", "Be specific about what happened and how you'll do better."), + ("burnout", "I feel exhausted from working nonstop.", "Rest is productive too. Consider a short, real break."), + ], + "professional_chat": [ + ("email rewrite", "Can you make this email polite?", "Yes. Keep it short, clear, and respectful."), + ("meeting plan", "How should I run a meeting?", "Set a goal, list topics, and end with action items."), + ("status update", "I need to report progress.", "Say what is done, what is blocked, and what comes next."), + ("giving feedback", "How do I give feedback kindly?", "Be specific, focus on the work, and suggest a next step."), + ("declining a request", "How do I say no politely?", "Thank them, explain briefly, and offer an alternative if possible."), + ("negotiating a deadline", "I need more time on a project.", "Explain the reason and propose a new, realistic date early."), + ("onboarding a teammate", "How do I help a new hire settle in?", "Share key contacts, documents, and a short first-week plan."), + ("prioritizing tasks", "I have too many tasks today.", "Rank by deadline and impact, then start with the most urgent."), + ("summarizing a call", "How do I summarize a meeting?", "List decisions made, owners, and deadlines in a few lines."), + ("cold outreach", "How do I write a cold email?", "Keep it short, state the purpose, and make the ask clear."), + ("handling conflict", "A coworker disagreed with my plan.", "Ask about their concern directly and look for common ground."), + ("requesting resources", "How do I ask for more budget?", "Explain the need, the expected benefit, and the cost clearly."), + ], + } + items = scenarios[topic] + endings = [ + "The best next step is to act carefully and review the result.", + "The best next step is to keep it simple and adjust later.", + "The best next step is to ask for help if anything is unclear.", + "The best next step is to write it down so it isn't forgotten.", + "The best next step is to check in again after trying it.", + ] + period = combinatorial_period(len(items), 9, len(endings)) + blocks = [] + for index in range(min(count, period)): + item_i, turn, ending_i = mixed_radix_pick(index, len(items), 9, len(endings)) + title, user_text, assistant_text = items[item_i] + blocks.append( + f"Conversation: {title}.\n" + f"User: {user_text}\n" + f"Assistant: {assistant_text}\n" + f"User: Can you explain simply?\n" + f"Assistant: Yes. I will use short steps.\n" + f"Assistant: First, understand the problem.\n" + f"Assistant: Second, choose a small action.\n" + f"Assistant: Third, check the result.\n" + f"User: What should I avoid?\n" + f"Assistant: Avoid guessing when facts are missing.\n" + f"Assistant: Ask a clear question if needed.\n" + f"User: Give me a final answer.\n" + f"Assistant: {endings[ending_i]}\n" + f"This dialogue teaches helpful conversation turn {turn}." + ) + return blocks + + +def instruction_fine_tune_blocks(count: int, topic: str) -> list[str]: + """Create instruction fine-tuning corpus blocks. + + Args: + count: Number of blocks to generate. + topic: Instruction task group. + + Returns: + Instruction training blocks. + """ + + tasks = { + "writing_tasks": [ + ("Rewrite this sentence in simpler English.", "The child rapidly moved across the room.", "The child ran across the room."), + ("Summarize this passage.", "Mina planted seeds. She watered them. After many days, leaves grew.", "Mina planted and cared for seeds until they grew leaves."), + ("Make this polite.", "Send the report now.", "Please send the report when you have a moment."), + ("Shorten this sentence.", "Due to the fact that it was raining, we decided to stay inside.", "Because it was raining, we stayed inside."), + ("Fix the grammar.", "She don't like the plan.", "She doesn't like the plan."), + ("Make this more formal.", "Hey, can you send that file?", "Could you please send the file at your convenience?"), + ("Combine these sentences.", "The dog barked. The dog ran to the door.", "The dog barked and ran to the door."), + ("Add a stronger verb.", "The team did a good job on the project.", "The team excelled on the project."), + ("Remove redundancy.", "In my opinion, I think the plan is good.", "I think the plan is good."), + ("Write a topic sentence.", "Details about rainforests having high rainfall and diverse species.", "Rainforests are defined by heavy rainfall and remarkable species diversity."), + ], + "reasoning_tasks": [ + ("Solve the word problem.", "A box has 6 pens. Ravi adds 4 pens. How many pens are there?", "There are 10 pens."), + ("Choose the safer action.", "A wire is broken. Should Tara touch it or call an adult?", "Tara should call an adult."), + ("Find the cause.", "The lamp does not turn on. The bulb is loose.", "The loose bulb may be the cause."), + ("Solve the word problem.", "Lena has 15 stickers and gives 6 away. How many are left?", "9 stickers are left."), + ("Order the steps.", "Steps: pour water, boil water, add tea leaves, given out of order.", "Pour water, boil water, add tea leaves."), + ("Spot the contradiction.", "The store is open every day. The store is closed on Sundays.", "These two statements contradict each other."), + ("Draw a conclusion.", "All birds in the flock flew south. It is now winter here.", "The birds likely migrated for winter."), + ("Find the missing step.", "Recipe skips from 'mix batter' to 'serve cake' with nothing baked.", "The recipe is missing a baking step."), + ("Compare two options.", "Option A costs less but takes longer. Option B costs more but is faster.", "Choose based on whether time or cost matters more."), + ("Explain the pattern.", "2, 4, 6, 8, ...", "The pattern adds 2 to get each next number."), + ], + "coding_tasks": [ + ("Write a Python function that adds two numbers.", "Use parameters a and b.", "def add(a, b):\n return a + b"), + ("Explain this code.", "print(len([1, 2, 3]))", "It creates a list with three items and prints its length, which is 3."), + ("Fix the bug.", "for i in range(3):\nprint(i)", "Indent the print line inside the loop."), + ("Write a function that returns the max of two numbers.", "Use parameters a and b.", "def maximum(a, b):\n return a if a > b else b"), + ("Explain this code.", "x = [n * n for n in range(5)]", "It builds a list of squares for numbers 0 through 4 using a list comprehension."), + ("Fix the bug.", "def greet(name)\n print('Hello ' + name)", "Add a colon after the function signature: def greet(name):"), + ("Write a function that checks if a number is even.", "Use one parameter n.", "def is_even(n):\n return n % 2 == 0"), + ("Explain this code.", "total = sum([1, 2, 3])", "It adds up the numbers in the list, giving a total of 6."), + ("Fix the bug.", "if x = 5:\n print('five')", "Use == for comparison instead of =: if x == 5:"), + ("Write a function that reverses a string.", "Use one parameter text.", "def reverse(text):\n return text[::-1]"), + ], + "classification_tasks": [ + ("Classify the sentence.", "The sky is cloudy today.", "Category: weather observation."), + ("Classify the request.", "Can you help me debug this error?", "Category: coding help."), + ("Classify the emotion.", "I am proud because I finished the project.", "Emotion: proud."), + ("Classify the sentence.", "Water boils at 100 degrees Celsius.", "Category: science fact."), + ("Classify the request.", "Please summarize this article for me.", "Category: writing help."), + ("Classify the emotion.", "I felt nervous before the interview.", "Emotion: nervous."), + ("Classify the sentence.", "Paris is the capital of France.", "Category: geography fact."), + ("Classify the request.", "Can you check my math homework?", "Category: math help."), + ("Classify the emotion.", "I was relieved when the test was over.", "Emotion: relieved."), + ("Classify the sentence.", "The stock market fell sharply today.", "Category: financial news."), + ], + "format_following": [ + ("Answer with two bullet points.", "Give two safe cooking tips.", "- Wash your hands.\n- Turn off the stove after cooking."), + ("Return only the number.", "What is 8 plus 5?", "13"), + ("Use a short answer.", "Why do plants need light?", "Plants use light to make food."), + ("Answer with two bullet points.", "Give two tips for studying.", "- Take short breaks.\n- Review notes daily."), + ("Return only the number.", "What is 12 minus 7?", "5"), + ("Use a short answer.", "Why do we wear seatbelts?", "Seatbelts help prevent injury in a crash."), + ("Answer with three bullet points.", "List three parts of a plant.", "- Roots\n- Stem\n- Leaves"), + ("Return only the word.", "What do bees produce?", "Honey"), + ("Use one sentence.", "What is gravity?", "Gravity is the force that pulls objects toward each other."), + ("Answer in a single word.", "What gas do plants release during photosynthesis?", "Oxygen"), + ], + } + items = tasks[topic] + closers = [ + "This instruction sample teaches format control.", + "This instruction sample teaches staying on topic.", + "This instruction sample teaches concise responses.", + "This instruction sample teaches following the exact request.", + ] + period = combinatorial_period(len(items), 13, len(closers)) + blocks = [] + for index in range(min(count, period)): + item_i, turn, closer_i = mixed_radix_pick(index, len(items), 13, len(closers)) + instruction, input_text, output_text = items[item_i] + blocks.append( + f"Instruction: {instruction}\n" + f"Input: {input_text}\n" + f"Response: {output_text}\n" + f"The response follows the instruction.\n" + f"The response stays focused on the user request.\n" + f"The response avoids extra unrelated text.\n" + f"If information is missing, ask one clear question.\n" + f"If the task is simple, answer directly.\n" + f"If the task needs steps, use short ordered steps.\n" + f"{closers[closer_i]} (variant {turn})" + ) + return blocks + + diff --git a/tools/curriculum_shared.py b/tools/curriculum_shared.py new file mode 100644 index 0000000..845bf4b --- /dev/null +++ b/tools/curriculum_shared.py @@ -0,0 +1,200 @@ +from __future__ import annotations + +from pathlib import Path + +NAMES = [ + "Mina", "Ravi", "Lena", "Omar", "Sara", "Tara", "Jin", "Ada", + "Leo", "Nia", "Sam", "Priya", +] + +MODIFIERS = ["carefully", "quickly", "calmly", "clearly", "patiently"] + +CLOSERS = [ + "What did you notice?", + "Why do you think that happens?", + "How could you check this?", + "What would change if one detail changed?", + "Explain it back in your own words.", +] + +CODE_LANGUAGE_SPECS = { + "python": { + "label": "Python", + "ext": "py", + "comment": "#", + "types": ["list[int]", "dict[str, int]", "str", "tuple[int, int]", "set[str]"], + "containers": ["list", "dictionary", "set", "tuple", "file"], + "errors": ["IndexError", "KeyError", "TypeError", "ValueError", "NameError"], + }, + "javascript": { + "label": "JavaScript", + "ext": "js", + "comment": "//", + "types": ["Array", "Object", "string", "number", "Promise"], + "containers": ["array", "object", "map", "set", "DOM node"], + "errors": ["TypeError", "ReferenceError", "RangeError", "SyntaxError", "Promise rejection"], + }, + "typescript": { + "label": "TypeScript", + "ext": "ts", + "comment": "//", + "types": ["number[]", "Record", "string", "Promise", "ReadonlyArray"], + "containers": ["typed array", "record", "interface", "union", "generic"], + "errors": ["type mismatch", "undefined value", "narrowing error", "implicit any", "async error"], + }, + "java": { + "label": "Java", + "ext": "java", + "comment": "//", + "types": ["List", "Map", "String", "Optional", "Set"], + "containers": ["ArrayList", "HashMap", "HashSet", "class", "stream"], + "errors": ["NullPointerException", "IndexOutOfBoundsException", "IllegalArgumentException", "ClassCastException", "IOException"], + }, + "csharp": { + "label": "C#", + "ext": "cs", + "comment": "//", + "types": ["List", "Dictionary", "string", "Task", "IEnumerable"], + "containers": ["List", "Dictionary", "HashSet", "class", "LINQ query"], + "errors": ["NullReferenceException", "IndexOutOfRangeException", "InvalidOperationException", "ArgumentException", "async deadlock"], + }, + "cpp": { + "label": "C++", + "ext": "cpp", + "comment": "//", + "types": ["vector", "unordered_map", "string", "unique_ptr", "optional"], + "containers": ["vector", "unordered_map", "set", "struct", "iterator"], + "errors": ["segmentation fault", "dangling pointer", "out_of_range", "memory leak", "undefined behavior"], + }, + "rust": { + "label": "Rust", + "ext": "rs", + "comment": "//", + "types": ["Vec", "HashMap", "String", "Option", "Result"], + "containers": ["Vec", "HashMap", "slice", "struct", "iterator"], + "errors": ["borrow checker error", "panic", "lifetime error", "unwrap failure", "type mismatch"], + }, + "go": { + "label": "Go", + "ext": "go", + "comment": "//", + "types": ["[]int", "map[string]int", "string", "error", "chan int"], + "containers": ["slice", "map", "struct", "goroutine", "channel"], + "errors": ["nil pointer", "index out of range", "data race", "ignored error", "deadlock"], + }, + "sql": { + "label": "SQL", + "ext": "sql", + "comment": "--", + "types": ["INTEGER", "TEXT", "TIMESTAMP", "BOOLEAN", "DECIMAL"], + "containers": ["table", "index", "view", "join", "transaction"], + "errors": ["missing index", "duplicate key", "bad join", "null value", "slow query"], + }, + "bash": { + "label": "Bash", + "ext": "sh", + "comment": "#", + "types": ["string", "array", "exit code", "path", "environment variable"], + "containers": ["loop", "function", "pipe", "process", "file"], + "errors": ["missing quote", "bad path", "nonzero exit", "unset variable", "permission denied"], + }, +} + +CODE_TASKS = [ + "parse input", + "validate data", + "filter a collection", + "count repeated values", + "read a file safely", + "write a small helper", + "handle an error", + "sort records", + "cache a result", + "format output", + "test an edge case", + "split work into functions", +] + +CODE_PATTERNS = [ + "loop", + "function", + "guard clause", + "map lookup", + "unit test", + "small class", + "command handler", + "parser", + "retry step", + "cleanup step", +] + + +def mixed_radix_pick(index: int, *sizes: int) -> list[int]: + """Decompose an index into independent per-axis picks. + + Unlike applying ``index % len(list)`` to several lists at once (which + repeats after ``lcm`` of the list lengths -- often a tiny number), this + treats ``index`` as a mixed-radix counter across every axis. The combined + period is the *product* of all axis sizes, so a handful of modest lists + (say four lists of 15-20 items) already yields a combinatorial space of + tens of thousands of unique combinations before anything repeats. + + Args: + index: Zero-based block index. + *sizes: Length of each axis, in the same order picks are needed. + + Returns: + One pick per axis, each in ``range(0, size)``. + """ + + picks = [] + remaining = index + for size in sizes: + size = max(1, size) + picks.append(remaining % size) + remaining //= size + return picks + +def combinatorial_period(*sizes: int) -> int: + """Return the number of unique combinations `mixed_radix_pick` can produce.""" + + period = 1 + for size in sizes: + period *= max(1, size) + return period + +def write_blocks(path: Path, blocks: list[str], min_unique_ratio: float = 0.9) -> None: + """Write plain-text corpus blocks, guarding against templated duplication. + + A generator that technically returns ``count`` blocks but only cycles + through a handful of unique strings silently produces a dataset that is + almost entirely duplicate data -- wasted disk, wasted training compute, + and a validation split that can't mean anything because train and + validation end up full of the same repeated content. This raises loudly + instead of writing a file that *looks* like a real corpus but isn't. + + Args: + path: Output text file. + blocks: Corpus blocks. + min_unique_ratio: Minimum allowed fraction of unique blocks. Raise + the ratio for categories that should have high diversity; lower + it only for content that is legitimately formulaic. + + Raises: + ValueError: If the unique-block ratio falls below ``min_unique_ratio``. + """ + + if blocks: + unique_ratio = len(set(blocks)) / len(blocks) + if unique_ratio < min_unique_ratio: + raise ValueError( + f"{path}: only {unique_ratio:.1%} of {len(blocks)} blocks are unique " + f"(minimum required: {min_unique_ratio:.0%}). Widen the source " + "vocabulary/axes in the generator instead of shipping a " + "duplicate-heavy file." + ) + + path.parent.mkdir(parents=True, exist_ok=True) + 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 new file mode 100644 index 0000000..1b5305a --- /dev/null +++ b/tools/curriculum_subjects.py @@ -0,0 +1,396 @@ +from __future__ import annotations + +try: + from .curriculum_shared import * +except ImportError: + from curriculum_shared import * + +def language_blocks(count: int) -> list[str]: + """Create language teaching blocks.""" + + subjects = [ + "Mina", "Ravi", "Lena", "Omar", "Sara", "Tara", "Jin", "Ada", + "Leo", "Nia", "Sam", "Priya", "Yuki", "Noah", "Ines", "Kofi", + ] + verbs = [ + "reads", "writes", "asks", "answers", "walks", "listens", + "draws", "explains", "counts", "builds", "shares", "practices", + ] + objects = [ + "a book", "a note", "a question", "a sentence", "a story", + "a message", "a poem", "a list", "a letter", "a riddle", + ] + qualities = [ + "clear", "short", "kind", "useful", "simple", "careful", + "honest", "tidy", "friendly", "direct", + ] + templates = [ + ( + "{subject} {verb} {obj}.\n" + "The sentence has a subject.\n" + "The subject tells who acts.\n" + "The verb tells what happens.\n" + "The object receives the action.\n" + "A {quality} sentence is easy to read.\n" + "What does {subject} do?\n" + "{subject} {verb} {obj}.\n" + "Rewrite the idea with fewer words.\n" + "{subject} {verb}." + ), + ( + "Today, {subject} {verb} {obj}.\n" + "Notice the subject and the verb.\n" + "{subject} is the subject; {verb} is the verb.\n" + "{obj} is the object of the sentence.\n" + "A {quality} sentence keeps its meaning obvious.\n" + "Try shortening it: {subject} {verb}.\n" + "Now expand it again with a detail of your own." + ), + ( + "Sentence practice: {subject} {verb} {obj}.\n" + "Underline the subject, then the verb, then the object.\n" + "A {quality} sentence usually has one clear idea.\n" + "Ask: who is acting? {subject}.\n" + "Ask: what do they do? {verb}.\n" + "Ask: what receives it? {obj}." + ), + ] + period = combinatorial_period(len(subjects), len(verbs), len(objects), len(qualities), len(templates)) + blocks = [] + for index in range(min(count, period)): + s_i, v_i, o_i, q_i, t_i = mixed_radix_pick( + index, len(subjects), len(verbs), len(objects), len(qualities), len(templates) + ) + blocks.append( + templates[t_i].format( + subject=subjects[s_i], verb=verbs[v_i], obj=objects[o_i], quality=qualities[q_i] + ) + ) + return blocks + + +def math_blocks(count: int) -> list[str]: + """Create math teaching blocks.""" + + nouns = [ + "pencils", "marbles", "stickers", "apples", "coins", "buttons", + "stones", "cards", "stamps", "beads", + ] + templates = [ + ( + "A box has {a} {noun}.\n" + "Another box has {b} {noun}.\n" + "{a} plus {b} equals {total}.\n" + "Together the boxes have {total} {noun}.\n" + "If there are {a} groups of {b}, multiply.\n" + "{a} times {b} equals {product}.\n" + "Addition joins amounts.\n" + "Multiplication joins equal groups.\n" + "Check the answer by counting carefully." + ), + ( + "You start with {a} {noun}.\n" + "You receive {b} more {noun}.\n" + "How many {noun} in total? Add {a} and {b}.\n" + "{a} + {b} = {total}.\n" + "If instead you had {a} equal groups of {b} {noun} each, multiply.\n" + "{a} x {b} = {product}.\n" + "Addition combines amounts; multiplication combines equal groups." + ), + ( + "There are {a} {noun} in one pile and {b} {noun} in another.\n" + "Combined, that is {total} {noun} ({a} + {b} = {total}).\n" + "If you arranged {a} rows of {b} {noun}, the total by multiplication is {product}.\n" + "Recount to double check: does {total} match, and does {product} match?" + ), + ] + a_values = list(range(2, 302)) + b_values = list(range(1, 201)) + period = combinatorial_period(len(a_values), len(b_values), len(nouns), len(templates)) + blocks = [] + for index in range(min(count, period)): + a_i, b_i, n_i, t_i = mixed_radix_pick(index, len(a_values), len(b_values), len(nouns), len(templates)) + a, b, noun = a_values[a_i], b_values[b_i], nouns[n_i] + blocks.append( + templates[t_i].format(a=a, b=b, noun=noun, total=a + b, product=a * b) + ) + return blocks + + +def science_blocks(count: int) -> list[str]: + """Create science teaching blocks.""" + + topics = [ + ("plant", "roots take water from soil", "leaves use sunlight"), + ("heart", "the heart pumps blood", "blood carries oxygen"), + ("battery", "a battery stores energy", "a wire can carry electricity"), + ("cloud", "warm air can hold water vapor", "cool air can form clouds"), + ("magnet", "a magnet pulls some metals", "iron is attracted to magnets"), + ("moon", "the Moon moves around Earth", "moonlight is reflected sunlight"), + ("volcano", "melted rock rises from below", "pressure can cause an eruption"), + ("river", "water flows from high to low ground", "rivers carry sediment downstream"), + ("seed", "a seed holds a tiny plant", "water and warmth help it sprout"), + ("lightning", "charge can build up in clouds", "a spark jumps between charges"), + ("skeleton", "bones support the body", "joints let bones move"), + ("sound", "sound travels as vibrations", "vibrations move through air"), + ("mirror", "a mirror reflects light", "the reflected image looks reversed"), + ("compass", "a compass needle is a small magnet", "it points toward magnetic north"), + ("insect", "many insects have six legs", "some insects go through metamorphosis"), + ("ice", "water freezes at zero degrees Celsius", "ice is less dense than liquid water"), + ("gravity", "gravity pulls objects toward Earth", "heavier objects still fall at the same rate"), + ("photosynthesis", "plants use sunlight to make food", "the process also releases oxygen"), + ("erosion", "wind and water wear down rock over time", "erosion can reshape landscapes slowly"), + ("circuit", "a circuit needs a complete loop", "a broken loop stops the current"), + ] + period = combinatorial_period(len(topics), len(NAMES), len(MODIFIERS), len(CLOSERS)) + blocks = [] + for index in range(min(count, period)): + topic_i, name_i, mod_i, close_i = mixed_radix_pick( + index, len(topics), len(NAMES), len(MODIFIERS), len(CLOSERS) + ) + name, fact_one, fact_two = topics[topic_i] + student, modifier, closer = NAMES[name_i], MODIFIERS[mod_i], CLOSERS[close_i] + blocks.append( + f"{student} studies a {name}.\n" + f"{student} observes {modifier}.\n" + f"{fact_one.capitalize()}.\n" + f"{fact_two.capitalize()}.\n" + f"An observation tells what we notice.\n" + f"A question asks why it happens.\n" + f"A test can compare two cases.\n" + f"{closer}" + ) + return blocks + + +def geography_history_blocks(count: int) -> list[str]: + """Create geography and history teaching blocks.""" + + places = [ + ("India", "New Delhi", "Asia", "the Himalayas"), + ("Egypt", "Cairo", "Africa", "the Nile River"), + ("Japan", "Tokyo", "Asia", "many islands"), + ("France", "Paris", "Europe", "the Seine River"), + ("Brazil", "Brasilia", "South America", "the Amazon region"), + ("Kenya", "Nairobi", "Africa", "the Great Rift Valley"), + ("Canada", "Ottawa", "North America", "vast northern forests"), + ("Australia", "Canberra", "Oceania", "large desert interior"), + ("Peru", "Lima", "South America", "the Andes mountains"), + ("Norway", "Oslo", "Europe", "deep coastal fjords"), + ("Vietnam", "Hanoi", "Asia", "the Mekong Delta"), + ("Morocco", "Rabat", "Africa", "the Atlas Mountains"), + ("Mexico", "Mexico City", "North America", "central highland valleys"), + ("Turkey", "Ankara", "Europe/Asia", "the Bosphorus strait"), + ("Chile", "Santiago", "South America", "the Atacama Desert"), + ] + inventions = [ + "the wheel", "writing", "the compass", "the printing press", + "the steam engine", "the telescope", "the telegraph", "the light bulb", + ] + period = combinatorial_period(len(places), len(inventions), len(NAMES), len(CLOSERS)) + blocks = [] + for index in range(min(count, period)): + place_i, inv_i, name_i, close_i = mixed_radix_pick( + index, len(places), len(inventions), len(NAMES), len(CLOSERS) + ) + country, capital, continent, feature = places[place_i] + invention, student, closer = inventions[inv_i], NAMES[name_i], CLOSERS[close_i] + blocks.append( + f"{student} is learning about {country}.\n" + f"{country} is in {continent}.\n" + f"The capital city is {capital}.\n" + f"A map can show where {country} is.\n" + f"One known feature is {feature}.\n" + f"People in each place have culture.\n" + f"Culture includes food, language, music, and customs.\n" + f"History studies change over time.\n" + f"An important invention was {invention}.\n" + f"{closer}" + ) + return blocks + + +def reasoning_blocks(count: int) -> list[str]: + """Create reasoning teaching blocks.""" + + people = ["Tom", "Mina", "Ravi", "Lena", "Omar", "Sara", "Tara", "Jin"] + items = ["apples", "marbles", "coins", "stickers", "pencils", "cards"] + templates = [ + ( + "{person} has {a} {item}.\n" + "{person} gets {b} more {item}.\n" + "To find the total, add.\n" + "{a} plus {b} equals {total}.\n" + "{person} has {total} {item}.\n" + "If the number goes up, addition may help.\n" + "If the number goes down, subtraction may help.\n" + "Choose the operation from the story." + ), + ( + "{person} starts with {a} {item} and gives away {b}.\n" + "{a} minus {b} equals {diff}.\n" + "{person} now has {diff} {item}.\n" + "Watch the wording: 'gives away' signals subtraction.\n" + "Reread the story before picking an operation." + ), + ] + a_values = list(range(2, 101)) + b_values = list(range(1, 61)) + period = combinatorial_period(len(people), len(items), len(a_values), len(b_values), len(templates)) + blocks = [] + for index in range(min(count, period)): + p_i, i_i, a_i, b_i, t_i = mixed_radix_pick( + index, len(people), len(items), len(a_values), len(b_values), len(templates) + ) + person, item, a, b = people[p_i], items[i_i], a_values[a_i], b_values[b_i] + diff = max(a, b) - min(a, b) + if diff == 0: + diff = 1 # avoid a degenerate "gives away everything" sentence + blocks.append( + templates[t_i].format(person=person, item=item, a=a, b=b, total=a + b, diff=diff) + ) + return blocks + + +def social_blocks(count: int) -> list[str]: + """Create emotion and social reasoning blocks.""" + + feelings = [ + ("sad", "her toy broke", "a friend helps her fix it"), + ("proud", "he finished a hard task", "his practice helped"), + ("worried", "the room is dark", "she turns on a light"), + ("angry", "someone took his pencil", "he asks for it back calmly"), + ("happy", "the class works together", "teamwork feels good"), + ("nervous", "she has a test tomorrow", "a short review calms her down"), + ("embarrassed", "he tripped in front of others", "a friend jokes kindly and moves on"), + ("excited", "her team scored a goal", "she cheers for her teammates"), + ("frustrated", "the puzzle piece will not fit", "he takes a short break and tries again"), + ("lonely", "his friend moved away", "he writes a letter to stay in touch"), + ("grateful", "a neighbor helped carry groceries", "she says thank you"), + ("confused", "the instructions were unclear", "he asks a clarifying question"), + ] + period = combinatorial_period(len(feelings), len(NAMES), len(MODIFIERS), len(CLOSERS)) + blocks = [] + for index in range(min(count, period)): + feel_i, name_i, mod_i, close_i = mixed_radix_pick( + index, len(feelings), len(NAMES), len(MODIFIERS), len(CLOSERS) + ) + feeling, cause, response = feelings[feel_i] + student, modifier, closer = NAMES[name_i], MODIFIERS[mod_i], CLOSERS[close_i] + blocks.append( + f"{student} feels {feeling} because {cause}.\n" + f"A feeling often has a cause.\n" + f"The feeling can change.\n" + f"Then {response}.\n" + f"{student} handles it {modifier}.\n" + f"Kind words can help people feel safe.\n" + f"Listening shows respect.\n" + f"{closer}" + ) + return blocks + + +def everyday_blocks(count: int) -> list[str]: + """Create everyday knowledge and ethics blocks.""" + + tasks = [ + ("cook rice", "wash the rice", "turn off the stove"), + ("cross a road", "look both ways", "wait for vehicles to stop"), + ("save money", "count income", "spend less than you earn"), + ("clean a room", "put sharp things away", "wipe wet floors"), + ("visit a doctor", "explain symptoms", "follow safe advice"), + ("pack a bag", "list what is needed", "check the list before leaving"), + ("plant a garden", "prepare the soil", "water on a regular schedule"), + ("fix a flat tire", "find a safe spot to stop", "use the right tools carefully"), + ("write a budget", "list all expenses", "compare expenses to income"), + ("host a guest", "prepare a clean space", "ask about any needs in advance"), + ("borrow an item", "ask permission first", "return it in good condition"), + ("resolve a disagreement", "listen to the other side", "look for a fair compromise"), + ] + period = combinatorial_period(len(tasks), len(NAMES), len(MODIFIERS), len(CLOSERS)) + blocks = [] + for index in range(min(count, period)): + task_i, name_i, mod_i, close_i = mixed_radix_pick( + index, len(tasks), len(NAMES), len(MODIFIERS), len(CLOSERS) + ) + task, first, safe = tasks[task_i] + person, modifier, closer = NAMES[name_i], MODIFIERS[mod_i], CLOSERS[close_i] + blocks.append( + f"{person} needs to {task}.\n" + f"First, {person} {modifier} {first}.\n" + f"Good planning avoids mistakes.\n" + f"{person} should {safe}.\n" + f"Safety protects people.\n" + f"Responsibility means doing what should be done.\n" + f"Honesty and care help a community.\n" + f"{closer}" + ) + return blocks + + +def computer_blocks(count: int) -> list[str]: + """Create computer science teaching blocks.""" + + ideas = [ + ("variable", "stores a value", "x = 5"), + ("loop", "repeats steps", "for item in items"), + ("function", "groups steps", "def add(a, b)"), + ("list", "keeps items in order", "numbers = [1, 2, 3]"), + ("dictionary", "connects keys to values", "scores = {'Mina': 9}"), + ("algorithm", "is a set of steps", "sort the numbers"), + ("conditional", "chooses a path based on a check", "if score > 50"), + ("recursion", "calls itself on a smaller case", "factorial(n - 1)"), + ("class", "groups data and behavior", "class Counter"), + ("array index", "points to one item's position", "items[0]"), + ("boolean", "is either true or false", "is_ready = True"), + ("string", "stores text", "name = 'Mina'"), + ] + period = combinatorial_period(len(ideas), len(NAMES), len(MODIFIERS), len(CLOSERS)) + blocks = [] + for index in range(min(count, period)): + idea_i, name_i, mod_i, close_i = mixed_radix_pick( + index, len(ideas), len(NAMES), len(MODIFIERS), len(CLOSERS) + ) + idea, meaning, example = ideas[idea_i] + student, modifier, closer = NAMES[name_i], MODIFIERS[mod_i], CLOSERS[close_i] + blocks.append( + f"A {idea} {meaning}.\n" + f"Example: {example}.\n" + f"{student} reads the code {modifier}.\n" + f"A programmer reads errors carefully.\n" + f"Debugging means finding the cause.\n" + f"Testing checks if code works.\n" + f"Small steps make hard problems easier.\n" + f"{closer}" + ) + return blocks + + +def code_blocks(count: int) -> list[str]: + """Create code explanation corpus blocks.""" + + operations = [ + ("+", "adds", lambda x, y: x + y), + ("-", "subtracts", lambda x, y: x - y), + ("*", "multiplies", lambda x, y: x * y), + ] + period = combinatorial_period(300, 200, len(operations)) + blocks = [] + for index in range(min(count, period)): + x_i, y_i, op_i = mixed_radix_pick(index, 300, 200, len(operations)) + value, other = x_i + 1, y_i + 1 + symbol, verb, func = operations[op_i] + blocks.append( + "Python example.\n" + f"x = {value}\n" + f"y = {other}\n" + f"print(x {symbol} y)\n" + f"x stores {value}.\n" + f"y stores {other}.\n" + f"The {symbol} sign {verb} the numbers.\n" + f"The program prints {func(value, other)}.\n" + "This example teaches variables and arithmetic." + ) + return blocks + + diff --git a/tools/generate_default_curriculum.py b/tools/generate_default_curriculum.py index 3031b7e..a8d7c0f 100644 --- a/tools/generate_default_curriculum.py +++ b/tools/generate_default_curriculum.py @@ -1,1355 +1,39 @@ from __future__ import annotations -import random from pathlib import Path - -ROOT = Path(__file__).resolve().parents[1] / "llm_trainer" / "default_data" - -# Shared secondary axes reused across the small "conceptual" categories below -# (science, geography, social, everyday, computer science, language). They -# exist purely to multiply the combinatorial space so a category built from a -# modest, hand-written topic list still produces thousands of genuinely -# distinct blocks instead of the same handful repeated on a loop. -NAMES = [ - "Mina", "Ravi", "Lena", "Omar", "Sara", "Tara", "Jin", "Ada", - "Leo", "Nia", "Sam", "Priya", -] -MODIFIERS = ["carefully", "quickly", "calmly", "clearly", "patiently"] -CLOSERS = [ - "What did you notice?", - "Why do you think that happens?", - "How could you check this?", - "What would change if one detail changed?", - "Explain it back in your own words.", -] - - -def mixed_radix_pick(index: int, *sizes: int) -> list[int]: - """Decompose an index into independent per-axis picks. - - Unlike applying ``index % len(list)`` to several lists at once (which - repeats after ``lcm`` of the list lengths -- often a tiny number), this - treats ``index`` as a mixed-radix counter across every axis. The combined - period is the *product* of all axis sizes, so a handful of modest lists - (say four lists of 15-20 items) already yields a combinatorial space of - tens of thousands of unique combinations before anything repeats. - - Args: - index: Zero-based block index. - *sizes: Length of each axis, in the same order picks are needed. - - Returns: - One pick per axis, each in ``range(0, size)``. - """ - - picks = [] - remaining = index - for size in sizes: - size = max(1, size) - picks.append(remaining % size) - remaining //= size - return picks - - -def combinatorial_period(*sizes: int) -> int: - """Return the number of unique combinations `mixed_radix_pick` can produce.""" - - period = 1 - for size in sizes: - period *= max(1, size) - return period - - -def write_blocks(path: Path, blocks: list[str], min_unique_ratio: float = 0.9) -> None: - """Write plain-text corpus blocks, guarding against templated duplication. - - A generator that technically returns ``count`` blocks but only cycles - through a handful of unique strings silently produces a dataset that is - almost entirely duplicate data -- wasted disk, wasted training compute, - and a validation split that can't mean anything because train and - validation end up full of the same repeated content. This raises loudly - instead of writing a file that *looks* like a real corpus but isn't. - - Args: - path: Output text file. - blocks: Corpus blocks. - min_unique_ratio: Minimum allowed fraction of unique blocks. Raise - the ratio for categories that should have high diversity; lower - it only for content that is legitimately formulaic. - - Raises: - ValueError: If the unique-block ratio falls below ``min_unique_ratio``. - """ - - if blocks: - unique_ratio = len(set(blocks)) / len(blocks) - if unique_ratio < min_unique_ratio: - raise ValueError( - f"{path}: only {unique_ratio:.1%} of {len(blocks)} blocks are unique " - f"(minimum required: {min_unique_ratio:.0%}). Widen the source " - "vocabulary/axes in the generator instead of shipping a " - "duplicate-heavy file." - ) - - path.parent.mkdir(parents=True, exist_ok=True) - 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) - - -def language_blocks(count: int) -> list[str]: - """Create language teaching blocks.""" - - subjects = [ - "Mina", "Ravi", "Lena", "Omar", "Sara", "Tara", "Jin", "Ada", - "Leo", "Nia", "Sam", "Priya", "Yuki", "Noah", "Ines", "Kofi", - ] - verbs = [ - "reads", "writes", "asks", "answers", "walks", "listens", - "draws", "explains", "counts", "builds", "shares", "practices", - ] - objects = [ - "a book", "a note", "a question", "a sentence", "a story", - "a message", "a poem", "a list", "a letter", "a riddle", - ] - qualities = [ - "clear", "short", "kind", "useful", "simple", "careful", - "honest", "tidy", "friendly", "direct", - ] - templates = [ - ( - "{subject} {verb} {obj}.\n" - "The sentence has a subject.\n" - "The subject tells who acts.\n" - "The verb tells what happens.\n" - "The object receives the action.\n" - "A {quality} sentence is easy to read.\n" - "What does {subject} do?\n" - "{subject} {verb} {obj}.\n" - "Rewrite the idea with fewer words.\n" - "{subject} {verb}." - ), - ( - "Today, {subject} {verb} {obj}.\n" - "Notice the subject and the verb.\n" - "{subject} is the subject; {verb} is the verb.\n" - "{obj} is the object of the sentence.\n" - "A {quality} sentence keeps its meaning obvious.\n" - "Try shortening it: {subject} {verb}.\n" - "Now expand it again with a detail of your own." - ), - ( - "Sentence practice: {subject} {verb} {obj}.\n" - "Underline the subject, then the verb, then the object.\n" - "A {quality} sentence usually has one clear idea.\n" - "Ask: who is acting? {subject}.\n" - "Ask: what do they do? {verb}.\n" - "Ask: what receives it? {obj}." - ), - ] - period = combinatorial_period(len(subjects), len(verbs), len(objects), len(qualities), len(templates)) - blocks = [] - for index in range(min(count, period)): - s_i, v_i, o_i, q_i, t_i = mixed_radix_pick( - index, len(subjects), len(verbs), len(objects), len(qualities), len(templates) - ) - blocks.append( - templates[t_i].format( - subject=subjects[s_i], verb=verbs[v_i], obj=objects[o_i], quality=qualities[q_i] - ) - ) - return blocks - - -def math_blocks(count: int) -> list[str]: - """Create math teaching blocks.""" - - nouns = [ - "pencils", "marbles", "stickers", "apples", "coins", "buttons", - "stones", "cards", "stamps", "beads", - ] - templates = [ - ( - "A box has {a} {noun}.\n" - "Another box has {b} {noun}.\n" - "{a} plus {b} equals {total}.\n" - "Together the boxes have {total} {noun}.\n" - "If there are {a} groups of {b}, multiply.\n" - "{a} times {b} equals {product}.\n" - "Addition joins amounts.\n" - "Multiplication joins equal groups.\n" - "Check the answer by counting carefully." - ), - ( - "You start with {a} {noun}.\n" - "You receive {b} more {noun}.\n" - "How many {noun} in total? Add {a} and {b}.\n" - "{a} + {b} = {total}.\n" - "If instead you had {a} equal groups of {b} {noun} each, multiply.\n" - "{a} x {b} = {product}.\n" - "Addition combines amounts; multiplication combines equal groups." - ), - ( - "There are {a} {noun} in one pile and {b} {noun} in another.\n" - "Combined, that is {total} {noun} ({a} + {b} = {total}).\n" - "If you arranged {a} rows of {b} {noun}, the total by multiplication is {product}.\n" - "Recount to double check: does {total} match, and does {product} match?" - ), - ] - a_values = list(range(2, 302)) - b_values = list(range(1, 201)) - period = combinatorial_period(len(a_values), len(b_values), len(nouns), len(templates)) - blocks = [] - for index in range(min(count, period)): - a_i, b_i, n_i, t_i = mixed_radix_pick(index, len(a_values), len(b_values), len(nouns), len(templates)) - a, b, noun = a_values[a_i], b_values[b_i], nouns[n_i] - blocks.append( - templates[t_i].format(a=a, b=b, noun=noun, total=a + b, product=a * b) - ) - return blocks - - -def science_blocks(count: int) -> list[str]: - """Create science teaching blocks.""" - - topics = [ - ("plant", "roots take water from soil", "leaves use sunlight"), - ("heart", "the heart pumps blood", "blood carries oxygen"), - ("battery", "a battery stores energy", "a wire can carry electricity"), - ("cloud", "warm air can hold water vapor", "cool air can form clouds"), - ("magnet", "a magnet pulls some metals", "iron is attracted to magnets"), - ("moon", "the Moon moves around Earth", "moonlight is reflected sunlight"), - ("volcano", "melted rock rises from below", "pressure can cause an eruption"), - ("river", "water flows from high to low ground", "rivers carry sediment downstream"), - ("seed", "a seed holds a tiny plant", "water and warmth help it sprout"), - ("lightning", "charge can build up in clouds", "a spark jumps between charges"), - ("skeleton", "bones support the body", "joints let bones move"), - ("sound", "sound travels as vibrations", "vibrations move through air"), - ("mirror", "a mirror reflects light", "the reflected image looks reversed"), - ("compass", "a compass needle is a small magnet", "it points toward magnetic north"), - ("insect", "many insects have six legs", "some insects go through metamorphosis"), - ("ice", "water freezes at zero degrees Celsius", "ice is less dense than liquid water"), - ("gravity", "gravity pulls objects toward Earth", "heavier objects still fall at the same rate"), - ("photosynthesis", "plants use sunlight to make food", "the process also releases oxygen"), - ("erosion", "wind and water wear down rock over time", "erosion can reshape landscapes slowly"), - ("circuit", "a circuit needs a complete loop", "a broken loop stops the current"), - ] - period = combinatorial_period(len(topics), len(NAMES), len(MODIFIERS), len(CLOSERS)) - blocks = [] - for index in range(min(count, period)): - topic_i, name_i, mod_i, close_i = mixed_radix_pick( - index, len(topics), len(NAMES), len(MODIFIERS), len(CLOSERS) - ) - name, fact_one, fact_two = topics[topic_i] - student, modifier, closer = NAMES[name_i], MODIFIERS[mod_i], CLOSERS[close_i] - blocks.append( - f"{student} studies a {name}.\n" - f"{student} observes {modifier}.\n" - f"{fact_one.capitalize()}.\n" - f"{fact_two.capitalize()}.\n" - f"An observation tells what we notice.\n" - f"A question asks why it happens.\n" - f"A test can compare two cases.\n" - f"{closer}" - ) - return blocks - - -def geography_history_blocks(count: int) -> list[str]: - """Create geography and history teaching blocks.""" - - places = [ - ("India", "New Delhi", "Asia", "the Himalayas"), - ("Egypt", "Cairo", "Africa", "the Nile River"), - ("Japan", "Tokyo", "Asia", "many islands"), - ("France", "Paris", "Europe", "the Seine River"), - ("Brazil", "Brasilia", "South America", "the Amazon region"), - ("Kenya", "Nairobi", "Africa", "the Great Rift Valley"), - ("Canada", "Ottawa", "North America", "vast northern forests"), - ("Australia", "Canberra", "Oceania", "large desert interior"), - ("Peru", "Lima", "South America", "the Andes mountains"), - ("Norway", "Oslo", "Europe", "deep coastal fjords"), - ("Vietnam", "Hanoi", "Asia", "the Mekong Delta"), - ("Morocco", "Rabat", "Africa", "the Atlas Mountains"), - ("Mexico", "Mexico City", "North America", "central highland valleys"), - ("Turkey", "Ankara", "Europe/Asia", "the Bosphorus strait"), - ("Chile", "Santiago", "South America", "the Atacama Desert"), - ] - inventions = [ - "the wheel", "writing", "the compass", "the printing press", - "the steam engine", "the telescope", "the telegraph", "the light bulb", - ] - period = combinatorial_period(len(places), len(inventions), len(NAMES), len(CLOSERS)) - blocks = [] - for index in range(min(count, period)): - place_i, inv_i, name_i, close_i = mixed_radix_pick( - index, len(places), len(inventions), len(NAMES), len(CLOSERS) - ) - country, capital, continent, feature = places[place_i] - invention, student, closer = inventions[inv_i], NAMES[name_i], CLOSERS[close_i] - blocks.append( - f"{student} is learning about {country}.\n" - f"{country} is in {continent}.\n" - f"The capital city is {capital}.\n" - f"A map can show where {country} is.\n" - f"One known feature is {feature}.\n" - f"People in each place have culture.\n" - f"Culture includes food, language, music, and customs.\n" - f"History studies change over time.\n" - f"An important invention was {invention}.\n" - f"{closer}" - ) - return blocks - - -def reasoning_blocks(count: int) -> list[str]: - """Create reasoning teaching blocks.""" - - people = ["Tom", "Mina", "Ravi", "Lena", "Omar", "Sara", "Tara", "Jin"] - items = ["apples", "marbles", "coins", "stickers", "pencils", "cards"] - templates = [ - ( - "{person} has {a} {item}.\n" - "{person} gets {b} more {item}.\n" - "To find the total, add.\n" - "{a} plus {b} equals {total}.\n" - "{person} has {total} {item}.\n" - "If the number goes up, addition may help.\n" - "If the number goes down, subtraction may help.\n" - "Choose the operation from the story." - ), - ( - "{person} starts with {a} {item} and gives away {b}.\n" - "{a} minus {b} equals {diff}.\n" - "{person} now has {diff} {item}.\n" - "Watch the wording: 'gives away' signals subtraction.\n" - "Reread the story before picking an operation." - ), - ] - a_values = list(range(2, 101)) - b_values = list(range(1, 61)) - period = combinatorial_period(len(people), len(items), len(a_values), len(b_values), len(templates)) - blocks = [] - for index in range(min(count, period)): - p_i, i_i, a_i, b_i, t_i = mixed_radix_pick( - index, len(people), len(items), len(a_values), len(b_values), len(templates) - ) - person, item, a, b = people[p_i], items[i_i], a_values[a_i], b_values[b_i] - diff = max(a, b) - min(a, b) - if diff == 0: - diff = 1 # avoid a degenerate "gives away everything" sentence - blocks.append( - templates[t_i].format(person=person, item=item, a=a, b=b, total=a + b, diff=diff) - ) - return blocks - - -def social_blocks(count: int) -> list[str]: - """Create emotion and social reasoning blocks.""" - - feelings = [ - ("sad", "her toy broke", "a friend helps her fix it"), - ("proud", "he finished a hard task", "his practice helped"), - ("worried", "the room is dark", "she turns on a light"), - ("angry", "someone took his pencil", "he asks for it back calmly"), - ("happy", "the class works together", "teamwork feels good"), - ("nervous", "she has a test tomorrow", "a short review calms her down"), - ("embarrassed", "he tripped in front of others", "a friend jokes kindly and moves on"), - ("excited", "her team scored a goal", "she cheers for her teammates"), - ("frustrated", "the puzzle piece will not fit", "he takes a short break and tries again"), - ("lonely", "his friend moved away", "he writes a letter to stay in touch"), - ("grateful", "a neighbor helped carry groceries", "she says thank you"), - ("confused", "the instructions were unclear", "he asks a clarifying question"), - ] - period = combinatorial_period(len(feelings), len(NAMES), len(MODIFIERS), len(CLOSERS)) - blocks = [] - for index in range(min(count, period)): - feel_i, name_i, mod_i, close_i = mixed_radix_pick( - index, len(feelings), len(NAMES), len(MODIFIERS), len(CLOSERS) - ) - feeling, cause, response = feelings[feel_i] - student, modifier, closer = NAMES[name_i], MODIFIERS[mod_i], CLOSERS[close_i] - blocks.append( - f"{student} feels {feeling} because {cause}.\n" - f"A feeling often has a cause.\n" - f"The feeling can change.\n" - f"Then {response}.\n" - f"{student} handles it {modifier}.\n" - f"Kind words can help people feel safe.\n" - f"Listening shows respect.\n" - f"{closer}" - ) - return blocks - - -def everyday_blocks(count: int) -> list[str]: - """Create everyday knowledge and ethics blocks.""" - - tasks = [ - ("cook rice", "wash the rice", "turn off the stove"), - ("cross a road", "look both ways", "wait for vehicles to stop"), - ("save money", "count income", "spend less than you earn"), - ("clean a room", "put sharp things away", "wipe wet floors"), - ("visit a doctor", "explain symptoms", "follow safe advice"), - ("pack a bag", "list what is needed", "check the list before leaving"), - ("plant a garden", "prepare the soil", "water on a regular schedule"), - ("fix a flat tire", "find a safe spot to stop", "use the right tools carefully"), - ("write a budget", "list all expenses", "compare expenses to income"), - ("host a guest", "prepare a clean space", "ask about any needs in advance"), - ("borrow an item", "ask permission first", "return it in good condition"), - ("resolve a disagreement", "listen to the other side", "look for a fair compromise"), - ] - period = combinatorial_period(len(tasks), len(NAMES), len(MODIFIERS), len(CLOSERS)) - blocks = [] - for index in range(min(count, period)): - task_i, name_i, mod_i, close_i = mixed_radix_pick( - index, len(tasks), len(NAMES), len(MODIFIERS), len(CLOSERS) - ) - task, first, safe = tasks[task_i] - person, modifier, closer = NAMES[name_i], MODIFIERS[mod_i], CLOSERS[close_i] - blocks.append( - f"{person} needs to {task}.\n" - f"First, {person} {modifier} {first}.\n" - f"Good planning avoids mistakes.\n" - f"{person} should {safe}.\n" - f"Safety protects people.\n" - f"Responsibility means doing what should be done.\n" - f"Honesty and care help a community.\n" - f"{closer}" - ) - return blocks - - -def computer_blocks(count: int) -> list[str]: - """Create computer science teaching blocks.""" - - ideas = [ - ("variable", "stores a value", "x = 5"), - ("loop", "repeats steps", "for item in items"), - ("function", "groups steps", "def add(a, b)"), - ("list", "keeps items in order", "numbers = [1, 2, 3]"), - ("dictionary", "connects keys to values", "scores = {'Mina': 9}"), - ("algorithm", "is a set of steps", "sort the numbers"), - ("conditional", "chooses a path based on a check", "if score > 50"), - ("recursion", "calls itself on a smaller case", "factorial(n - 1)"), - ("class", "groups data and behavior", "class Counter"), - ("array index", "points to one item's position", "items[0]"), - ("boolean", "is either true or false", "is_ready = True"), - ("string", "stores text", "name = 'Mina'"), - ] - period = combinatorial_period(len(ideas), len(NAMES), len(MODIFIERS), len(CLOSERS)) - blocks = [] - for index in range(min(count, period)): - idea_i, name_i, mod_i, close_i = mixed_radix_pick( - index, len(ideas), len(NAMES), len(MODIFIERS), len(CLOSERS) - ) - idea, meaning, example = ideas[idea_i] - student, modifier, closer = NAMES[name_i], MODIFIERS[mod_i], CLOSERS[close_i] - blocks.append( - f"A {idea} {meaning}.\n" - f"Example: {example}.\n" - f"{student} reads the code {modifier}.\n" - f"A programmer reads errors carefully.\n" - f"Debugging means finding the cause.\n" - f"Testing checks if code works.\n" - f"Small steps make hard problems easier.\n" - f"{closer}" - ) - return blocks - - -def code_blocks(count: int) -> list[str]: - """Create code explanation corpus blocks.""" - - operations = [ - ("+", "adds", lambda x, y: x + y), - ("-", "subtracts", lambda x, y: x - y), - ("*", "multiplies", lambda x, y: x * y), - ] - period = combinatorial_period(300, 200, len(operations)) - blocks = [] - for index in range(min(count, period)): - x_i, y_i, op_i = mixed_radix_pick(index, 300, 200, len(operations)) - value, other = x_i + 1, y_i + 1 - symbol, verb, func = operations[op_i] - blocks.append( - "Python example.\n" - f"x = {value}\n" - f"y = {other}\n" - f"print(x {symbol} y)\n" - f"x stores {value}.\n" - f"y stores {other}.\n" - f"The {symbol} sign {verb} the numbers.\n" - f"The program prints {func(value, other)}.\n" - "This example teaches variables and arithmetic." - ) - return blocks - - -def programming_deep_blocks(count: int, topic: str) -> list[str]: - """Create programming-focused corpus blocks. - - Args: - count: Number of blocks to generate. - topic: Programming topic name. - - Returns: - Generated corpus blocks. - """ - - python_examples = [ - ( - "Python list filtering", - "numbers = [1, 2, 3, 4, 5]\n" - "even = []\n" - "for number in numbers:\n" - " if number % 2 == 0:\n" - " even.append(number)\n" - "print(even)", - "The list stores numbers in order.\nThe loop checks each number.\nThe percent operator gives the remainder.\nA remainder of zero means the number is even.", - ), - ( - "Python function", - "def area(width, height):\n" - " return width * height\n\n" - "result = area(6, 4)\n" - "print(result)", - "The function receives width and height.\nThe return statement sends back the answer.\nThe area is twenty four.", - ), - ( - "Python dictionary", - "scores = {'Mina': 8, 'Ravi': 9}\n" - "scores['Lena'] = 7\n" - "for name, score in scores.items():\n" - " print(name, score)", - "A dictionary maps keys to values.\nThe key is a name.\nThe value is a score.\nThe items method gives key and value pairs.", - ), - ] - javascript_examples = [ - ( - "JavaScript array map", - "const prices = [10, 20, 30];\n" - "const doubled = prices.map(price => price * 2);\n" - "console.log(doubled);", - "An array keeps values in order.\nThe map method creates a new array.\nThe arrow function runs once for each value.", - ), - ( - "JavaScript async function", - "async function loadUser(id) {\n" - " const response = await fetch(`/users/${id}`);\n" - " return response.json();\n" - "}", - "The async keyword allows await.\nAwait pauses until the promise settles.\nThis is useful for network requests.", - ), - ] - java_examples = [ - ( - "Java class", - "class Counter {\n" - " private int value = 0;\n" - " void increment() {\n" - " value++;\n" - " }\n" - " int getValue() {\n" - " return value;\n" - " }\n" - "}", - "A class groups data and behavior.\nThe field stores the count.\nThe method changes the count.\nPrivate data is hidden from outside code.", - ), - ] - cpp_examples = [ - ( - "C++ vector loop", - "#include \n" - "#include \n\n" - "int main() {\n" - " std::vector values{1, 2, 3};\n" - " int total = 0;\n" - " for (int value : values) {\n" - " total += value;\n" - " }\n" - " std::cout << total << '\\n';\n" - "}", - "A vector stores many values.\nThe range loop visits each value.\nThe total variable accumulates the sum.", - ), - ( - "C pointer safety", - "#include \n\n" - "int main(void) {\n" - " int value = 5;\n" - " int *ptr = &value;\n" - " printf(\"%d\\n\", *ptr);\n" - " return 0;\n" - "}", - "A pointer stores an address.\nThe address operator gets the address.\nThe star operator reads the value at the address.", - ), - ] - rust_go_examples = [ - ( - "Rust ownership", - "fn main() {\n" - " let name = String::from(\"Mina\");\n" - " print_name(&name);\n" - " println!(\"{}\", name);\n" - "}\n\n" - "fn print_name(value: &String) {\n" - " println!(\"{}\", value);\n" - "}", - "The ampersand borrows the string.\nBorrowing lets a function read without taking ownership.\nThe original value can still be used later.", - ), - ( - "Go error handling", - "file, err := os.Open(\"data.txt\")\n" - "if err != nil {\n" - " return err\n" - "}\n" - "defer file.Close()", - "Go returns errors as values.\nThe code checks the error immediately.\nThe defer statement closes the file later.", - ), - ] - sql_shell_examples = [ - ( - "SQL selection", - "SELECT name, age\n" - "FROM users\n" - "WHERE age >= 18\n" - "ORDER BY name;", - "The SELECT clause chooses columns.\nThe FROM clause chooses a table.\nThe WHERE clause filters rows.\nThe ORDER BY clause sorts the result.", - ), - ( - "Bash pipeline", - "cat access.log | grep ERROR | sort | uniq -c", - "A pipeline sends output to the next command.\nGrep filters matching lines.\nSort groups similar lines.\nUniq counts repeated lines.", - ), - ( - "PowerShell pipeline", - "Get-ChildItem -File | Where-Object { $_.Length -gt 1MB } | Select-Object Name, Length", - "PowerShell passes objects through the pipeline.\nWhere-Object filters objects.\nSelect-Object chooses properties to display.", - ), - ] - web_examples = [ - ( - "HTML form", - "
    \n" - " \n" - " \n" - " \n" - "
    ", - "A form collects input.\nA label tells the user what to enter.\nA button submits the form.", - ), - ( - "CSS button", - ".button {\n" - " background: #222;\n" - " color: white;\n" - " padding: 8px 12px;\n" - "}\n" - ".button:hover {\n" - " background: #444;\n" - "}", - "CSS changes how elements look.\nThe hover rule runs when the pointer is over the button.", - ), - ] - algorithm_examples = [ - ( - "Binary search", - "def binary_search(values, target):\n" - " low = 0\n" - " high = len(values) - 1\n" - " while low <= high:\n" - " mid = (low + high) // 2\n" - " if values[mid] == target:\n" - " return mid\n" - " if values[mid] < target:\n" - " low = mid + 1\n" - " else:\n" - " high = mid - 1\n" - " return -1", - "Binary search works on sorted data.\nEach step removes half of the remaining choices.\nThis makes it faster than checking every item.", - ), - ( - "Queue with list", - "from collections import deque\n\n" - "queue = deque()\n" - "queue.append('first')\n" - "queue.append('second')\n" - "item = queue.popleft()\n" - "print(item)", - "A queue is first in, first out.\nAppend adds to the back.\nPopleft removes from the front.", - ), - ] - debugging_examples = [ - ( - "Read the traceback", - "Traceback says the error line.\nStart at the last line.\nFind the exception name.\nThen inspect the code near that line.\nA NameError often means a variable name is missing or misspelled.", - "Debugging starts with evidence.\nDo not guess first.\nRead the error.\nReproduce the bug.\nChange one thing.\nRun the test again.", - ), - ( - "Off by one error", - "for index in range(len(items)):\n" - " print(items[index])", - "Indexes start at zero in many languages.\nThe last index is length minus one.\nAn off by one error reads before the start or after the end.", - ), - ] - sets = { - "python": python_examples, - "javascript_web": javascript_examples + web_examples, - "java_csharp": java_examples, - "c_cpp_systems": cpp_examples, - "rust_go": rust_go_examples, - "sql_shell": sql_shell_examples, - "algorithms": algorithm_examples, - "debugging": debugging_examples, - "full_stack": web_examples + sql_shell_examples + javascript_examples, - "data_structures": algorithm_examples + python_examples, - "software_engineering": debugging_examples + java_examples + rust_go_examples, - "mixed_language": python_examples + javascript_examples + cpp_examples + rust_go_examples + sql_shell_examples, - } - examples = sets[topic] - blocks = [] - for index in range(count): - title, code, explanation = examples[index % len(examples)] - scenario = index % 11 - blocks.append( - f"{title}.\n" - f"Example number {index + 1}.\n" - f"{code}\n" - f"{explanation}\n" - f"The programmer should name variables clearly.\n" - f"The program should handle expected input.\n" - f"The program should fail clearly when input is wrong.\n" - f"A small test should check the normal case.\n" - f"A second test should check an edge case.\n" - f"If scenario {scenario} changes, update the test first.\n" - f"Good code is readable, correct, and easy to change." - ) - return blocks - - -def conversation_fine_tune_blocks(count: int, topic: str) -> list[str]: - """Create conversation fine-tuning corpus blocks. - - Args: - count: Number of blocks to generate. - topic: Conversation scenario group. - - Returns: - Conversation training blocks. - """ - - scenarios = { - "daily_help": [ - ("planning a day", "I have school and chores today.", "Make a short list. Do the urgent task first."), - ("finding a lost key", "I lost my key.", "Check the last place you used it. Then check pockets and bags."), - ("choosing food", "I want a simple dinner.", "Rice, vegetables, and eggs can make a quick meal."), - ("packing for a trip", "I don't know what to pack.", "List essentials first: documents, clothes, chargers."), - ("running late", "I'm running late for an appointment.", "Call ahead if you can, then focus on getting there safely."), - ("organizing a closet", "My closet is a mess.", "Sort into keep, donate, and discard piles first."), - ("choosing a gift", "I need a gift for a friend.", "Think about what they enjoy doing in their free time."), - ("scheduling chores", "I keep forgetting chores.", "Write them on a weekly checklist you can see each morning."), - ("saving time", "Mornings feel rushed.", "Prepare clothes and bags the night before."), - ("handling a spill", "I spilled water on my desk.", "Unplug electronics first, then blot the liquid with a cloth."), - ("choosing a route", "Traffic is bad on my usual route.", "Check a map app for a lighter route before leaving."), - ("staying hydrated", "I forget to drink water.", "Keep a bottle within sight as a visual reminder."), - ], - "learning_tutor": [ - ("asking about fractions", "What is one half?", "One half means one part out of two equal parts."), - ("grammar help", "When do I use went?", "Use went for the past. I went to school yesterday."), - ("science question", "Why does ice melt?", "Ice melts when it gains heat. It becomes liquid water."), - ("history question", "Why do maps matter in history?", "Maps show where events happened and how borders changed."), - ("vocabulary help", "What does 'concise' mean?", "Concise means expressing something clearly in few words."), - ("multiplication help", "What is 7 times 8?", "7 times 8 is 56."), - ("reading comprehension", "How do I find the main idea?", "Look at the first and last sentences of the paragraph."), - ("study habits", "How should I study for a test?", "Break the material into small chunks and review daily."), - ("essay structure", "How do I start an essay?", "Open with a clear sentence stating your main point."), - ("geometry help", "What is a right angle?", "A right angle measures exactly ninety degrees."), - ("spelling help", "How do I remember 'necessary'?", "Remember: one collar, two sleeves -- one C, two S's."), - ("note taking", "How should I take notes?", "Write short phrases, not full sentences, and use bullet points."), - ], - "coding_mentor": [ - ("debugging Python", "My loop prints too many lines.", "Check the range. The stop value may be too large."), - ("understanding variables", "What is a variable?", "A variable is a name that stores a value."), - ("reading an error", "What does NameError mean?", "It usually means Python cannot find that variable name."), - ("understanding functions", "Why use functions?", "Functions let you reuse steps without repeating code."), - ("list vs tuple", "What's the difference between a list and a tuple?", "Lists can change after creation; tuples cannot."), - ("indexing errors", "I got an IndexError.", "You likely tried to access a position beyond the list's length."), - ("understanding loops", "When do I use a while loop?", "Use a while loop when you don't know how many times to repeat in advance."), - ("git basics", "What does git commit do?", "It saves a snapshot of your staged changes to the project history."), - ("testing code", "Why write tests?", "Tests catch mistakes early and confirm code still works after changes."), - ("naming variables", "How should I name variables?", "Use short, descriptive names that explain what the value represents."), - ("recursion confusion", "I don't understand recursion.", "A recursive function solves a small case, then calls itself on a smaller one."), - ("reading stack traces", "The stack trace is confusing.", "Start from the bottom-most line pointing to your own code."), - ], - "empathy_support": [ - ("nervous before exam", "I feel nervous about my exam.", "That is normal. Study small parts and take breaks."), - ("friend conflict", "My friend ignored me.", "Ask calmly what happened. Listen before deciding."), - ("mistake at work", "I made a mistake.", "Own it, fix what you can, and learn the cause."), - ("feeling overwhelmed", "I have too much to do.", "Pick one task, finish it, then move to the next."), - ("disappointment", "I didn't get the result I wanted.", "It's okay to feel disappointed. Consider what to try differently."), - ("homesickness", "I miss home.", "That feeling is common. Reach out to family when you can."), - ("public speaking fear", "I'm scared to speak in front of others.", "Practice out loud a few times; familiarity reduces nerves."), - ("comparison worry", "I feel behind compared to others.", "Everyone moves at a different pace. Focus on your own progress."), - ("difficult feedback", "I got harsh feedback.", "Take a breath, look for the useful part, and set the rest aside."), - ("change anxiety", "Things are changing and I feel unsettled.", "Focus on what stays the same and what you can control."), - ("apologizing", "I need to apologize but don't know how.", "Be specific about what happened and how you'll do better."), - ("burnout", "I feel exhausted from working nonstop.", "Rest is productive too. Consider a short, real break."), - ], - "professional_chat": [ - ("email rewrite", "Can you make this email polite?", "Yes. Keep it short, clear, and respectful."), - ("meeting plan", "How should I run a meeting?", "Set a goal, list topics, and end with action items."), - ("status update", "I need to report progress.", "Say what is done, what is blocked, and what comes next."), - ("giving feedback", "How do I give feedback kindly?", "Be specific, focus on the work, and suggest a next step."), - ("declining a request", "How do I say no politely?", "Thank them, explain briefly, and offer an alternative if possible."), - ("negotiating a deadline", "I need more time on a project.", "Explain the reason and propose a new, realistic date early."), - ("onboarding a teammate", "How do I help a new hire settle in?", "Share key contacts, documents, and a short first-week plan."), - ("prioritizing tasks", "I have too many tasks today.", "Rank by deadline and impact, then start with the most urgent."), - ("summarizing a call", "How do I summarize a meeting?", "List decisions made, owners, and deadlines in a few lines."), - ("cold outreach", "How do I write a cold email?", "Keep it short, state the purpose, and make the ask clear."), - ("handling conflict", "A coworker disagreed with my plan.", "Ask about their concern directly and look for common ground."), - ("requesting resources", "How do I ask for more budget?", "Explain the need, the expected benefit, and the cost clearly."), - ], - } - items = scenarios[topic] - endings = [ - "The best next step is to act carefully and review the result.", - "The best next step is to keep it simple and adjust later.", - "The best next step is to ask for help if anything is unclear.", - "The best next step is to write it down so it isn't forgotten.", - "The best next step is to check in again after trying it.", - ] - period = combinatorial_period(len(items), 9, len(endings)) - blocks = [] - for index in range(min(count, period)): - item_i, turn, ending_i = mixed_radix_pick(index, len(items), 9, len(endings)) - title, user_text, assistant_text = items[item_i] - blocks.append( - f"Conversation: {title}.\n" - f"User: {user_text}\n" - f"Assistant: {assistant_text}\n" - f"User: Can you explain simply?\n" - f"Assistant: Yes. I will use short steps.\n" - f"Assistant: First, understand the problem.\n" - f"Assistant: Second, choose a small action.\n" - f"Assistant: Third, check the result.\n" - f"User: What should I avoid?\n" - f"Assistant: Avoid guessing when facts are missing.\n" - f"Assistant: Ask a clear question if needed.\n" - f"User: Give me a final answer.\n" - f"Assistant: {endings[ending_i]}\n" - f"This dialogue teaches helpful conversation turn {turn}." - ) - return blocks - - -def instruction_fine_tune_blocks(count: int, topic: str) -> list[str]: - """Create instruction fine-tuning corpus blocks. - - Args: - count: Number of blocks to generate. - topic: Instruction task group. - - Returns: - Instruction training blocks. - """ - - tasks = { - "writing_tasks": [ - ("Rewrite this sentence in simpler English.", "The child rapidly moved across the room.", "The child ran across the room."), - ("Summarize this passage.", "Mina planted seeds. She watered them. After many days, leaves grew.", "Mina planted and cared for seeds until they grew leaves."), - ("Make this polite.", "Send the report now.", "Please send the report when you have a moment."), - ("Shorten this sentence.", "Due to the fact that it was raining, we decided to stay inside.", "Because it was raining, we stayed inside."), - ("Fix the grammar.", "She don't like the plan.", "She doesn't like the plan."), - ("Make this more formal.", "Hey, can you send that file?", "Could you please send the file at your convenience?"), - ("Combine these sentences.", "The dog barked. The dog ran to the door.", "The dog barked and ran to the door."), - ("Add a stronger verb.", "The team did a good job on the project.", "The team excelled on the project."), - ("Remove redundancy.", "In my opinion, I think the plan is good.", "I think the plan is good."), - ("Write a topic sentence.", "Details about rainforests having high rainfall and diverse species.", "Rainforests are defined by heavy rainfall and remarkable species diversity."), - ], - "reasoning_tasks": [ - ("Solve the word problem.", "A box has 6 pens. Ravi adds 4 pens. How many pens are there?", "There are 10 pens."), - ("Choose the safer action.", "A wire is broken. Should Tara touch it or call an adult?", "Tara should call an adult."), - ("Find the cause.", "The lamp does not turn on. The bulb is loose.", "The loose bulb may be the cause."), - ("Solve the word problem.", "Lena has 15 stickers and gives 6 away. How many are left?", "9 stickers are left."), - ("Order the steps.", "Steps: pour water, boil water, add tea leaves, given out of order.", "Pour water, boil water, add tea leaves."), - ("Spot the contradiction.", "The store is open every day. The store is closed on Sundays.", "These two statements contradict each other."), - ("Draw a conclusion.", "All birds in the flock flew south. It is now winter here.", "The birds likely migrated for winter."), - ("Find the missing step.", "Recipe skips from 'mix batter' to 'serve cake' with nothing baked.", "The recipe is missing a baking step."), - ("Compare two options.", "Option A costs less but takes longer. Option B costs more but is faster.", "Choose based on whether time or cost matters more."), - ("Explain the pattern.", "2, 4, 6, 8, ...", "The pattern adds 2 to get each next number."), - ], - "coding_tasks": [ - ("Write a Python function that adds two numbers.", "Use parameters a and b.", "def add(a, b):\n return a + b"), - ("Explain this code.", "print(len([1, 2, 3]))", "It creates a list with three items and prints its length, which is 3."), - ("Fix the bug.", "for i in range(3):\nprint(i)", "Indent the print line inside the loop."), - ("Write a function that returns the max of two numbers.", "Use parameters a and b.", "def maximum(a, b):\n return a if a > b else b"), - ("Explain this code.", "x = [n * n for n in range(5)]", "It builds a list of squares for numbers 0 through 4 using a list comprehension."), - ("Fix the bug.", "def greet(name)\n print('Hello ' + name)", "Add a colon after the function signature: def greet(name):"), - ("Write a function that checks if a number is even.", "Use one parameter n.", "def is_even(n):\n return n % 2 == 0"), - ("Explain this code.", "total = sum([1, 2, 3])", "It adds up the numbers in the list, giving a total of 6."), - ("Fix the bug.", "if x = 5:\n print('five')", "Use == for comparison instead of =: if x == 5:"), - ("Write a function that reverses a string.", "Use one parameter text.", "def reverse(text):\n return text[::-1]"), - ], - "classification_tasks": [ - ("Classify the sentence.", "The sky is cloudy today.", "Category: weather observation."), - ("Classify the request.", "Can you help me debug this error?", "Category: coding help."), - ("Classify the emotion.", "I am proud because I finished the project.", "Emotion: proud."), - ("Classify the sentence.", "Water boils at 100 degrees Celsius.", "Category: science fact."), - ("Classify the request.", "Please summarize this article for me.", "Category: writing help."), - ("Classify the emotion.", "I felt nervous before the interview.", "Emotion: nervous."), - ("Classify the sentence.", "Paris is the capital of France.", "Category: geography fact."), - ("Classify the request.", "Can you check my math homework?", "Category: math help."), - ("Classify the emotion.", "I was relieved when the test was over.", "Emotion: relieved."), - ("Classify the sentence.", "The stock market fell sharply today.", "Category: financial news."), - ], - "format_following": [ - ("Answer with two bullet points.", "Give two safe cooking tips.", "- Wash your hands.\n- Turn off the stove after cooking."), - ("Return only the number.", "What is 8 plus 5?", "13"), - ("Use a short answer.", "Why do plants need light?", "Plants use light to make food."), - ("Answer with two bullet points.", "Give two tips for studying.", "- Take short breaks.\n- Review notes daily."), - ("Return only the number.", "What is 12 minus 7?", "5"), - ("Use a short answer.", "Why do we wear seatbelts?", "Seatbelts help prevent injury in a crash."), - ("Answer with three bullet points.", "List three parts of a plant.", "- Roots\n- Stem\n- Leaves"), - ("Return only the word.", "What do bees produce?", "Honey"), - ("Use one sentence.", "What is gravity?", "Gravity is the force that pulls objects toward each other."), - ("Answer in a single word.", "What gas do plants release during photosynthesis?", "Oxygen"), - ], - } - items = tasks[topic] - closers = [ - "This instruction sample teaches format control.", - "This instruction sample teaches staying on topic.", - "This instruction sample teaches concise responses.", - "This instruction sample teaches following the exact request.", - ] - period = combinatorial_period(len(items), 13, len(closers)) - blocks = [] - for index in range(min(count, period)): - item_i, turn, closer_i = mixed_radix_pick(index, len(items), 13, len(closers)) - instruction, input_text, output_text = items[item_i] - blocks.append( - f"Instruction: {instruction}\n" - f"Input: {input_text}\n" - f"Response: {output_text}\n" - f"The response follows the instruction.\n" - f"The response stays focused on the user request.\n" - f"The response avoids extra unrelated text.\n" - f"If information is missing, ask one clear question.\n" - f"If the task is simple, answer directly.\n" - f"If the task needs steps, use short ordered steps.\n" - f"{closers[closer_i]} (variant {turn})" - ) - return blocks - - -CODE_LANGUAGE_SPECS = { - "python": { - "label": "Python", - "ext": "py", - "comment": "#", - "types": ["list[int]", "dict[str, int]", "str", "tuple[int, int]", "set[str]"], - "containers": ["list", "dictionary", "set", "tuple", "file"], - "errors": ["IndexError", "KeyError", "TypeError", "ValueError", "NameError"], - }, - "javascript": { - "label": "JavaScript", - "ext": "js", - "comment": "//", - "types": ["Array", "Object", "string", "number", "Promise"], - "containers": ["array", "object", "map", "set", "DOM node"], - "errors": ["TypeError", "ReferenceError", "RangeError", "SyntaxError", "Promise rejection"], - }, - "typescript": { - "label": "TypeScript", - "ext": "ts", - "comment": "//", - "types": ["number[]", "Record", "string", "Promise", "ReadonlyArray"], - "containers": ["typed array", "record", "interface", "union", "generic"], - "errors": ["type mismatch", "undefined value", "narrowing error", "implicit any", "async error"], - }, - "java": { - "label": "Java", - "ext": "java", - "comment": "//", - "types": ["List", "Map", "String", "Optional", "Set"], - "containers": ["ArrayList", "HashMap", "HashSet", "class", "stream"], - "errors": ["NullPointerException", "IndexOutOfBoundsException", "IllegalArgumentException", "ClassCastException", "IOException"], - }, - "csharp": { - "label": "C#", - "ext": "cs", - "comment": "//", - "types": ["List", "Dictionary", "string", "Task", "IEnumerable"], - "containers": ["List", "Dictionary", "HashSet", "class", "LINQ query"], - "errors": ["NullReferenceException", "IndexOutOfRangeException", "InvalidOperationException", "ArgumentException", "async deadlock"], - }, - "cpp": { - "label": "C++", - "ext": "cpp", - "comment": "//", - "types": ["vector", "unordered_map", "string", "unique_ptr", "optional"], - "containers": ["vector", "unordered_map", "set", "struct", "iterator"], - "errors": ["segmentation fault", "dangling pointer", "out_of_range", "memory leak", "undefined behavior"], - }, - "rust": { - "label": "Rust", - "ext": "rs", - "comment": "//", - "types": ["Vec", "HashMap", "String", "Option", "Result"], - "containers": ["Vec", "HashMap", "slice", "struct", "iterator"], - "errors": ["borrow checker error", "panic", "lifetime error", "unwrap failure", "type mismatch"], - }, - "go": { - "label": "Go", - "ext": "go", - "comment": "//", - "types": ["[]int", "map[string]int", "string", "error", "chan int"], - "containers": ["slice", "map", "struct", "goroutine", "channel"], - "errors": ["nil pointer", "index out of range", "data race", "ignored error", "deadlock"], - }, - "sql": { - "label": "SQL", - "ext": "sql", - "comment": "--", - "types": ["INTEGER", "TEXT", "TIMESTAMP", "BOOLEAN", "DECIMAL"], - "containers": ["table", "index", "view", "join", "transaction"], - "errors": ["missing index", "duplicate key", "bad join", "null value", "slow query"], - }, - "bash": { - "label": "Bash", - "ext": "sh", - "comment": "#", - "types": ["string", "array", "exit code", "path", "environment variable"], - "containers": ["loop", "function", "pipe", "process", "file"], - "errors": ["missing quote", "bad path", "nonzero exit", "unset variable", "permission denied"], - }, -} - - -CODE_TASKS = [ - "parse input", - "validate data", - "filter a collection", - "count repeated values", - "read a file safely", - "write a small helper", - "handle an error", - "sort records", - "cache a result", - "format output", - "test an edge case", - "split work into functions", -] - -CODE_PATTERNS = [ - "loop", - "function", - "guard clause", - "map lookup", - "unit test", - "small class", - "command handler", - "parser", - "retry step", - "cleanup step", -] - - -def code_training_block(language: str, index: int) -> str: - """Create one base-training code explanation block. - - Args: - language: Key from CODE_LANGUAGE_SPECS. - index: Unique deterministic block number. - - Returns: - Plain text code teaching block. - """ - - spec = CODE_LANGUAGE_SPECS[language] - task_i, pattern_i, type_i, container_i, error_i, name_i = mixed_radix_pick( - index, - len(CODE_TASKS), - len(CODE_PATTERNS), - len(spec["types"]), - len(spec["containers"]), - len(spec["errors"]), - 997, +try: + from .curriculum_subjects import ( + code_blocks, computer_blocks, everyday_blocks, geography_history_blocks, + language_blocks, math_blocks, reasoning_blocks, science_blocks, + social_blocks, ) - label = spec["label"] - comment = spec["comment"] - task = CODE_TASKS[task_i] - pattern = CODE_PATTERNS[pattern_i] - type_name = spec["types"][type_i] - container = spec["containers"][container_i] - error = spec["errors"][error_i] - unique = f"{language}_{name_i}_{index}" - if language == "python": - snippet = ( - f"def process_{unique}(items: list[int]) -> int:\n" - f" total = 0\n" - f" for value in items:\n" - f" if value >= 0:\n" - f" total += value\n" - f" return total\n\n" - f"assert process_{unique}([1, -2, 3]) == 4" - ) - elif language in {"javascript", "typescript"}: - annotation = ": number[]" if language == "typescript" else "" - return_type = ": number" if language == "typescript" else "" - snippet = ( - f"function process_{unique}(items{annotation}){return_type} {{\n" - f" let total = 0;\n" - f" for (const value of items) {{\n" - f" if (value >= 0) total += value;\n" - f" }}\n" - f" return total;\n" - f"}}\n\n" - f"console.assert(process_{unique}([1, -2, 3]) === 4);" - ) - elif language == "java": - snippet = ( - f"static int process{unique.title().replace('_', '')}(java.util.List items) {{\n" - f" int total = 0;\n" - f" for (int value : items) {{\n" - f" if (value >= 0) total += value;\n" - f" }}\n" - f" return total;\n" - f"}}" - ) - elif language == "csharp": - snippet = ( - f"static int Process{unique.title().replace('_', '')}(IEnumerable items) {{\n" - f" var total = 0;\n" - f" foreach (var value in items) {{\n" - f" if (value >= 0) total += value;\n" - f" }}\n" - f" return total;\n" - f"}}" - ) - elif language == "cpp": - snippet = ( - f"int process_{unique}(const std::vector& items) {{\n" - f" int total = 0;\n" - f" for (int value : items) {{\n" - f" if (value >= 0) total += value;\n" - f" }}\n" - f" return total;\n" - f"}}" - ) - elif language == "rust": - snippet = ( - f"fn process_{unique}(items: &[i32]) -> i32 {{\n" - f" let mut total = 0;\n" - f" for value in items {{\n" - f" if *value >= 0 {{ total += *value; }}\n" - f" }}\n" - f" total\n" - f"}}" - ) - elif language == "go": - snippet = ( - f"func process{unique.title().replace('_', '')}(items []int) int {{\n" - f" total := 0\n" - f" for _, value := range items {{\n" - f" if value >= 0 {{ total += value }}\n" - f" }}\n" - f" return total\n" - f"}}" - ) - elif language == "sql": - snippet = ( - f"SELECT user_id, SUM(amount) AS total_{name_i}\n" - f"FROM payments\n" - f"WHERE amount >= 0\n" - f"GROUP BY user_id\n" - f"ORDER BY total_{name_i} DESC;" - ) - else: - snippet = ( - f"process_{unique}() {{\n" - f" local total=0\n" - f" for value in \"$@\"; do\n" - f" if [ \"$value\" -ge 0 ]; then total=$((total + value)); fi\n" - f" done\n" - f" printf '%s\\n' \"$total\"\n" - f"}}" - ) - return ( - f"{label} example {index}.\n" - f"Goal: teach how to {task} with a {pattern}.\n" - f"The example uses a {container} and a {type_name} value.\n" - f"```{spec['ext']}\n{snippet}\n```\n" - f"{comment} Read the code from top to bottom.\n" - f"The function receives data, skips invalid values, and returns one clear result.\n" - f"A common mistake in this topic is {error}.\n" - f"Check the empty input case before trusting the code.\n" - f"Keep names descriptive, keep steps small, and test one behavior at a time.\n" + from .curriculum_finetune import ( + conversation_fine_tune_blocks, instruction_fine_tune_blocks, + programming_deep_blocks, ) - - -def code_fine_tune_block(language: str, index: int) -> str: - """Create one code fine-tuning instruction block. - - Args: - language: Key from CODE_LANGUAGE_SPECS. - index: Unique deterministic block number. - - Returns: - Instruction-style code fine-tuning block. - """ - - spec = CODE_LANGUAGE_SPECS[language] - task_i, pattern_i, type_i, container_i, error_i, variant = mixed_radix_pick( - index, - len(CODE_TASKS), - len(CODE_PATTERNS), - len(spec["types"]), - len(spec["containers"]), - len(spec["errors"]), - 2003, + from .curriculum_code import ( + CODE_LANGUAGE_SPECS, code_fine_tune_block, code_training_block, + write_target_bytes, ) - label = spec["label"] - task = CODE_TASKS[task_i] - pattern = CODE_PATTERNS[pattern_i] - type_name = spec["types"][type_i] - container = spec["containers"][container_i] - error = spec["errors"][error_i] - unique = f"{language}_{variant}_{index}" - if language == "sql": - response_code = ( - f"SELECT category, COUNT(*) AS count_{variant}\n" - f"FROM events\n" - f"WHERE created_at >= CURRENT_DATE - INTERVAL '7 days'\n" - f"GROUP BY category\n" - f"ORDER BY count_{variant} DESC;" - ) - bug_code = "SELECT category, COUNT(*) FROM events WHERE created_at >= CURRENT_DATE - INTERVAL '7 days';" - elif language == "bash": - response_code = ( - f"count_{unique}() {{\n" - f" local path=\"$1\"\n" - f" if [ ! -f \"$path\" ]; then return 1; fi\n" - f" grep -c \"ERROR\" \"$path\"\n" - f"}}" - ) - bug_code = "grep -c ERROR $path" - elif language == "python": - response_code = ( - f"def solve_{unique}(items: list[int], limit: int) -> list[int]:\n" - f" result: list[int] = []\n" - f" for value in items:\n" - f" if 0 <= value <= limit:\n" - f" result.append(value)\n" - f" return result\n\n" - f"assert solve_{unique}([1, -1, 5], 3) == [1]" - ) - bug_code = f"def solve_{unique}(items, limit):\n return [x for x in items if x <= limit]" - elif language in {"javascript", "typescript"}: - annotation = ": number[]" if language == "typescript" else "" - limit_annotation = ": number" if language == "typescript" else "" - return_annotation = ": number[]" if language == "typescript" else "" - response_code = ( - f"function solve_{unique}(items{annotation}, limit{limit_annotation}){return_annotation} {{\n" - f" const result = [];\n" - f" for (const value of items) {{\n" - f" if (value >= 0 && value <= limit) result.push(value);\n" - f" }}\n" - f" return result;\n" - f"}}\n\n" - f"console.assert(JSON.stringify(solve_{unique}([1, -1, 5], 3)) === JSON.stringify([1]));" - ) - bug_code = f"function solve_{unique}(items, limit) {{ return items.filter(x => x <= limit); }}" - elif language == "java": - method = f"solve{unique.title().replace('_', '')}" - response_code = ( - f"static java.util.List {method}(java.util.List items, int limit) {{\n" - f" java.util.List result = new java.util.ArrayList<>();\n" - f" for (int value : items) {{\n" - f" if (value >= 0 && value <= limit) result.add(value);\n" - f" }}\n" - f" return result;\n" - f"}}" - ) - bug_code = f"static java.util.List {method}(java.util.List items, int limit) {{ return null; }}" - elif language == "csharp": - method = f"Solve{unique.title().replace('_', '')}" - response_code = ( - f"static List {method}(IEnumerable items, int limit) {{\n" - f" var result = new List();\n" - f" foreach (var value in items) {{\n" - f" if (value >= 0 && value <= limit) result.Add(value);\n" - f" }}\n" - f" return result;\n" - f"}}" - ) - bug_code = f"static List {method}(IEnumerable items, int limit) => items.Where(x => x <= limit).ToList();" - elif language == "cpp": - response_code = ( - f"std::vector solve_{unique}(const std::vector& items, int limit) {{\n" - f" std::vector result;\n" - f" for (int value : items) {{\n" - f" if (value >= 0 && value <= limit) result.push_back(value);\n" - f" }}\n" - f" return result;\n" - f"}}" - ) - bug_code = f"std::vector solve_{unique}(std::vector& items, int limit) {{ return items; }}" - elif language == "rust": - response_code = ( - f"fn solve_{unique}(items: &[i32], limit: i32) -> Vec {{\n" - f" items.iter()\n" - f" .copied()\n" - f" .filter(|value| *value >= 0 && *value <= limit)\n" - f" .collect()\n" - f"}}" - ) - bug_code = f"fn solve_{unique}(items: Vec, limit: i32) -> Vec {{ items }}" - elif language == "go": - method = f"solve{unique.title().replace('_', '')}" - response_code = ( - f"func {method}(items []int, limit int) []int {{\n" - f" result := make([]int, 0, len(items))\n" - f" for _, value := range items {{\n" - f" if value >= 0 && value <= limit {{ result = append(result, value) }}\n" - f" }}\n" - f" return result\n" - f"}}" - ) - bug_code = f"func {method}(items []int, limit int) []int {{ return items }}" - return ( - f"Instruction: Write {label} code to {task}.\n" - f"User context: Use a {pattern}. The input involves a {container}. The important type is {type_name}.\n" - f"Response:\n" - f"```{spec['ext']}\n{response_code}\n```\n" - f"Explanation: The solution separates input handling from the core operation.\n" - f"It names the result clearly and keeps each step small.\n" - f"Edge case: empty input should return a safe default or a clear error.\n" - f"Debugging example:\n" - f"```{spec['ext']}\n{bug_code}\n```\n" - f"The likely issue is {error}.\n" - f"Fix: validate inputs, check boundaries, and test the smallest failing case first.\n" - f"Final answer: use the shown pattern, then add tests for normal, empty, and invalid inputs.\n" + from .curriculum_shared import write_blocks +except ImportError: + from curriculum_subjects import ( + code_blocks, computer_blocks, everyday_blocks, geography_history_blocks, + language_blocks, math_blocks, reasoning_blocks, science_blocks, + social_blocks, ) + from curriculum_finetune import ( + conversation_fine_tune_blocks, instruction_fine_tune_blocks, + programming_deep_blocks, + ) + from curriculum_code import ( + CODE_LANGUAGE_SPECS, code_fine_tune_block, code_training_block, + write_target_bytes, + ) + from curriculum_shared import write_blocks - -def write_target_bytes(path: Path, block_factory, target_bytes: int) -> None: - """Write generated corpus blocks until a file reaches the target size. - - Args: - path: Output file path. - block_factory: Callable accepting a block index and returning text. - target_bytes: Minimum UTF-8 byte size to write. - """ - - path.parent.mkdir(parents=True, exist_ok=True) - with path.open("w", encoding="utf-8", newline="\n") as handle: - index = 0 - while handle.tell() < target_bytes: - handle.write(block_factory(index)) - handle.write("\n\n") - index += 1 - +ROOT = Path(__file__).resolve().parents[1] / "engine" / "default_data" def main() -> None: """Generate the expanded default curriculum."""