Spaces:
Sleeping
Sleeping
| """Background training worker. | |
| Runs *real* gradient-descent training of a small 3-layer MLP with NumPy | |
| (no GPU required). Each epoch persists loss to the DB so clients can stream | |
| metrics. To fine-tune real models, replace `_run` with a TRL/transformers | |
| training call (see README "Fine-tuning real models"). | |
| """ | |
| import datetime | |
| import threading | |
| import numpy as np | |
| from .. import models as orm | |
| from ..database import SessionLocal | |
| _workers = {} | |
| _cancel = set() | |
| def _forward_backward(X, y, W1, b1, W2, b2): | |
| z = X @ W1 + b1 | |
| h = np.tanh(z) | |
| logits = h @ W2 + b2 | |
| ex = np.exp(logits - logits.max(axis=1, keepdims=True)) | |
| p = ex / ex.sum(axis=1, keepdims=True) | |
| loss = -np.log(p[np.arange(len(y)), y]).mean() | |
| dp = p.copy() | |
| dp[np.arange(len(y)), y] -= 1.0 | |
| dp /= len(y) | |
| dW2 = h.T @ dp | |
| db2 = dp.sum(0) | |
| dh = dp @ W2.T | |
| dz = dh * (1.0 - np.tanh(z) ** 2) | |
| dW1 = X.T @ dz | |
| db1 = dz.sum(0) | |
| return loss, dW1, db1, dW2, db2 | |
| def _run(job_id: str): | |
| db = SessionLocal() | |
| try: | |
| job = db.get(orm.TrainingJob, job_id) | |
| if job is None: | |
| return | |
| cfg = job.config or {} | |
| rng = np.random.default_rng(int(cfg.get("seed", 42))) | |
| n, d, k = 600, 8, 4 | |
| X = rng.normal(size=(n, d)) | |
| true_W = rng.normal(size=(d, d)) | |
| y = (np.argmax(X @ true_W + 0.4 * rng.normal(size=(n, d)), axis=1) % k).astype(int) | |
| W1 = rng.normal(size=(d, 32)) * 0.1 | |
| b1 = np.zeros(32) | |
| W2 = rng.normal(size=(32, k)) * 0.1 | |
| b2 = np.zeros(k) | |
| lr = float(cfg.get("lr", 0.05)) | |
| epochs = int(cfg.get("epochs", 20)) | |
| bs = int(cfg.get("batch_size", 32)) | |
| job.status = "running" | |
| job.started_at = datetime.datetime.utcnow() | |
| loss_history = (job.metrics or {}).get("loss") if isinstance(job.metrics, dict) else None | |
| loss_history = list(loss_history) if loss_history else [] | |
| db.commit() | |
| for ep in range(epochs): | |
| if job_id in _cancel: | |
| job.status = "cancelled" | |
| job.finished_at = datetime.datetime.utcnow() | |
| db.commit() | |
| return | |
| perm = rng.permutation(n) | |
| tot, nb = 0.0, 0 | |
| for i in range(0, n, bs): | |
| idx = perm[i:i + bs] | |
| loss, dW1, db1, dW2, db2 = _forward_backward(X[idx], y[idx], W1, b1, W2, b2) | |
| W2 -= lr * dW2 | |
| b2 -= lr * db2 | |
| W1 -= lr * dW1 | |
| b1 -= lr * db1 | |
| tot += loss | |
| nb += 1 | |
| avg = float(tot / nb) | |
| loss_history.append(avg) | |
| job.current_epoch = ep + 1 | |
| job.loss = avg | |
| job.metrics = {"loss": loss_history[-200:], "lr": lr, "batch_size": bs} | |
| db.commit() | |
| h = np.tanh(X @ W1 + b1) | |
| preds = np.argmax(h @ W2 + b2, axis=1) | |
| acc = float((preds == y).mean()) | |
| job.status = "completed" | |
| job.finished_at = datetime.datetime.utcnow() | |
| job.metrics = {"loss": loss_history[-200:], "accuracy": acc} | |
| db.commit() | |
| except Exception as e: # noqa: BLE001 | |
| try: | |
| job.status = "failed" | |
| job.error = str(e) | |
| job.finished_at = datetime.datetime.utcnow() | |
| db.commit() | |
| except Exception: | |
| pass | |
| finally: | |
| db.close() | |
| def launch(job_id: str) -> bool: | |
| """Start a training thread for the job; returns False if already running.""" | |
| if job_id in _workers and _workers[job_id].is_alive(): | |
| return False | |
| t = threading.Thread(target=_run, args=(job_id,), daemon=True) | |
| _workers[job_id] = t | |
| t.start() | |
| return True | |
| def cancel(job_id: str) -> None: | |
| _cancel.add(job_id) | |
| def is_running(job_id: str) -> bool: | |
| return job_id in _workers and _workers[job_id].is_alive() | |