rustgen / app.py
Rishabh Mishra
Add ZeroGPU support (spaces + @spaces.GPU on translate)
16e09d1
Raw
History Blame Contribute Delete
1.48 kB
"""Hugging Face Space entry point for the RustGen demo (ZeroGPU).
Runs the real Qwen2.5-Coder-1.5B backend with the compile-guided cascade and
TF-IDF retrieval; rustc (installed via packages.txt) powers the verification
panel. The rustgen package and data/rust_corpus_qwen.jsonl sit next to this file
(see prepare.sh), so no install step is needed.
ZeroGPU: the model generation runs inside `translate()`, which we wrap with
`@spaces.GPU` at startup. That both allocates a GPU slice for the duration of a
request and lets the ZeroGPU runtime detect a GPU function at boot. The wrap is a
no-op if `spaces` is unavailable (a plain CPU Space, or running locally).
"""
import os
os.environ.setdefault("RUSTGEN_BACKEND", "hf")
os.environ.setdefault("RUSTGEN_RAG_BACKEND", "tfidf")
from rustgen.app import demo as demo_module
try:
import spaces
# Wrap the module-level generation handler so it executes on an allocated
# GPU. build_demo() looks up `translate` by name at wiring time, so replacing
# it here (before build_demo runs) is enough — no change to rustgen itself.
demo_module.translate = spaces.GPU(duration=120)(demo_module.translate)
print("ZeroGPU: translate() wrapped with @spaces.GPU (duration=120s)")
except Exception as exc: # not a ZeroGPU Space / spaces not installed
print(f"ZeroGPU wrapper not applied ({type(exc).__name__}: {exc}) — CPU path")
demo = demo_module.build_demo()
if __name__ == "__main__":
demo.launch()