Spaces:
Running
Running
File size: 4,397 Bytes
7191aa9 | 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 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | """
InferRoute Streaming Benchmark Loader (Hugging Face streaming=True).
Supports zero-disk-download streaming evaluation from:
1. allenai/WildChat-4.8M (3.2M real ChatGPT conversations up to Aug 2025, ideal for production traffic simulation)
2. HuggingFaceH4/no_robots (10,000 category-labeled instructions, ideal for task-aware routing correctness)
3. OpenAssistant/oasst1 (Multi-turn conversation trees)
4. lmsys/lmsys-chat-1m (Gated multi-model arena dataset)
"""
import os
import sys
import json
import time
import asyncio
from typing import Generator, Dict, Any, Optional
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# Lazy import datasets
try:
from datasets import load_dataset
HAS_HF_DATASETS = True
except ImportError:
HAS_HF_DATASETS = False
def stream_wildchat_prompts(limit: int = 1000) -> Generator[Dict[str, Any], None, None]:
"""
Streams real user prompts from allenai/WildChat-4.8M with streaming=True.
Zero 15GB repo download required.
"""
if not HAS_HF_DATASETS:
print("[WARN] datasets library not installed. Cannot stream WildChat-4.8M.")
return
print(f"[STREAM] Connecting to allenai/WildChat-4.8M (streaming=True)...")
try:
ds = load_dataset("allenai/WildChat-4.8M", split="train", streaming=True)
count = 0
for row in ds:
conversation = row.get("conversation", [])
user_messages = [
msg.get("content", "")
for msg in conversation
if isinstance(msg, dict) and msg.get("role") == "user" and msg.get("content")
]
if user_messages:
prompt = user_messages[0]
count += 1
yield {
"id": f"wildchat_{count:06d}",
"source": "allenai/WildChat-4.8M",
"category": "wildchat_real_user",
"prompt": prompt.strip(),
"requires_json": ("json" in prompt.lower() or "schema" in prompt.lower())
}
if count >= limit:
break
print(f"[OK] Streamed {count:,} real prompts from allenai/WildChat-4.8M")
except Exception as e:
print(f"[ERR] WildChat streaming exception: {e}")
def stream_no_robots_prompts(limit: int = 1000) -> Generator[Dict[str, Any], None, None]:
"""
Streams category-labeled instructions from HuggingFaceH4/no_robots with streaming=True.
Used for task-aware routing correctness benchmarking.
"""
if not HAS_HF_DATASETS:
print("[WARN] datasets library not installed. Cannot stream no_robots.")
return
print(f"[STREAM] Connecting to HuggingFaceH4/no_robots (streaming=True)...")
try:
ds = load_dataset("HuggingFaceH4/no_robots", split="train", streaming=True)
count = 0
for row in ds:
prompt = row.get("prompt", "")
category = row.get("category", "general")
if prompt:
count += 1
yield {
"id": f"norobots_{count:06d}",
"source": "HuggingFaceH4/no_robots",
"category": category,
"prompt": prompt.strip(),
"requires_json": category in ["coding", "summarization", "classification"]
}
if count >= limit:
break
print(f"[OK] Streamed {count:,} labeled prompts from HuggingFaceH4/no_robots")
except Exception as e:
print(f"[ERR] NoRobots streaming exception: {e}")
def run_streaming_benchmark(limit: int = 500, dataset_choice: str = "wildchat"):
"""
Executes a real-time streaming benchmark directly feeding prompts into InferRoute.
"""
if dataset_choice == "wildchat":
generator = stream_wildchat_prompts(limit=limit)
else:
generator = stream_no_robots_prompts(limit=limit)
prompts = list(generator)
print(f"[BENCHMARK] Loaded {len(prompts)} prompts from HF stream. Ready for gateway evaluation.")
return prompts
if __name__ == "__main__":
prompts = run_streaming_benchmark(limit=10, dataset_choice="norobots")
for p in prompts[:3]:
print("--- Sample Prompt ---")
print(f"ID: {p['id']} | Category: {p['category']}")
print(f"Content: {p['prompt'][:100]}...\n")
|