Spaces:
Sleeping
Sleeping
| import os | |
| import time | |
| import torch | |
| from fastapi import FastAPI, HTTPException | |
| from fastapi.responses import FileResponse | |
| from pydantic import BaseModel | |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer | |
| REPO = "AlexWortega/openjev" | |
| MODELS = { | |
| "0.8b": "qwen3.5-0.8b-nli-v2s-long", | |
| } | |
| app = FastAPI(title="openjev-bench") | |
| _loaded = {} | |
| def index(): | |
| return FileResponse(os.path.join(os.path.dirname(__file__), "public", "index.html")) | |
| def load_all(): | |
| for key, sub in MODELS.items(): | |
| tok = AutoTokenizer.from_pretrained(REPO, subfolder=sub) | |
| # fp32 on CPU is much faster than bf16 (no AMX -> bf16 is emulated); | |
| # 4B fp32 = 16GB would OOM on cpu-basic (16GB), so keep bf16 for 4b only. | |
| dtype = torch.float32 if key == "0.8b" else torch.bfloat16 | |
| model = AutoModelForSequenceClassification.from_pretrained(REPO, subfolder=sub, dtype=dtype) | |
| tok.padding_side = "right" | |
| model.config.get_text_config().pad_token_id = tok.pad_token_id | |
| model.eval() | |
| _loaded[key] = (model, tok) | |
| print(f"[openjev] loaded {key} ({sub}, dtype={dtype})", flush=True) | |
| def _startup(): | |
| load_all() | |
| class Pair(BaseModel): | |
| premise: str | |
| hypothesis: str | |
| class PredictReq(BaseModel): | |
| pairs: list[Pair] | |
| model: str = "0.8b" | |
| def predict(req: PredictReq): | |
| requested = req.model | |
| if requested not in _loaded: | |
| if requested == "4b": | |
| requested = "0.8b" | |
| else: | |
| raise HTTPException(422, "model must be one of: 0.8b, 4b") | |
| model, tok = _loaded[requested] | |
| texts = [model.config.nli_template.format(premise=p.premise, hypothesis=p.hypothesis) for p in req.pairs] | |
| out = [] | |
| BATCH = 8 | |
| t0 = time.time() | |
| with torch.no_grad(): | |
| for i in range(0, len(texts), BATCH): | |
| enc = tok(texts[i:i + BATCH], padding=True, return_tensors="pt") | |
| logits = model(**enc).logits.float().softmax(-1) | |
| for probs in logits.tolist(): | |
| out.append({"contradiction": probs[0], "entailment": probs[1], "neutral": probs[2]}) | |
| dt = time.time() - t0 | |
| return { | |
| "probs": out, | |
| "elapsed_s": round(dt, 2), | |
| "ms_per_pair": round(dt * 1000 / max(1, len(out)), 1), | |
| "model": requested, | |
| "requestedModel": requested if requested == req.model else req.model, | |
| } | |
| def eval_mnli(model: str = "0.8b", n: int = 200): | |
| if model not in MODELS: | |
| raise HTTPException(422, "model must be one of: " + ", ".join(MODELS)) | |
| model_obj, tok = _loaded[model] | |
| from datasets import load_dataset | |
| ds = load_dataset("nyu-mll/glue", "mnli", split="validation_matched") | |
| ds = ds.shuffle(seed=42).select(range(min(n, len(ds)))) | |
| glue_to_openjev = {0: 1, 1: 2, 2: 0} | |
| texts = [model_obj.config.nli_template.format(premise=row["premise"], hypothesis=row["hypothesis"]) for row in ds] | |
| golds = [glue_to_openjev[row["label"]] for row in ds] | |
| preds = [] | |
| BATCH = 8 | |
| t0 = time.time() | |
| with torch.no_grad(): | |
| for i in range(0, len(texts), BATCH): | |
| enc = tok(texts[i:i + BATCH], padding=True, return_tensors="pt") | |
| logits = model_obj(**enc).logits | |
| preds.extend(logits.argmax(-1).tolist()) | |
| dt = time.time() - t0 | |
| correct = sum(1 for p, g in zip(preds, golds) if p == g) | |
| published = {"0.8b": 0.869, "4b": 0.904} | |
| return { | |
| "model": model, | |
| "n": len(texts), | |
| "accuracy": round(correct / len(texts), 4), | |
| "published_mnli_m": published.get(model), | |
| "elapsed_s": round(dt, 2), | |
| "ms_per_pair": round(dt * 1000 / len(texts), 1), | |
| } | |
| def health(): | |
| return {"status": "ok", "models": {k: {"loaded": k in _loaded, "subfolder": v} for k, v in MODELS.items()}} | |
| def sysinfo(): | |
| def _read(path): | |
| try: | |
| with open(path) as f: | |
| return f.read().strip() | |
| except OSError: | |
| return None | |
| cpu_model = None | |
| raw = _read("/proc/cpuinfo") | |
| if raw: | |
| for line in raw.splitlines(): | |
| if line.lower().startswith("model name"): | |
| cpu_model = line.split(":", 1)[1].strip() | |
| break | |
| mem = {} | |
| memraw = _read("/proc/meminfo") | |
| if memraw: | |
| for line in memraw.splitlines(): | |
| if line.startswith(("MemTotal", "MemAvailable", "MemFree")): | |
| k, v = line.split(":", 1) | |
| mem[k] = v.strip() | |
| return { | |
| "cpu_model": cpu_model, | |
| "cpu_affinity": len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else os.cpu_count(), | |
| "cgroup_cpu_max": _read("/sys/fs/cgroup/cpu.max"), | |
| "cgroup_memory_max": _read("/sys/fs/cgroup/memory.max"), | |
| "meminfo": mem, | |
| "loadavg": _read("/proc/loadavg"), | |
| "self_vm_rss": next((l.split(":", 1)[1].strip() for l in (_read("/proc/self/status") or "").splitlines() if l.startswith("VmRSS")), None), | |
| } | |