| |
| |
| |
| import os |
| import sys |
| from pathlib import Path |
|
|
| import gradio as gr |
| import spaces |
| import torch |
| from fastapi import HTTPException |
| from fastapi.responses import FileResponse, HTMLResponse |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer |
|
|
| ROOT = Path(__file__).resolve().parent |
| MODEL = sys.argv[1] if len(sys.argv) > 1 else os.environ.get("PANCODE_MODEL", "tiagozip/pancode") |
| LOCAL = Path(MODEL).resolve() if Path(MODEL).is_dir() else None |
| CHUNK_CHARS = 6000 |
| MAX_CHARS = 40000 |
| MAX_CHUNKS = 12 |
|
|
| dev = "cuda" if os.environ.get("SPACES_ZERO_GPU") or torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" |
| tok = AutoTokenizer.from_pretrained(MODEL) |
| model = AutoModelForSequenceClassification.from_pretrained(MODEL, dtype=torch.float32).to(dev).eval() |
| CONFIG = { |
| "model": "local" if LOCAL else MODEL, |
| "local": bool(LOCAL), |
| "heads": model.config.pancode_heads, |
| "thresholds": model.config.pancode_thresholds, |
| "eval": model.config.pancode_eval, |
| } |
| PAGE = (ROOT / "index.html").read_text() |
| ASSETS = {p.name for p in (ROOT / "assets").glob("*.woff2")} |
|
|
| app = gr.Server() |
|
|
|
|
| @app.api(name="score") |
| @spaces.GPU(duration=20) |
| def score(chunks: list[str]) -> list[list[float]]: |
| if not 0 < len(chunks) <= MAX_CHUNKS or max(map(len, chunks)) > CHUNK_CHARS or sum(map(len, chunks)) > MAX_CHARS: |
| raise gr.Error(f"Send at most {MAX_CHARS:,} characters, in chunks of up to {CHUNK_CHARS:,}.") |
| enc = tok(chunks, truncation=True, max_length=CONFIG["thresholds"]["max_len"], padding=True, return_tensors="pt").to(dev) |
| with torch.no_grad(): |
| return model(**enc).logits.float().cpu().tolist() |
|
|
|
|
| @app.get("/", response_class=HTMLResponse) |
| def page(): |
| return PAGE |
|
|
|
|
| @app.get("/config.json") |
| def config(): |
| return CONFIG |
|
|
|
|
| @app.get("/assets/{name}") |
| def asset(name: str): |
| if name not in ASSETS: |
| raise HTTPException(404) |
| return FileResponse(ROOT / "assets" / name) |
|
|
|
|
| @app.get("/models/local/{path:path}") |
| def model_file(path: str): |
| f = LOCAL and (LOCAL / path).resolve() |
| if not f or not f.is_relative_to(LOCAL) or not f.is_file(): |
| raise HTTPException(404) |
| return FileResponse(f) |
|
|
|
|
| app.launch() |
|
|
