"""Run FleetWorker-E4B on one task-board item. Use the project field; ignore priority (see the model card).

    pip install "transformers>=5" torch
    python inference_example.py "title: Add retry with backoff to the nightly export job"
"""
import json
import sys

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

REPO = "ThakiCloud/FleetWorker-E4B"
LABELS = ["ai-experiment-measurement", "gitops-deployment-health", "synthetic-data-pipeline", "company-operations-agents", "ondevice-small-model", "todo-board-tooling", "enterprise-ops-benchmark", "shop-assistant-console", "commerce-growth-agents", "platform-strategy-planning", "decision-routing-research", "legal-decision-models", "secops-automation-console", "multimodal-api-gateway", "character-ai-app", "telecom-phy-research", "engineering-blog-publishing", "generative-video-training", "semiconductor-design-llm", "agent-trajectory-distillation", "korean-humanization-corpus", "edge-inference-box", "unattended-agent-runs", "personal-agent-system", "music-generation-rooms"]
SYSTEM = "You triage a task-board item. Output exactly one compact JSON object with keys \"project\" and \"priority\" and nothing else. \"project\" is one of: ai-experiment-measurement, gitops-deployment-health, synthetic-data-pipeline, company-operations-agents, ondevice-small-model, todo-board-tooling, enterprise-ops-benchmark, shop-assistant-console, commerce-growth-agents, platform-strategy-planning, decision-routing-research, legal-decision-models, secops-automation-console, multimodal-api-gateway, character-ai-app, telecom-phy-research, engineering-blog-publishing, generative-video-training, semiconductor-design-llm, agent-trajectory-distillation, korean-humanization-corpus, edge-inference-box, unattended-agent-runs, personal-agent-system, music-generation-rooms, or \"other\"; \"unknown\" if the text gives no way to tell. \"priority\" is one of: P0, P1, P2, P3."

tok = AutoTokenizer.from_pretrained(REPO)
model = AutoModelForCausalLM.from_pretrained(REPO, dtype=torch.bfloat16, device_map="auto").eval()
stop = [i for i in {tok.eos_token_id, tok.convert_tokens_to_ids("<turn|>")} if isinstance(i, int)]
item = sys.argv[1] if len(sys.argv) > 1 else "title: Add retry with backoff to the nightly export job"
messages = [{"role": "system", "content": SYSTEM}, {"role": "user", "content": item}]
prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
ids = tok(prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
with torch.no_grad():
    out = model.generate(**ids, max_new_tokens=64, do_sample=False, eos_token_id=stop)
text = tok.decode(out[0, ids["input_ids"].shape[1]:], skip_special_tokens=False).split("<turn|>")[0].strip()
pred = json.loads(text)
print({"project": pred.get("project"), "valid_label": pred.get("project") in LABELS + ["other", "unknown"]})