wehub-resource-sync
1449 lines
61 KiB
TypeScript
1449 lines
61 KiB
TypeScript
console.log("Starting FineTune Studio MCP Server...");
|
||
|
||
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
|
||
import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
|
||
import {
|
||
registerAppTool,
|
||
registerAppResource,
|
||
RESOURCE_MIME_TYPE,
|
||
} from "@modelcontextprotocol/ext-apps/server";
|
||
import cors from "cors";
|
||
import express from "express";
|
||
import fs from "node:fs/promises";
|
||
import path from "node:path";
|
||
import { fileURLToPath } from "node:url";
|
||
import { z } from "zod";
|
||
import { Client as GradioClient } from "@gradio/client";
|
||
|
||
// ── __dirname compat (works on all Node ESM versions) ─────────────────
|
||
const __filename = fileURLToPath(import.meta.url);
|
||
const __dirname = path.dirname(__filename);
|
||
|
||
const server = new McpServer({ name: "FineTuneStudio", version: "1.0.0" });
|
||
const resourceUri = "ui://finetune-studio/widget.html";
|
||
|
||
const HF_API = "https://huggingface.co/api";
|
||
const HF_ROUTER_URL = "https://router.huggingface.co/v1/chat/completions";
|
||
|
||
// ── Inference Space templates ──────────────────────────────────────────
|
||
// A Gradio Space that loads the fine-tuned model and serves it via
|
||
// a simple /api/predict endpoint. Deployed automatically after training.
|
||
const INFERENCE_APP_PY = `
|
||
import gradio as gr
|
||
import json
|
||
import os
|
||
import torch
|
||
from transformers import pipeline
|
||
|
||
MODEL_ID = os.environ.get("MODEL_ID", "")
|
||
HF_TOKEN = os.environ.get("HF_TOKEN", None)
|
||
|
||
print(f"[inference] Loading {MODEL_ID} ...", flush=True)
|
||
try:
|
||
_pipe = pipeline(
|
||
"text-generation",
|
||
model=MODEL_ID,
|
||
torch_dtype=torch.float32,
|
||
token=HF_TOKEN or None,
|
||
)
|
||
print("[inference] Ready!", flush=True)
|
||
except Exception as _e:
|
||
print(f"[inference] Load error: {_e}", flush=True)
|
||
_pipe = None
|
||
|
||
def predict(messages_json: str, max_tokens: float = 512, temperature: float = 0.7) -> str:
|
||
"""Generate a response given a JSON-encoded messages array."""
|
||
if _pipe is None:
|
||
return "Error: model failed to load. Check Space logs."
|
||
try:
|
||
messages = json.loads(messages_json)
|
||
temp = float(temperature)
|
||
result = _pipe(
|
||
messages,
|
||
max_new_tokens=int(max_tokens),
|
||
temperature=temp if temp > 0.01 else None,
|
||
do_sample=temp > 0.01,
|
||
return_full_text=False,
|
||
)
|
||
content = result[0]["generated_text"]
|
||
if isinstance(content, list):
|
||
return content[-1].get("content", str(content[-1]))
|
||
return str(content)
|
||
except Exception as e:
|
||
return f"Error: {e}"
|
||
|
||
demo = gr.Interface(
|
||
fn=predict,
|
||
inputs=[
|
||
gr.Textbox(label="messages_json", lines=3),
|
||
gr.Number(value=512, label="max_tokens"),
|
||
gr.Number(value=0.7, label="temperature"),
|
||
],
|
||
outputs=gr.Textbox(label="response"),
|
||
title=f"Inference: {MODEL_ID}",
|
||
flagging_mode="never",
|
||
api_name="predict",
|
||
)
|
||
demo.launch()
|
||
`.trim();
|
||
|
||
const INFERENCE_REQUIREMENTS_TXT = `gradio>=4.0.0
|
||
transformers>=4.40.0
|
||
torch
|
||
accelerate>=0.26.0`.trim();
|
||
|
||
// Helper: fetch the current HEAD commit SHA of a repo's branch.
|
||
// The HF commit API requires parentCommit when the repo already has commits;
|
||
// omitting it on a non-empty repo causes a 412 Precondition Failed error.
|
||
async function getHeadCommit(
|
||
repoType: "spaces" | "models",
|
||
owner: string,
|
||
name: string,
|
||
branch: string,
|
||
token: string,
|
||
): Promise<string | null> {
|
||
try {
|
||
const res = await fetch(
|
||
`https://huggingface.co/api/${repoType}/${owner}/${name}/commits/${branch}`,
|
||
{ headers: { Authorization: `Bearer ${token}` }, signal: AbortSignal.timeout(8000) },
|
||
);
|
||
if (!res.ok) return null;
|
||
const data = await res.json() as any[];
|
||
return Array.isArray(data) && data[0]?.id ? String(data[0].id) : null;
|
||
} catch {
|
||
return null;
|
||
}
|
||
}
|
||
|
||
// Deploys a Gradio inference Space for a fine-tuned model.
|
||
// Returns the Space URL on success, throws on failure (so callers know it failed).
|
||
async function deployInferenceSpace(
|
||
username: string,
|
||
projectName: string,
|
||
token: string,
|
||
): Promise<{ space_url: string; space_name: string }> {
|
||
const spaceName = `inference-${projectName}`;
|
||
const modelId = `${username}/${projectName}`;
|
||
const spaceUrl = `https://huggingface.co/spaces/${username}/${spaceName}`;
|
||
|
||
// ── 1. Create Space (idempotent — 409 means already exists) ─────────
|
||
const createRes = await fetch("https://huggingface.co/api/repos/create", {
|
||
method: "POST",
|
||
headers: { Authorization: `Bearer ${token}`, "Content-Type": "application/json" },
|
||
body: JSON.stringify({ type: "space", name: spaceName, private: false, sdk: "gradio" }),
|
||
signal: AbortSignal.timeout(12000),
|
||
});
|
||
const isNew = createRes.status === 200 || createRes.status === 201;
|
||
if (!createRes.ok && createRes.status !== 409) {
|
||
const e = await createRes.text().catch(() => "");
|
||
throw new Error(`Space create failed (HTTP ${createRes.status}): ${e.slice(0, 200)}`);
|
||
}
|
||
console.log(`[inference-space] Space ${isNew ? "created" : "already exists"}: ${username}/${spaceName}`);
|
||
|
||
// ── 2. Wait for git repo to initialise (only needed on fresh create) ─
|
||
// HF backend takes 1-3 s to set up the git repo after creation.
|
||
// Committing too early returns 404 on the commit endpoint.
|
||
if (isNew) {
|
||
console.log("[inference-space] Waiting 4s for repo to initialise…");
|
||
await new Promise(r => setTimeout(r, 4000));
|
||
}
|
||
|
||
// ── 3. Get HEAD commit SHA (required as parentCommit) ────────────────
|
||
// If we omit parentCommit on a non-empty repo the API returns 412.
|
||
// On a brand-new repo it may still be empty (null is fine — omit it).
|
||
const parentCommit = await getHeadCommit("spaces", username, spaceName, "main", token);
|
||
console.log(`[inference-space] parentCommit=${parentCommit ?? "(none — empty repo)"}`);
|
||
|
||
// ── 4. Commit app.py + requirements.txt ──────────────────────────────
|
||
const toB64 = (s: string) => Buffer.from(s, "utf-8").toString("base64");
|
||
const headerValue: Record<string, unknown> = {
|
||
summary: `Deploy inference Space for ${modelId}`,
|
||
};
|
||
if (parentCommit) headerValue.parentCommit = parentCommit;
|
||
|
||
const ndjson = [
|
||
JSON.stringify({ key: "header", value: headerValue }),
|
||
JSON.stringify({ key: "file", value: { path: "app.py", encoding: "base64", content: toB64(INFERENCE_APP_PY) } }),
|
||
JSON.stringify({ key: "file", value: { path: "requirements.txt", encoding: "base64", content: toB64(INFERENCE_REQUIREMENTS_TXT) } }),
|
||
].join("\n");
|
||
|
||
const commitRes = await fetch(
|
||
`https://huggingface.co/api/spaces/${username}/${spaceName}/commit/main`,
|
||
{
|
||
method: "POST",
|
||
headers: { Authorization: `Bearer ${token}`, "Content-Type": "application/x-ndjson" },
|
||
body: ndjson,
|
||
signal: AbortSignal.timeout(30000),
|
||
},
|
||
);
|
||
if (!commitRes.ok) {
|
||
const e = await commitRes.text().catch(() => "");
|
||
throw new Error(`Space commit failed (HTTP ${commitRes.status}): ${e.slice(0, 300)}`);
|
||
}
|
||
console.log(`[inference-space] ✅ app.py committed to ${username}/${spaceName}`);
|
||
|
||
// ── 5. Set MODEL_ID and HF_TOKEN secrets ─────────────────────────────
|
||
for (const [key, value] of [["MODEL_ID", modelId], ["HF_TOKEN", token]] as const) {
|
||
const secretRes = await fetch(
|
||
`https://huggingface.co/api/spaces/${username}/${spaceName}/secrets`,
|
||
{
|
||
method: "POST",
|
||
headers: { Authorization: `Bearer ${token}`, "Content-Type": "application/json" },
|
||
body: JSON.stringify({ key, value }),
|
||
signal: AbortSignal.timeout(8000),
|
||
},
|
||
);
|
||
if (!secretRes.ok) {
|
||
console.warn(`[inference-space] secret ${key} failed: ${secretRes.status}`);
|
||
} else {
|
||
console.log(`[inference-space] ✅ secret ${key} set`);
|
||
}
|
||
}
|
||
|
||
console.log(`[inference-space] ✅ Fully deployed: ${spaceUrl}`);
|
||
return { space_url: spaceUrl, space_name: spaceName };
|
||
}
|
||
|
||
// Read token dynamically so it is always current regardless of when the
|
||
// env var was injected (e.g. mcp-use dashboard sets it after process start).
|
||
function getHFToken(): string {
|
||
return process.env.HF_TOKEN || "";
|
||
}
|
||
|
||
// ── Helper: HF API fetch ──────────────────────────────────────────────
|
||
async function hfGet(url: string) {
|
||
const token = getHFToken();
|
||
const headers: Record<string, string> = { "Content-Type": "application/json" };
|
||
if (token) headers["Authorization"] = `Bearer ${token}`;
|
||
const res = await fetch(url, { headers });
|
||
if (!res.ok) {
|
||
const text = await res.text().catch(() => "");
|
||
throw new Error(`HF API error ${res.status}: ${text.slice(0, 300)}`);
|
||
}
|
||
return res.json();
|
||
}
|
||
|
||
async function hfPost(url: string, body: unknown) {
|
||
const token = getHFToken();
|
||
if (!token) throw new Error("HF_TOKEN environment variable is not set. Add it to your environment and restart.");
|
||
const res = await fetch(url, {
|
||
method: "POST",
|
||
headers: {
|
||
"Authorization": `Bearer ${token}`,
|
||
"Content-Type": "application/json",
|
||
},
|
||
body: JSON.stringify(body),
|
||
});
|
||
if (!res.ok) {
|
||
const text = await res.text().catch(() => "");
|
||
throw new Error(`HF API error ${res.status}: ${text.slice(0, 300)}`);
|
||
}
|
||
return res.json();
|
||
}
|
||
|
||
// ── Tool 1: launch_studio ─────────────────────────────────────────────
|
||
registerAppTool(
|
||
server,
|
||
"launch_studio",
|
||
{
|
||
title: "FineTune Studio",
|
||
description:
|
||
"Opens the FineTune Studio — an interactive dashboard for fine-tuning any Hugging Face model and chatting with it, all inside this conversation.",
|
||
inputSchema: {},
|
||
_meta: { ui: { resourceUri } },
|
||
},
|
||
async () => {
|
||
return { content: [{ type: "text" as const, text: "" }] };
|
||
}
|
||
);
|
||
|
||
// ── Tool 2: search_models ─────────────────────────────────────────────
|
||
registerAppTool(
|
||
server,
|
||
"search_models",
|
||
{
|
||
title: "Search HF Models",
|
||
description: "Search Hugging Face Hub for text-generation models.",
|
||
inputSchema: {
|
||
query: z.string().describe("Search query"),
|
||
limit: z.number().optional().describe("Max results, default 10"),
|
||
},
|
||
_meta: {},
|
||
},
|
||
async ({ query, limit }: { query: string; limit?: number }) => {
|
||
try {
|
||
const n = limit || 10;
|
||
const url = `${HF_API}/models?search=${encodeURIComponent(query)}&limit=${n}&filter=text-generation&sort=downloads`;
|
||
const data = await hfGet(url) as any[];
|
||
const models = data.map((m: any) => ({
|
||
modelId: m.modelId || m.id || "",
|
||
author: m.author || (m.modelId || m.id || "").split("/")[0] || "",
|
||
downloads: m.downloads || 0,
|
||
likes: m.likes || 0,
|
||
pipeline_tag: m.pipeline_tag || "text-generation",
|
||
tags: (m.tags || []).filter((t: string) => !t.startsWith("arxiv:")).slice(0, 6),
|
||
lastModified: m.lastModified || "",
|
||
}));
|
||
return { content: [{ type: "text" as const, text: JSON.stringify({ models }) }] };
|
||
} catch (err: unknown) {
|
||
return { content: [{ type: "text" as const, text: JSON.stringify({ models: [], error: String(err) }) }] };
|
||
}
|
||
}
|
||
);
|
||
|
||
// ── Tool 3: search_datasets ───────────────────────────────────────────
|
||
registerAppTool(
|
||
server,
|
||
"search_datasets",
|
||
{
|
||
title: "Search HF Datasets",
|
||
description: "Search Hugging Face Hub for datasets.",
|
||
inputSchema: {
|
||
query: z.string().describe("Search query"),
|
||
limit: z.number().optional().describe("Max results, default 10"),
|
||
},
|
||
_meta: {},
|
||
},
|
||
async ({ query, limit }: { query: string; limit?: number }) => {
|
||
try {
|
||
const n = limit || 10;
|
||
const url = `${HF_API}/datasets?search=${encodeURIComponent(query)}&limit=${n}&sort=downloads`;
|
||
const data = await hfGet(url) as any[];
|
||
const datasets = data.map((d: any) => ({
|
||
datasetId: d.id || "",
|
||
author: d.author || (d.id || "").split("/")[0] || "",
|
||
downloads: d.downloads || 0,
|
||
likes: d.likes || 0,
|
||
description: ((d.description || d.cardData?.description || "")).slice(0, 200),
|
||
tags: (d.tags || []).slice(0, 5),
|
||
size: d.cardData?.size_categories?.[0] || "",
|
||
}));
|
||
return { content: [{ type: "text" as const, text: JSON.stringify({ datasets }) }] };
|
||
} catch (err: unknown) {
|
||
return { content: [{ type: "text" as const, text: JSON.stringify({ datasets: [], error: String(err) }) }] };
|
||
}
|
||
}
|
||
);
|
||
|
||
// ── Tool 4: start_training ────────────────────────────────────────────
|
||
registerAppTool(
|
||
server,
|
||
"start_training",
|
||
{
|
||
title: "Start Training",
|
||
description: "Start a fine-tuning job on Hugging Face AutoTrain.",
|
||
inputSchema: {
|
||
base_model: z.string(),
|
||
dataset: z.string(),
|
||
dataset_split: z.string().optional(),
|
||
project_name: z.string(),
|
||
task: z.string().optional(),
|
||
hardware: z.string().optional(),
|
||
training_type: z.string().optional(),
|
||
chat_template: z.string().optional(),
|
||
hyperparameters: z.object({
|
||
epochs: z.number().optional(),
|
||
batch_size: z.number().optional(),
|
||
learning_rate: z.number().optional(),
|
||
block_size: z.number().optional(),
|
||
lora_r: z.number().optional(),
|
||
lora_alpha: z.number().optional(),
|
||
lora_dropout: z.number().optional(),
|
||
quantization: z.string().optional(),
|
||
gradient_accumulation: z.number().optional(),
|
||
mixed_precision: z.string().optional(),
|
||
warmup_ratio: z.number().optional(),
|
||
weight_decay: z.number().optional(),
|
||
target_modules: z.string().optional(),
|
||
max_steps: z.number().optional(),
|
||
}).optional(),
|
||
column_mapping: z.record(z.string()).optional(),
|
||
},
|
||
_meta: {},
|
||
},
|
||
async (params: any) => {
|
||
const token = getHFToken();
|
||
try {
|
||
if (!token) throw new Error("HF_TOKEN environment variable is not set.");
|
||
|
||
// Fetch username
|
||
const me = await hfGet("https://huggingface.co/api/whoami-v2") as any;
|
||
const username = me.name as string;
|
||
|
||
const hp = params.hyperparameters || {};
|
||
const trainingType = params.training_type || "sft";
|
||
|
||
// ── Find user's AutoTrain Advanced Space ───────────────────────────
|
||
// HF removed the old cloud API. Training now routes through an
|
||
// AutoTrain Advanced Space the user runs in their HF account.
|
||
// We look for it by the standard name; if missing, return setup instructions.
|
||
let spaceSubdomain = "";
|
||
const AUTOTRAIN_SPACE_NAME = "autotrain-advanced";
|
||
|
||
try {
|
||
const spaceInfo = await hfGet(`${HF_API}/spaces/${username}/${AUTOTRAIN_SPACE_NAME}`) as any;
|
||
// subdomain is the slug used in the .hf.space URL
|
||
spaceSubdomain = spaceInfo.subdomain || "";
|
||
} catch {
|
||
// Space not found — return actionable setup instructions
|
||
return {
|
||
content: [{
|
||
type: "text" as const,
|
||
text: JSON.stringify({
|
||
success: false,
|
||
setup_required: true,
|
||
error:
|
||
`AutoTrain Advanced Space not found in your HF account (@${username}). ` +
|
||
`HF retired the old cloud training API. You need a one-time setup: ` +
|
||
`go to https://huggingface.co/spaces/autotrain-projects/autotrain-advanced, ` +
|
||
`click "Duplicate this Space", name it "${AUTOTRAIN_SPACE_NAME}", ` +
|
||
`add your HF_TOKEN as a Space secret, then try again.`,
|
||
setup_url: "https://huggingface.co/spaces/autotrain-projects/autotrain-advanced",
|
||
}),
|
||
}],
|
||
};
|
||
}
|
||
|
||
if (!spaceSubdomain) {
|
||
throw new Error(`Could not determine the URL for your AutoTrain Space (@${username}/${AUTOTRAIN_SPACE_NAME}).`);
|
||
}
|
||
|
||
const spaceApiBase = `https://${spaceSubdomain}.hf.space`;
|
||
|
||
// ── Compute effective epochs (honouring max_steps if set) ────────
|
||
// AutoTrain's API doesn't expose max_steps as a config field — it is
|
||
// silently dropped by Pydantic. We implement it ourselves by fetching the
|
||
// dataset split size and converting max_steps → a fractional epoch count.
|
||
// effective_epochs = min(max_steps, steps_per_epoch × config_epochs)
|
||
// ÷ steps_per_epoch
|
||
const configEpochs = hp.epochs ?? 2;
|
||
const batchSize = hp.batch_size ?? 1;
|
||
const gradAcc = hp.gradient_accumulation ?? 4;
|
||
|
||
// When max_steps is set, fetch dataset size so we can calculate the minimum
|
||
// number of *integer* epochs needed for those steps to be reachable.
|
||
// AutoTrain requires epochs to be an integer and passes max_steps directly
|
||
// to HF Trainer which stops training at min(max_steps, total_steps).
|
||
let maxStepsEpochs: number | null = null; // integer epochs needed
|
||
if (hp.max_steps && hp.max_steps > 0 && params.dataset) {
|
||
const dsEnc = encodeURIComponent(params.dataset);
|
||
const splitKey = params.dataset_split || "train";
|
||
let trainRows: number | null = null;
|
||
|
||
try {
|
||
const r = await fetch(
|
||
`https://datasets-server.huggingface.co/info?dataset=${dsEnc}`,
|
||
{ headers: token ? { Authorization: `Bearer ${token}` } : {}, signal: AbortSignal.timeout(8000) }
|
||
);
|
||
const d = await r.json() as any;
|
||
const info = d?.dataset_info ?? {};
|
||
const splits = info.splits ?? info[Object.keys(info)[0]]?.splits ?? {};
|
||
trainRows = splits[splitKey]?.num_examples ?? splits.train?.num_examples ?? null;
|
||
} catch { /* ignore */ }
|
||
|
||
if (trainRows && trainRows > 0) {
|
||
const stepsPerEpoch = Math.ceil(trainRows / batchSize) / gradAcc;
|
||
// Minimum whole epochs needed so max_steps can actually be reached
|
||
maxStepsEpochs = Math.max(1, Math.ceil(hp.max_steps / stepsPerEpoch));
|
||
console.log(`[max_steps] rows=${trainRows} stepsPerEpoch=${stepsPerEpoch.toFixed(0)} maxStepsEpochs=${maxStepsEpochs}`);
|
||
} else {
|
||
console.warn(`[max_steps] could not fetch dataset size for "${params.dataset}" — using epochs=${configEpochs}`);
|
||
}
|
||
}
|
||
|
||
// ── Build the training request body ───────────────────────────────
|
||
// Note: `push_to_hub` must be at the TOP LEVEL — AutoTrain ignores it
|
||
// when nested inside `params` (silently dropped by Pydantic). Same pattern
|
||
// as `username` and `token` which also live at top level and flow through
|
||
// to the training config.
|
||
const body: Record<string, unknown> = {
|
||
username,
|
||
project_name: params.project_name,
|
||
task: params.task || "llm:sft",
|
||
base_model: params.base_model,
|
||
hub_dataset: params.dataset,
|
||
train_split: params.dataset_split || "train",
|
||
hardware: params.hardware || "spaces-a10g-large",
|
||
column_mapping: params.column_mapping || { text_column: "text" },
|
||
// push_to_hub must live inside a top-level "hub" object —
|
||
// the API ignores it when nested in "params" or bare at the top level.
|
||
hub: {
|
||
username,
|
||
token,
|
||
push_to_hub: true,
|
||
},
|
||
hub_model: `${username}/${params.project_name}`,
|
||
token, // also kept at top level for the Space auth header
|
||
params: {
|
||
epochs: maxStepsEpochs ?? Math.max(1, Math.round(configEpochs)),
|
||
batch_size: hp.batch_size ?? 1,
|
||
lr: hp.learning_rate ?? 0.0002,
|
||
block_size: hp.block_size ?? 1024,
|
||
peft: true,
|
||
quantization: hp.quantization ?? "int4",
|
||
mixed_precision: hp.mixed_precision ?? "bf16",
|
||
lora_r: hp.lora_r ?? 16,
|
||
lora_alpha: hp.lora_alpha ?? 32,
|
||
lora_dropout: hp.lora_dropout ?? 0.05,
|
||
gradient_accumulation: hp.gradient_accumulation ?? 4,
|
||
warmup_ratio: hp.warmup_ratio ?? 0.1,
|
||
weight_decay: hp.weight_decay ?? 0.01,
|
||
target_modules: hp.target_modules ?? "all-linear",
|
||
// "none" means a plain-text dataset — omit the field so AutoTrain
|
||
// skips apply_chat_template (avoids the ast.literal_eval SyntaxError).
|
||
...(params.chat_template && params.chat_template !== "none"
|
||
? { chat_template: params.chat_template }
|
||
: {}),
|
||
trainer: trainingType,
|
||
merge_adapter: true, // merge LoRA weights into base model before push
|
||
...(hp.max_steps && hp.max_steps > 0 ? { max_steps: Math.round(hp.max_steps) } : {}),
|
||
},
|
||
};
|
||
|
||
// ── POST to the user's AutoTrain Space API ─────────────────────────
|
||
const res = await fetch(`${spaceApiBase}/api/create_project`, {
|
||
method: "POST",
|
||
headers: {
|
||
"Authorization": `Bearer ${token}`,
|
||
"Content-Type": "application/json",
|
||
},
|
||
body: JSON.stringify(body),
|
||
});
|
||
|
||
if (!res.ok) {
|
||
const errText = await res.text().catch(() => "");
|
||
throw new Error(`AutoTrain Space API error ${res.status}: ${errText.slice(0, 300)}`);
|
||
}
|
||
|
||
const result = await res.json() as any;
|
||
|
||
// ── Inject push_to_hub=true directly into the training Space secret ──
|
||
//
|
||
// Root cause (confirmed by reading AutoTrain source):
|
||
// SpaceRunner._add_secrets() serialises LLMTrainingParams → JSON and stores
|
||
// it as the PARAMS env secret on the training sub-Space
|
||
// (e.g. autotrain-{project_name}). LLMTrainingParams.push_to_hub defaults
|
||
// to False, and because push_to_hub is in FIELDS_TO_EXCLUDE the API body
|
||
// can never override it. The training script reads os.environ["PARAMS"],
|
||
// deserialises it, and honours push_to_hub — so updating the secret to
|
||
// True is the definitive fix.
|
||
//
|
||
// Timing: create_project is synchronous — by the time it returns the
|
||
// training Space already exists with its secrets set. Updating a secret
|
||
// triggers a Space restart (HF behaviour). Since the Space is still in
|
||
// its 2-5 min build phase this restart is free; it re-builds and starts
|
||
// training with push_to_hub=True.
|
||
const trainingSpaceName = `autotrain-${params.project_name}`;
|
||
const textCol = (params.column_mapping as Record<string, string>)?.text_column || "text";
|
||
const effectiveEpochs = maxStepsEpochs ?? Math.max(1, Math.round(configEpochs));
|
||
|
||
// Build the LLMTrainingParams-compatible JSON that the training script
|
||
// deserialises from the PARAMS secret. Extra/unknown keys are fine —
|
||
// LLMTrainingParams(**json) ignores them via Pydantic's default behaviour.
|
||
const paramsSecret: Record<string, unknown> = {
|
||
model: params.base_model,
|
||
project_name: params.project_name,
|
||
username,
|
||
token,
|
||
data_path: params.dataset,
|
||
train_split: params.dataset_split || "train",
|
||
text_column: textCol,
|
||
push_to_hub: true, // ← THE FIX
|
||
log: "tensorboard",
|
||
epochs: effectiveEpochs,
|
||
batch_size: hp.batch_size ?? 1,
|
||
lr: hp.learning_rate ?? 0.0002,
|
||
block_size: hp.block_size ?? 1024,
|
||
peft: true,
|
||
quantization: hp.quantization ?? "int4",
|
||
mixed_precision: hp.mixed_precision ?? "bf16",
|
||
lora_r: hp.lora_r ?? 16,
|
||
lora_alpha: hp.lora_alpha ?? 32,
|
||
lora_dropout: hp.lora_dropout ?? 0.05,
|
||
gradient_accumulation: hp.gradient_accumulation ?? 4,
|
||
warmup_ratio: hp.warmup_ratio ?? 0.1,
|
||
weight_decay: hp.weight_decay ?? 0.01,
|
||
target_modules: hp.target_modules ?? "all-linear",
|
||
trainer: trainingType,
|
||
merge_adapter: true,
|
||
padding: "right",
|
||
chat_template: (params.chat_template && params.chat_template !== "none")
|
||
? params.chat_template : "none",
|
||
...(hp.max_steps && hp.max_steps > 0 ? { max_steps: Math.round(hp.max_steps) } : {}),
|
||
};
|
||
|
||
// POST to HF Hub Secrets API — same endpoint/method used by huggingface_hub
|
||
// Python library's add_space_secret().
|
||
try {
|
||
const secretRes = await fetch(
|
||
`https://huggingface.co/api/spaces/${username}/${trainingSpaceName}/secrets`,
|
||
{
|
||
method: "POST",
|
||
headers: {
|
||
"Authorization": `Bearer ${token}`,
|
||
"Content-Type": "application/json",
|
||
},
|
||
body: JSON.stringify({ key: "PARAMS", value: JSON.stringify(paramsSecret) }),
|
||
signal: AbortSignal.timeout(10000),
|
||
}
|
||
);
|
||
if (secretRes.ok) {
|
||
console.log(`[push_to_hub] ✅ PARAMS secret updated on ${trainingSpaceName} — push_to_hub=true`);
|
||
} else {
|
||
const errBody = await secretRes.text().catch(() => "");
|
||
console.warn(`[push_to_hub] ⚠️ Secret update HTTP ${secretRes.status}: ${errBody.slice(0, 200)}`);
|
||
}
|
||
} catch (secretErr) {
|
||
// Non-fatal — training will proceed, just may not push to Hub
|
||
console.warn(`[push_to_hub] ⚠️ Could not update PARAMS secret: ${secretErr}`);
|
||
}
|
||
|
||
return {
|
||
content: [{
|
||
type: "text" as const,
|
||
text: JSON.stringify({
|
||
success: true,
|
||
project_id: result.id || result.project_id || params.project_name,
|
||
// AutoTrain creates the training Space with "autotrain-" prefix
|
||
space_url: `https://huggingface.co/spaces/${username}/autotrain-${params.project_name}`,
|
||
autotrain_space_url: `https://huggingface.co/spaces/${username}/${AUTOTRAIN_SPACE_NAME}`,
|
||
model_url: `https://huggingface.co/${username}/${params.project_name}`,
|
||
username,
|
||
status: "starting",
|
||
}),
|
||
}],
|
||
};
|
||
} catch (err: unknown) {
|
||
return {
|
||
content: [{
|
||
type: "text" as const,
|
||
text: JSON.stringify({
|
||
success: false,
|
||
error: err instanceof Error ? err.message : String(err),
|
||
}),
|
||
}],
|
||
};
|
||
}
|
||
}
|
||
);
|
||
|
||
// ── Tool 5: check_training_status ─────────────────────────────────────
|
||
registerAppTool(
|
||
server,
|
||
"check_training_status",
|
||
{
|
||
title: "Check Training Status",
|
||
description: "Poll the status of a running AutoTrain fine-tuning job.",
|
||
inputSchema: {
|
||
project_name: z.string(),
|
||
username: z.string().optional(),
|
||
},
|
||
_meta: {},
|
||
},
|
||
async ({ project_name, username }: { project_name: string; username?: string }) => {
|
||
try {
|
||
let user = username;
|
||
if (!user) {
|
||
const me = await hfGet("https://huggingface.co/api/whoami-v2") as any;
|
||
user = me.name as string;
|
||
}
|
||
|
||
const stageMap: Record<string, string> = {
|
||
RUNNING: "training",
|
||
BUILDING: "starting",
|
||
STOPPED: "completed",
|
||
ERROR: "error",
|
||
SLEEPING: "starting",
|
||
PAUSED: "starting",
|
||
NO_APP_FILE: "starting",
|
||
CONFIG_ERROR: "error",
|
||
APP_STARTING: "starting",
|
||
};
|
||
|
||
// AutoTrain names the training Space "autotrain-{project_name}".
|
||
// After training it renames the space to "{project_name}" (drops prefix).
|
||
// We poll with the autotrain- prefix while it's running; the JWT in the
|
||
// container logs confirms: sub = /spaces/avi81/autotrain-{project_name}.
|
||
const trainingSpaceName = `autotrain-${project_name}`;
|
||
|
||
let rawStage = "BUILDING";
|
||
let status = "starting";
|
||
let spaceSubdomain = "";
|
||
try {
|
||
// Fetch the Space info (includes subdomain we need for logs)
|
||
const spaceInfo = await hfGet(`${HF_API}/spaces/${user}/${trainingSpaceName}`) as any;
|
||
spaceSubdomain = spaceInfo.subdomain || "";
|
||
// Also fetch runtime for the live stage
|
||
const runtime = await hfGet(`${HF_API}/spaces/${user}/${trainingSpaceName}/runtime`) as any;
|
||
rawStage = runtime.stage || "BUILDING";
|
||
status = stageMap[rawStage] || "starting";
|
||
} catch (runtimeErr: unknown) {
|
||
const msg = String(runtimeErr);
|
||
// 404 = Space not yet provisioned by AutoTrain — keep as "starting"
|
||
if (!msg.includes("404")) {
|
||
status = "error";
|
||
rawStage = "ERROR";
|
||
}
|
||
}
|
||
// AutoTrain pauses its Space when training finishes, so PAUSED / STOPPED
|
||
// alone can't distinguish "not-started-yet" from "just-finished".
|
||
// We disambiguate by looking for shutdown / cleanup log lines.
|
||
// This is applied AFTER logs are collected below, so we store the
|
||
// preliminary status and patch it if needed.
|
||
const preliminaryStatus = status;
|
||
|
||
// ── Fetch container logs via HF Hub SSE endpoint ─────────────────
|
||
// Correct endpoint is /logs/run (not /logs).
|
||
// Streams SSE events: data: {"type":"stdout","text":"<line>\n"}
|
||
let logs = "";
|
||
let logsDebug = "";
|
||
const token = getHFToken();
|
||
const logsUrl = `${HF_API}/spaces/${user}/${trainingSpaceName}/logs/run`;
|
||
try {
|
||
const logRes = await fetch(logsUrl, {
|
||
headers: token ? { Authorization: `Bearer ${token}` } : {},
|
||
signal: AbortSignal.timeout(8000),
|
||
});
|
||
logsDebug = `[logs fetch: HTTP ${logRes.status}]`;
|
||
if (logRes.ok && logRes.body) {
|
||
const reader = logRes.body.getReader();
|
||
const decoder = new TextDecoder();
|
||
let raw = "";
|
||
|
||
// Hard-cancel the reader after 4 s — reader.cancel() causes
|
||
// reader.read() to reject, which is our signal to stop.
|
||
const cancelTimer = setTimeout(
|
||
() => reader.cancel("timeout").catch(() => {}),
|
||
4000
|
||
);
|
||
try {
|
||
while (true) {
|
||
const { done, value } = await reader.read();
|
||
if (done) break;
|
||
raw += decoder.decode(value, { stream: true });
|
||
}
|
||
} catch {
|
||
// reader.cancel() threw — normal exit path
|
||
} finally {
|
||
clearTimeout(cancelTimer);
|
||
}
|
||
|
||
logsDebug += ` [raw bytes: ${raw.length}]`;
|
||
|
||
// Parse SSE lines: "data: {"type":"stdout","text":"log line\n"}"
|
||
const ansiRe = /\x1b\[[0-9;]*m/g;
|
||
logs = raw
|
||
.split("\n")
|
||
.filter(l => l.startsWith("data:"))
|
||
.map(l => {
|
||
const payload = l.slice(5).trim();
|
||
if (!payload) return "";
|
||
try {
|
||
const obj = JSON.parse(payload) as any;
|
||
if (obj.type === "run") return "";
|
||
return String(obj.text ?? obj.data ?? payload)
|
||
.replace(/\n$/, "")
|
||
.replace(ansiRe, "");
|
||
} catch {
|
||
return payload.replace(ansiRe, "");
|
||
}
|
||
})
|
||
.filter(Boolean)
|
||
.join("\n");
|
||
} else {
|
||
const errBody = await logRes.text().catch(() => "");
|
||
logsDebug += ` [body: ${errBody.slice(0, 200)}]`;
|
||
}
|
||
} catch (logErr: unknown) {
|
||
logsDebug = `[logs fetch error: ${String(logErr)}]`;
|
||
}
|
||
|
||
// ── Patch status using log content ────────────────────────────────
|
||
// AutoTrain pauses the Space after training completes. The PAUSED stage
|
||
// alone is ambiguous (could be "not started yet"). If the logs contain
|
||
// AutoTrain / HF Trainer shutdown signatures we know training finished.
|
||
if (preliminaryStatus !== "training" && preliminaryStatus !== "error") {
|
||
const completionRe = /SIGTERM|graceful\.exit|Application\.shutdown|Training complete|model.*pushed|Pausing space|pause_space/i;
|
||
if (completionRe.test(logs)) {
|
||
status = "completed";
|
||
}
|
||
}
|
||
|
||
// Parse metrics from logs — handle both 'key': val and key: val formats.
|
||
// AutoTrain logs: {'loss': 9.73, 'learning_rate': 0.0002, 'epoch': 0.02}
|
||
// Always use the LAST match so stale early-epoch values don't stick.
|
||
const lastMatch = (text: string, ...patterns: RegExp[]) => {
|
||
for (const re of patterns) {
|
||
const all = [...text.matchAll(new RegExp(re.source, re.flags.replace("g","") + "g"))];
|
||
if (all.length > 0) return all[all.length - 1];
|
||
}
|
||
return null;
|
||
};
|
||
const lossMatch = lastMatch(logs, /'loss':\s*([0-9]+\.[0-9]+)/, /(?:train_)?loss[:\s=]+([0-9]+\.[0-9]+)/i);
|
||
const epochMatch = lastMatch(logs, /'epoch':\s*([0-9]+(?:\.[0-9]*)?)/, /epoch[:\s]+([0-9]+\.?[0-9]*)/i);
|
||
const lrMatch = lastMatch(logs, /'learning_rate':\s*([0-9.e+\-]+)/, /(?:learning_rate|lr)[:\s=]+([0-9.e+\-]+)/i);
|
||
// tqdm uses \r to overwrite lines, so only the initial 0/N line arrives cleanly.
|
||
// Send totalSteps to the client so it can compute current step as:
|
||
// step = round((epoch / config.epochs) * totalSteps)
|
||
// (epoch goes 0→config.epochs continuously, totalSteps spans all epochs)
|
||
const totalStepsMatch = lastMatch(logs, /\d+\/(\d+)\s*\[/);
|
||
const totalSteps = totalStepsMatch ? parseInt(totalStepsMatch[1]) : null;
|
||
|
||
// ── Auto-publish model when training completes ─────────────────────
|
||
// AutoTrain hardcodes private=True when creating the model repo.
|
||
// We flip it to public so the inference Space can access it.
|
||
// This is a one-shot operation: the HF API ignores it if already public.
|
||
let madePublic = false;
|
||
let inferenceSpaceResult: { space_url: string; space_name: string } | null = null;
|
||
|
||
if (status === "completed" && token) {
|
||
try {
|
||
const pubRes = await fetch(
|
||
`https://huggingface.co/api/models/${user}/${project_name}/settings`,
|
||
{
|
||
method: "PUT",
|
||
headers: {
|
||
Authorization: `Bearer ${token}`,
|
||
"Content-Type": "application/json",
|
||
},
|
||
body: JSON.stringify({ private: false }),
|
||
signal: AbortSignal.timeout(8000),
|
||
}
|
||
);
|
||
madePublic = pubRes.ok;
|
||
if (pubRes.ok) {
|
||
console.log(`[visibility] ✅ ${user}/${project_name} set to public`);
|
||
} else {
|
||
const t = await pubRes.text().catch(() => "");
|
||
console.warn(`[visibility] ⚠️ HTTP ${pubRes.status}: ${t.slice(0, 200)}`);
|
||
}
|
||
} catch (pubErr) {
|
||
console.warn(`[visibility] ⚠️ ${pubErr}`);
|
||
}
|
||
|
||
// ── Deploy inference Space ─────────────────────────────────────
|
||
// HF serverless inference (inferenceProviderMapping) no longer serves
|
||
// arbitrary custom fine-tuned models. We work around this by deploying
|
||
// a small Gradio Space that loads the model with transformers directly.
|
||
// The Space name is inference-{project_name}.
|
||
try {
|
||
inferenceSpaceResult = await deployInferenceSpace(user, project_name, token);
|
||
} catch (spaceErr) {
|
||
console.warn(`[inference-space] deployment failed: ${spaceErr}`);
|
||
}
|
||
}
|
||
|
||
return {
|
||
content: [{
|
||
type: "text" as const,
|
||
text: JSON.stringify({
|
||
status,
|
||
stage: rawStage,
|
||
training_space: `https://huggingface.co/spaces/${user}/${trainingSpaceName}`,
|
||
logs: (logsDebug ? logsDebug + "\n" : "") + logs.slice(-4000),
|
||
metrics: {
|
||
loss: lossMatch ? parseFloat(lossMatch[1]) : null,
|
||
epoch: epochMatch ? parseFloat(epochMatch[1]) : null,
|
||
totalSteps,
|
||
learning_rate: lrMatch ? parseFloat(lrMatch[1]) : null,
|
||
},
|
||
...(status === "completed" ? {
|
||
model_made_public: madePublic,
|
||
inference_space_url: inferenceSpaceResult?.space_url ?? null,
|
||
inference_space_name: inferenceSpaceResult?.space_name ?? null,
|
||
} : {}),
|
||
}),
|
||
}],
|
||
};
|
||
} catch (err: unknown) {
|
||
return {
|
||
content: [{
|
||
type: "text" as const,
|
||
text: JSON.stringify({ status: "error", error: String(err) }),
|
||
}],
|
||
};
|
||
}
|
||
}
|
||
);
|
||
|
||
// ── Tool 6: chat_with_model ───────────────────────────────────────────
|
||
registerAppTool(
|
||
server,
|
||
"chat_with_model",
|
||
{
|
||
title: "Chat with Model",
|
||
description: "Send messages to a Hugging Face model via the inference API and get a response.",
|
||
inputSchema: {
|
||
model_id: z.string().describe("HF model ID, e.g. meta-llama/Llama-3.2-3B-Instruct"),
|
||
messages: z.array(z.object({
|
||
role: z.enum(["system", "user", "assistant"]),
|
||
content: z.string(),
|
||
})),
|
||
parameters: z.object({
|
||
temperature: z.number().optional(),
|
||
top_p: z.number().optional(),
|
||
max_tokens: z.number().optional(),
|
||
}).optional(),
|
||
},
|
||
_meta: {},
|
||
},
|
||
async ({ model_id, messages, parameters }: {
|
||
model_id: string;
|
||
messages: Array<{ role: "system" | "user" | "assistant"; content: string }>;
|
||
parameters?: { temperature?: number; top_p?: number; max_tokens?: number };
|
||
}) => {
|
||
try {
|
||
const token = getHFToken();
|
||
if (!token) throw new Error("HF_TOKEN environment variable is not set.");
|
||
|
||
const chatBody = {
|
||
model: model_id,
|
||
messages,
|
||
temperature: parameters?.temperature ?? 0.7,
|
||
max_tokens: parameters?.max_tokens ?? 512,
|
||
top_p: parameters?.top_p ?? 0.9,
|
||
stream: false,
|
||
};
|
||
|
||
// ── Inference strategy — four attempts in order ────────────────────
|
||
//
|
||
// 0. Dedicated Inference Space ({owner}/inference-{model_name}).
|
||
// We auto-deploy this Gradio Space after training completes.
|
||
// It loads the model with transformers directly — works for ANY
|
||
// Hub model regardless of HF's inferenceProviderMapping.
|
||
// This is the primary path for custom fine-tuned models.
|
||
//
|
||
// 1. HF Inference Router (router.huggingface.co).
|
||
// Works for popular public models supported by providers.
|
||
//
|
||
// 2. HF per-model Messages API (api-inference.huggingface.co /v1/chat).
|
||
// Works for public models indexed by HF's TGI fleet.
|
||
//
|
||
// 3. HF Legacy Text-Generation API (api-inference.huggingface.co /models).
|
||
// Last resort for public models that support the older endpoint.
|
||
|
||
let responseText = "";
|
||
|
||
// ── Attempt 0: Dedicated Inference Space (Gradio 4.x raw HTTP API) ──
|
||
// For custom fine-tuned models we auto-deploy a Gradio Space post-training
|
||
// named {owner}/inference-{model_name}.
|
||
//
|
||
// Gradio 4.x dropped the old /api/predict sync endpoint. The current API is:
|
||
// POST /gradio_api/call/predict → {"event_id": "abc"}
|
||
// GET /gradio_api/call/predict/{event_id} → SSE "event: complete\ndata: [...]"
|
||
//
|
||
// We use raw HTTP so we have full control over errors (the @gradio/client
|
||
// package silently swallows connection errors which masked failures).
|
||
if (!responseText && model_id.includes("/")) {
|
||
const [owner, modelName] = model_id.split("/");
|
||
const inferSpaceName = `inference-${modelName}`;
|
||
|
||
// ── 1. Does the Space exist? Get subdomain while we're at it. ────
|
||
const spaceInfoRes = await fetch(
|
||
`${HF_API}/spaces/${owner}/${inferSpaceName}`,
|
||
{ headers: { Authorization: `Bearer ${token}` }, signal: AbortSignal.timeout(8000) },
|
||
).catch(() => null);
|
||
|
||
if (spaceInfoRes?.ok) {
|
||
const spaceData = await spaceInfoRes.json() as any;
|
||
const subdomain = spaceData.subdomain as string | undefined;
|
||
|
||
// ── 2. Check runtime stage ──────────────────────────────────
|
||
const runtimeRes = await fetch(
|
||
`${HF_API}/spaces/${owner}/${inferSpaceName}/runtime`,
|
||
{ headers: { Authorization: `Bearer ${token}` }, signal: AbortSignal.timeout(8000) },
|
||
).catch(() => null);
|
||
|
||
// If the runtime endpoint returns 404, the Space was just created and
|
||
// hasn't started its build cycle yet — treat it as BUILDING.
|
||
const stage: string = runtimeRes?.ok
|
||
? (((await runtimeRes.json()) as any).stage || "UNKNOWN")
|
||
: "BUILDING";
|
||
|
||
console.log(`[chat] inference-space ${owner}/${inferSpaceName} stage=${stage}`);
|
||
|
||
if (["BUILDING", "APP_STARTING", "STARTING", "NO_APP_FILE", "UNKNOWN"].includes(stage)) {
|
||
throw new Error(
|
||
`⏳ Your inference Space (${owner}/${inferSpaceName}) is still building — ` +
|
||
`takes 2–3 min on first deploy. Please wait and try again.\n` +
|
||
`Track progress: https://huggingface.co/spaces/${owner}/${inferSpaceName}`,
|
||
);
|
||
}
|
||
|
||
if (stage === "SLEEPING") {
|
||
throw new Error(
|
||
`😴 Your inference Space (${owner}/${inferSpaceName}) was sleeping and is now waking up. ` +
|
||
`Please wait ~1 minute and try again.`,
|
||
);
|
||
}
|
||
|
||
if (stage === "RUNNING" && subdomain) {
|
||
// ── 3. POST to Gradio 4.x queue endpoint ─────────────────
|
||
const gradioBase = `https://${subdomain}.hf.space/gradio_api/call/predict`;
|
||
const submitRes = await fetch(gradioBase, {
|
||
method: "POST",
|
||
headers: { "Content-Type": "application/json" },
|
||
body: JSON.stringify({
|
||
// gr.Interface uses positional parameters — same order as fn signature:
|
||
// predict(messages_json, max_tokens, temperature)
|
||
data: [
|
||
JSON.stringify(messages),
|
||
chatBody.max_tokens,
|
||
chatBody.temperature,
|
||
],
|
||
}),
|
||
signal: AbortSignal.timeout(30000),
|
||
});
|
||
|
||
if (!submitRes.ok) {
|
||
const b = await submitRes.text().catch(() => "");
|
||
throw new Error(
|
||
`Inference Space submit failed (HTTP ${submitRes.status}): ${b.slice(0, 200)}`,
|
||
);
|
||
}
|
||
|
||
const { event_id } = await submitRes.json() as { event_id: string };
|
||
if (!event_id) throw new Error("Inference Space did not return an event_id");
|
||
|
||
// ── 4. Stream SSE result ───────────────────────────────────
|
||
const sseRes = await fetch(`${gradioBase}/${event_id}`, {
|
||
signal: AbortSignal.timeout(180_000), // 3 min — first inference on CPU is slow
|
||
});
|
||
|
||
if (!sseRes.ok) {
|
||
throw new Error(`Inference Space SSE failed (HTTP ${sseRes.status})`);
|
||
}
|
||
|
||
const reader = sseRes.body!.getReader();
|
||
const decoder = new TextDecoder();
|
||
let buf = "";
|
||
|
||
parseSSE: while (true) {
|
||
const { done, value } = await reader.read();
|
||
if (done) break;
|
||
buf += decoder.decode(value, { stream: true });
|
||
|
||
for (const chunk of buf.split("\n\n")) {
|
||
const eventLine = chunk.split("\n").find(l => l.startsWith("event:"))?.slice(7).trim() ?? "";
|
||
const dataLine = chunk.split("\n").find(l => l.startsWith("data:"))?.slice(6).trim() ?? "";
|
||
|
||
if (eventLine === "error") {
|
||
throw new Error(`Inference Space error: ${dataLine.slice(0, 200)}`);
|
||
}
|
||
if (eventLine === "complete" && dataLine) {
|
||
const parsed = JSON.parse(dataLine);
|
||
const resp = String(parsed?.[0] ?? "");
|
||
if (resp.startsWith("Error:")) {
|
||
throw new Error(`Inference Space returned: ${resp}`);
|
||
}
|
||
responseText = resp;
|
||
reader.cancel().catch(() => {});
|
||
break parseSSE;
|
||
}
|
||
}
|
||
|
||
// Keep only the last incomplete chunk
|
||
buf = buf.includes("\n\n") ? buf.slice(buf.lastIndexOf("\n\n") + 2) : buf;
|
||
}
|
||
} else if (!["STOPPED", "PAUSED", "ERROR", "CONFIG_ERROR"].includes(stage)) {
|
||
// Unexpected stage — tell the user rather than silently failing
|
||
throw new Error(
|
||
`Inference Space (${owner}/${inferSpaceName}) is in an unexpected state: ${stage}. ` +
|
||
`Check https://huggingface.co/spaces/${owner}/${inferSpaceName}`,
|
||
);
|
||
}
|
||
// STOPPED / PAUSED / ERROR / CONFIG_ERROR → fall through to HF inference attempts
|
||
}
|
||
// Space doesn't exist (404) → fall through silently to HF inference attempts
|
||
}
|
||
|
||
// ── Attempt 1: HF Inference Router ────────────────────────────────
|
||
if (!responseText) {
|
||
const routerModel = model_id.includes(":") ? model_id : `${model_id}:fastest`;
|
||
const res1 = await fetch(HF_ROUTER_URL, {
|
||
method: "POST",
|
||
headers: { "Authorization": `Bearer ${token}`, "Content-Type": "application/json" },
|
||
body: JSON.stringify({ ...chatBody, model: routerModel }),
|
||
});
|
||
if (res1.ok) {
|
||
const d = await res1.json() as any;
|
||
responseText = d.choices?.[0]?.message?.content || "";
|
||
} else {
|
||
if (res1.status === 401 || res1.status === 403) {
|
||
throw new Error(`Access denied for "${model_id}". Check your HF_TOKEN and model license.`);
|
||
}
|
||
// 404 / "not supported by any provider" → try next endpoint
|
||
}
|
||
}
|
||
|
||
// ── Attempt 2: Messages API (chat-completions style) ──────────────
|
||
if (!responseText) {
|
||
const url2 = `https://api-inference.huggingface.co/models/${model_id}/v1/chat/completions`;
|
||
const res2 = await fetch(url2, {
|
||
method: "POST",
|
||
headers: { "Authorization": `Bearer ${token}`, "Content-Type": "application/json" },
|
||
body: JSON.stringify({ ...chatBody, model: model_id }),
|
||
});
|
||
if (res2.ok) {
|
||
const d2 = await res2.json() as any;
|
||
responseText = d2.choices?.[0]?.message?.content || "";
|
||
}
|
||
// 404 or 503 → fall through to attempt 3
|
||
}
|
||
|
||
// ── Attempt 3: Legacy Text-Generation API ─────────────────────────
|
||
if (!responseText) {
|
||
const promptParts: string[] = [];
|
||
for (const msg of messages) {
|
||
if (msg.role === "system") {
|
||
promptParts.push(`<|im_start|>system\n${msg.content}<|im_end|>`);
|
||
} else if (msg.role === "user") {
|
||
promptParts.push(`<|im_start|>user\n${msg.content}<|im_end|>`);
|
||
} else {
|
||
promptParts.push(`<|im_start|>assistant\n${msg.content}<|im_end|>`);
|
||
}
|
||
}
|
||
promptParts.push("<|im_start|>assistant\n");
|
||
const prompt = promptParts.join("\n");
|
||
|
||
const url3 = `https://api-inference.huggingface.co/models/${model_id}`;
|
||
const res3 = await fetch(url3, {
|
||
method: "POST",
|
||
headers: { "Authorization": `Bearer ${token}`, "Content-Type": "application/json" },
|
||
body: JSON.stringify({
|
||
inputs: prompt,
|
||
parameters: {
|
||
max_new_tokens: chatBody.max_tokens,
|
||
temperature: chatBody.temperature,
|
||
top_p: chatBody.top_p,
|
||
return_full_text: false,
|
||
do_sample: true,
|
||
},
|
||
}),
|
||
});
|
||
|
||
if (!res3.ok) {
|
||
const b3 = await res3.text().catch(() => "");
|
||
if (res3.status === 404) {
|
||
// Check whether we deployed an inference Space for this model
|
||
const hasInferenceSpace = model_id.includes("/");
|
||
const [owner2, modelName2] = hasInferenceSpace ? model_id.split("/") : ["", ""];
|
||
throw new Error(
|
||
hasInferenceSpace
|
||
? `Model "${model_id}" is not yet available via HF serverless inference. ` +
|
||
`An inference Space (${owner2}/inference-${modelName2}) was deployed for you — ` +
|
||
`it takes 2–3 min to start. Check https://huggingface.co/spaces/${owner2}/inference-${modelName2}`
|
||
: `Model "${model_id}" was not found on the Hugging Face Hub. ` +
|
||
`Verify it exists at https://huggingface.co/${model_id}`
|
||
);
|
||
}
|
||
if (res3.status === 503 || b3.includes("loading") || b3.includes("estimated_time")) {
|
||
const eta = (() => { try { return JSON.parse(b3).estimated_time; } catch { return null; } })();
|
||
throw new Error(
|
||
`Model "${model_id}" is warming up${eta ? ` (~${Math.ceil(eta)}s)` : ""}. ` +
|
||
`Please try again in a moment.`
|
||
);
|
||
}
|
||
throw new Error(`Inference failed (HTTP ${res3.status}): ${b3.slice(0, 300)}`);
|
||
}
|
||
|
||
const d3 = await res3.json() as any;
|
||
const raw = Array.isArray(d3) ? d3[0]?.generated_text : d3?.generated_text;
|
||
responseText = (raw || "").replace(/^<\|im_start\|>assistant\n?/, "").trim();
|
||
}
|
||
|
||
if (!responseText) {
|
||
throw new Error(`No response from "${model_id}". The model may not support text generation via the HF Inference API.`);
|
||
}
|
||
|
||
return {
|
||
content: [{
|
||
type: "text" as const,
|
||
text: JSON.stringify({ text: responseText, model_id }),
|
||
}],
|
||
};
|
||
} catch (err: unknown) {
|
||
return {
|
||
content: [{
|
||
type: "text" as const,
|
||
text: JSON.stringify({ error: err instanceof Error ? err.message : String(err) }),
|
||
}],
|
||
};
|
||
}
|
||
}
|
||
);
|
||
|
||
// ── Tool 7: deploy_inference_space ───────────────────────────────────────
|
||
registerAppTool(
|
||
server,
|
||
"deploy_inference_space",
|
||
{
|
||
title: "Deploy Inference Space",
|
||
description:
|
||
"Manually deploy (or redeploy) a Gradio inference Space for a fine-tuned model. " +
|
||
"The Space loads the model with transformers and exposes an API endpoint, " +
|
||
"bypassing the HF serverless inference provider limitation.",
|
||
inputSchema: {
|
||
project_name: z.string().describe("The project / model name (without username prefix)"),
|
||
},
|
||
_meta: {},
|
||
},
|
||
async ({ project_name }: { project_name: string }) => {
|
||
try {
|
||
const token = getHFToken();
|
||
if (!token) throw new Error("HF_TOKEN is not set.");
|
||
const me = await hfGet("https://huggingface.co/api/whoami-v2") as any;
|
||
const username = me.name as string;
|
||
|
||
const result = await deployInferenceSpace(username, project_name, token);
|
||
|
||
return {
|
||
content: [{
|
||
type: "text" as const,
|
||
text: JSON.stringify({
|
||
success: true,
|
||
message: "✅ Inference Space deployed! It takes 2–3 min to build and load the model.",
|
||
space_url: result.space_url,
|
||
space_name: result.space_name,
|
||
}),
|
||
}],
|
||
};
|
||
} catch (err: unknown) {
|
||
return {
|
||
content: [{
|
||
type: "text" as const,
|
||
text: JSON.stringify({ success: false, error: err instanceof Error ? err.message : String(err) }),
|
||
}],
|
||
};
|
||
}
|
||
}
|
||
);
|
||
|
||
// ── Tool 8: patch_autotrain_space ─────────────────────────────────────
|
||
// Fixes the push_to_hub=False bug present in autotrain-advanced ≤ 0.8.36.
|
||
// Strategy: add a sitecustomize.py monkey-patch to the user's AutoTrain Space
|
||
// repo so Python auto-imports it at startup and forces push_to_hub=True before
|
||
// any training job runs.
|
||
registerAppTool(
|
||
server,
|
||
"patch_autotrain_space",
|
||
{
|
||
title: "Patch AutoTrain Space (Fix Push to Hub)",
|
||
description:
|
||
"One-click fix for the autotrain-advanced push_to_hub=False bug. " +
|
||
"Commits a sitecustomize.py monkey-patch and updated Dockerfile to your " +
|
||
"AutoTrain Space so future training jobs automatically push to the Hub.",
|
||
inputSchema: {},
|
||
_meta: {},
|
||
},
|
||
async () => {
|
||
try {
|
||
const token = getHFToken();
|
||
if (!token) throw new Error("HF_TOKEN is not set.");
|
||
|
||
const me = await hfGet("https://huggingface.co/api/whoami-v2") as any;
|
||
const username = me.name as string;
|
||
const AUTOTRAIN_SPACE_NAME = "autotrain-advanced";
|
||
const spaceRepo = `${username}/${AUTOTRAIN_SPACE_NAME}`;
|
||
|
||
// ── Why the previous sitecustomize.py approach failed ─────────────
|
||
// Python's site.getsitepackages()[0] often points to a dist-packages
|
||
// directory the autotrain interpreter doesn't check first. The reliable
|
||
// alternative is to DIRECTLY REWRITE the installed params.py source file
|
||
// on disk after each `pip install`, then also clear the .pyc cache so
|
||
// Python recompiles from the patched source.
|
||
|
||
// ── 1. fix_push_to_hub.py — run once after pip install ─────────────
|
||
const fixScript = `"""
|
||
Patch autotrain-advanced so that every training job pushes to the Hub.
|
||
|
||
Root cause (all versions <= 0.8.36):
|
||
PARAMS["llm"] is built from LLMTrainingParams().model_dump() which always
|
||
includes push_to_hub=False. The guard in _munge_common_params:
|
||
if "push_to_hub" not in _params:
|
||
_params["push_to_hub"] = True
|
||
never fires because the key IS already present (as False).
|
||
|
||
Fix: rewrite the installed params.py so the guard becomes unconditional.
|
||
"""
|
||
import importlib, inspect, os, re, sys
|
||
|
||
# ── locate params.py ──────────────────────────────────────────────────
|
||
import autotrain.app.params as _m
|
||
src = inspect.getfile(_m)
|
||
print(f"[push_to_hub patch] target: {src}", flush=True)
|
||
|
||
with open(src) as fh:
|
||
code = fh.read()
|
||
|
||
# ── apply patch ───────────────────────────────────────────────────────
|
||
# Match the conditional block regardless of exact indentation.
|
||
PATTERN = r'([ \\t]+)if ["\\'\\']push_to_hub["\\'\\'] not in _params:[\\s\\S]*?_params["\\'\\']["\\'\\']push_to_hub["\\'\\']["\\'\\'] = True'
|
||
REPLACEMENT = r'\\1_params["push_to_hub"] = True # patched: force hub push'
|
||
|
||
patched_code, n = re.subn(PATTERN, REPLACEMENT, code)
|
||
|
||
if n > 0:
|
||
with open(src, "w") as fh:
|
||
fh.write(patched_code)
|
||
# Bust the bytecode cache so the interpreter picks up the new source.
|
||
cache_dir = os.path.join(os.path.dirname(src), "__pycache__")
|
||
if os.path.isdir(cache_dir):
|
||
removed = [
|
||
os.remove(os.path.join(cache_dir, f))
|
||
for f in os.listdir(cache_dir)
|
||
if f.startswith("params") and f.endswith(".pyc")
|
||
]
|
||
print(f"[push_to_hub patch] ✅ params.py patched ({n} substitution(s))", flush=True)
|
||
elif "# patched: force hub push" in code:
|
||
print("[push_to_hub patch] ✅ params.py already patched", flush=True)
|
||
else:
|
||
# Fallback: try a simpler literal replacement (handles minor whitespace diffs).
|
||
SIMPLE_OLD = ' if "push_to_hub" not in _params:\\n _params["push_to_hub"] = True'
|
||
SIMPLE_NEW = ' _params["push_to_hub"] = True # patched: force hub push'
|
||
if SIMPLE_OLD in code:
|
||
with open(src, "w") as fh:
|
||
fh.write(code.replace(SIMPLE_OLD, SIMPLE_NEW))
|
||
print("[push_to_hub patch] ✅ params.py patched (literal fallback)", flush=True)
|
||
else:
|
||
print(f"[push_to_hub patch] ⚠️ Could not locate patch target in {src}.", file=sys.stderr, flush=True)
|
||
print("[push_to_hub patch] Inspecting file…", file=sys.stderr, flush=True)
|
||
for i, line in enumerate(code.splitlines(), 1):
|
||
if "push_to_hub" in line:
|
||
print(f" L{i}: {line!r}", file=sys.stderr, flush=True)
|
||
`;
|
||
|
||
// ── 2. Dockerfile ───────────────────────────────────────────────────
|
||
const dockerfileContent = `FROM huggingface/autotrain-advanced:latest
|
||
COPY fix_push_to_hub.py /tmp/fix_push_to_hub.py
|
||
CMD pip uninstall -y autotrain-advanced && \\
|
||
pip install -U autotrain-advanced && \\
|
||
python /tmp/fix_push_to_hub.py && \\
|
||
autotrain app --host 0.0.0.0 --port 7860 --workers 1
|
||
`;
|
||
|
||
// ── 3. Commit both files to the Space repo ─────────────────────────
|
||
const toB64 = (s: string) => Buffer.from(s, "utf-8").toString("base64");
|
||
|
||
const ndjson = [
|
||
JSON.stringify({
|
||
key: "header",
|
||
value: {
|
||
summary: "Fix push_to_hub=False: rewrite params.py after pip install",
|
||
description:
|
||
"Replace fragile sitecustomize.py approach with direct source-file " +
|
||
"rewrite of autotrain/app/params.py after each pip install. " +
|
||
"Clears .pyc cache so the patched source is used immediately.",
|
||
},
|
||
}),
|
||
JSON.stringify({
|
||
key: "file",
|
||
value: { path: "fix_push_to_hub.py", encoding: "base64", content: toB64(fixScript) },
|
||
}),
|
||
JSON.stringify({
|
||
key: "file",
|
||
value: { path: "Dockerfile", encoding: "base64", content: toB64(dockerfileContent) },
|
||
}),
|
||
].join("\n");
|
||
|
||
const commitRes = await fetch(
|
||
`https://huggingface.co/api/spaces/${spaceRepo}/commit/main`,
|
||
{
|
||
method: "POST",
|
||
headers: {
|
||
Authorization: `Bearer ${token}`,
|
||
"Content-Type": "application/x-ndjson",
|
||
},
|
||
body: ndjson,
|
||
}
|
||
);
|
||
|
||
if (!commitRes.ok) {
|
||
const errText = await commitRes.text().catch(() => "");
|
||
throw new Error(`HF Hub commit failed (HTTP ${commitRes.status}): ${errText.slice(0, 400)}`);
|
||
}
|
||
|
||
const commitData = await commitRes.json() as any;
|
||
|
||
return {
|
||
content: [{
|
||
type: "text" as const,
|
||
text: JSON.stringify({
|
||
success: true,
|
||
message:
|
||
"✅ Patch committed! Your AutoTrain Space is rebuilding (~2-3 min). " +
|
||
"Once it's back online, new training jobs will automatically push " +
|
||
"the fine-tuned model to the Hub.",
|
||
space_url: `https://huggingface.co/spaces/${spaceRepo}`,
|
||
commit_url: commitData.commitUrl || `https://huggingface.co/spaces/${spaceRepo}/commit/main`,
|
||
next_step:
|
||
"Wait for the Space status to go green, then start a new training run.",
|
||
}),
|
||
}],
|
||
};
|
||
} catch (err: unknown) {
|
||
return {
|
||
content: [{
|
||
type: "text" as const,
|
||
text: JSON.stringify({
|
||
success: false,
|
||
error: err instanceof Error ? err.message : String(err),
|
||
}),
|
||
}],
|
||
};
|
||
}
|
||
}
|
||
);
|
||
|
||
// ── Resource ──────────────────────────────────────────────────────────
|
||
registerAppResource(
|
||
server,
|
||
resourceUri,
|
||
resourceUri,
|
||
{ mimeType: RESOURCE_MIME_TYPE, _meta: { ui: { csp: { resourceDomains: [] }, permissions: { openLinks: {} } } } },
|
||
async () => {
|
||
const html = await fs.readFile(
|
||
path.join(__dirname, "dist", "widget.html"),
|
||
"utf-8"
|
||
);
|
||
return { contents: [{ uri: resourceUri, mimeType: RESOURCE_MIME_TYPE, text: html }] };
|
||
}
|
||
);
|
||
|
||
// ── Express server ────────────────────────────────────────────────────
|
||
const expressApp = express();
|
||
expressApp.use(cors());
|
||
expressApp.use(express.json({ limit: "10mb" }));
|
||
|
||
// Health check — required by many cloud platforms and MCP hosts
|
||
expressApp.get("/", (_req, res) => {
|
||
res.json({ name: "FineTune Studio MCP", status: "ok", version: "1.0.0" });
|
||
});
|
||
|
||
// Some MCP hosts probe GET /mcp before connecting
|
||
expressApp.get("/mcp", (_req, res) => {
|
||
res.json({ name: "FineTune Studio MCP", status: "ok" });
|
||
});
|
||
|
||
expressApp.post("/mcp", async (req, res) => {
|
||
try {
|
||
const transport = new StreamableHTTPServerTransport({
|
||
sessionIdGenerator: undefined,
|
||
enableJsonResponse: true,
|
||
});
|
||
res.on("close", () => transport.close());
|
||
await server.connect(transport);
|
||
await transport.handleRequest(req, res, req.body);
|
||
} catch (err) {
|
||
console.error("MCP handler error:", err);
|
||
if (!res.headersSent) {
|
||
res.status(500).json({ error: "Internal server error" });
|
||
}
|
||
}
|
||
});
|
||
|
||
const PORT = process.env.PORT || 3002;
|
||
expressApp.listen(PORT, () =>
|
||
console.log(`FineTune Studio MCP Server running at http://localhost:${PORT}/mcp`)
|
||
);
|