Spaces:
Sleeping
Sleeping
File size: 1,874 Bytes
4b3a33f | 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 52 53 54 55 56 | import time
import torch
import numpy as np
from sentence_transformers import SentenceTransformer
import psutil
import os
def benchmark_bge():
print("🚀 Starting BGE-M3 Efficiency Benchmark...")
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"💻 Device: {device}")
print("📥 Loading BAAI/bge-m3...")
start_load = time.time()
model = SentenceTransformer('BAAI/bge-m3', device=device)
print(f"⏱️ Load Time: {time.time() - start_load:.2f}s")
process = psutil.Process(os.getpid())
mem_info = process.memory_info()
print(f"📊 Memory Usage (RAM): {mem_info.rss / 1024 / 1024:.2f} MB")
sentences = [
"The quick brown fox jumps over the lazy dog.",
"Artificial intelligence is transforming the recruitment industry.",
"Candidate has 5 years of experience in Python and FastAPI.",
"Looking for a Senior Software Engineer with cloud expertise."
] * 25 # 100 sentences
batch_sizes = [1, 4, 8, 16, 32]
print("\n--- Latency vs Batch Size ---")
print(f"{'Batch Size':<12} | {'Time (s)':<10} | {'Sec/Sent':<10} | {'Throughput (sent/s)':<20}")
print("-" * 65)
for bs in batch_sizes:
start_time = time.time()
# Warmup
model.encode(sentences[:bs], batch_size=bs, show_progress_bar=False)
# Actual benchmark
start_time = time.time()
model.encode(sentences, batch_size=bs, show_progress_bar=False)
end_time = time.time()
total_time = end_time - start_time
sec_per_sent = total_time / len(sentences)
throughput = len(sentences) / total_time
print(f"{bs:<12} | {total_time:<10.3f} | {sec_per_sent:<10.4f} | {throughput:<20.2f}")
print("\n✅ Benchmark Complete.")
if __name__ == "__main__":
benchmark_bge()
|