Spaces:
Running
Running
File size: 1,728 Bytes
f7318fd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | """Local CPU benchmark with a fixed prompt and separate startup timings."""
import json
import resource
import sys
import time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from token_trail_backend.config import Settings # noqa: E402
from token_trail_backend.inference import InferenceEngine # noqa: E402
from token_trail_backend.schemas import TraceRequest # noqa: E402
def main() -> None:
engine = InferenceEngine(Settings())
started = time.perf_counter()
engine.load()
trace = engine.trace(
TraceRequest(prompt="The moon is", max_new_tokens=8, top_k=5, temperature=0)
)
steps = trace["steps"]
print(
json.dumps(
{
"model": trace["model"],
"fixed_prompt": "The moon is",
"tokenizer_load_ms": engine.tokenizer_load_ms,
"model_load_ms": engine.model_load_ms,
"prompt_tokenisation_ms": trace["timing"]["prompt_tokenisation_ms"],
"first_step_ms": trace["timing"]["first_step_ms"],
"per_step_ms": [step["step_ms"] for step in steps],
"total_trace_ms": trace["timing"]["total_trace_ms"],
"generated_token_count": len(steps),
"peak_process_memory_mb": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
/ 1024
/ 1024
if sys.platform == "darwin"
else resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024,
"whole_process_ms": (time.perf_counter() - started) * 1000,
},
indent=2,
)
)
engine.close()
if __name__ == "__main__":
main()
|