Spaces:
Running
Running
File size: 1,419 Bytes
173f28e 7c4c461 173f28e 7c4c461 173f28e 7c4c461 173f28e f56dbf3 173f28e f56dbf3 173f28e 7c4c461 173f28e f56dbf3 173f28e 7c4c461 | 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 | from __future__ import annotations
import json
import os
import subprocess
import sys
def _measure(model_id: str, backend: str, sentences: list[str], batch_size: int) -> float:
import psutil
from models import ModelConfig
from wrapper import load_model
process = psutil.Process(os.getpid())
baseline = process.memory_info().rss
cfg = ModelConfig(name="", model_id=model_id, backend=backend)
model = load_model(cfg)
model.encode(sentences, batch_size=batch_size)
peak = process.memory_info().rss
return round((peak - baseline) / (1024 * 1024), 1)
def evaluate_memory(model_id: str, sentences: list[str], batch_size: int = 64, backend: str = "sbert") -> float:
"""Return memory delta in MB, measured in an isolated subprocess."""
payload = json.dumps({
"model_id": model_id,
"backend": backend,
"sentences": sentences,
"batch_size": batch_size,
})
result = subprocess.run(
[sys.executable, "-m", "evals.memory"],
input=payload,
capture_output=True,
text=True,
check=True,
)
return float(result.stdout.strip().splitlines()[-1])
if __name__ == "__main__":
args = json.loads(sys.stdin.read())
mb = _measure(
model_id=args["model_id"],
backend=args["backend"],
sentences=args["sentences"],
batch_size=args["batch_size"],
)
print(mb)
|