Spaces:
Running
Running
| import sys | |
| import os | |
| import random | |
| from time import sleep | |
| from gradio_client import Client | |
| current_dir = os.path.dirname(os.path.abspath(__file__)) | |
| sys.path.append(os.path.dirname(current_dir)) | |
| from datasets import load_dataset | |
| def send_message(client, message): | |
| _apps = ["fintechgpt", "healthgpt", "mydoc", "knowledge-centre", "assistantgpt"] | |
| # Compute exponential weights for service names | |
| factor = 2 # Control the steepness of the exponential decrease | |
| app_weights = [factor ** -i for i in range(len(_apps))] | |
| # Select a service_name randomly based on the exponential weights | |
| service_name = random.choices(_apps, weights=app_weights, k=1)[0] | |
| # Define providers and corresponding models with weights | |
| providers_models = { | |
| "Meta AI": (["LLaMA 65B", "LLaMA 33B"], 10), | |
| "Mistral": (["Mistral 7B"], 8), | |
| "Anthropic": (["Claude 3.5 Sonnet"], 15), | |
| "OpenAI": (["GPT-3", "GPT-3.5", "GPT-4", "GPT-4o"], 25), | |
| "Google": (["Gemini Ultra", "Gemini Pro", "Gemini Nano", "Lamda", "Palm"], 20), | |
| "Databricks": (["Dolly"], 5), | |
| "IBM AI": (["Watson NLP"], 5), | |
| "Azure AI": (["Azure OpenAI", "Custom GPT-3.5"], 7), | |
| "Snowflake": (["Snowflake GPT"], 3), | |
| "Krutrim": (["Krutrim LLM"], 2), | |
| "Baidu": (["Ernie 4.0"], 10), | |
| "Stability AI": (["StableLM 3B", "StableLM 7B"], 4), | |
| "Microsoft": (["Orca", "Phi-1"], 8) | |
| } | |
| # Extract providers and their weights | |
| providers = list(providers_models.keys()) | |
| provider_weights = [providers_models[provider][1] for provider in providers] | |
| # Randomly select a provider based on weights | |
| selected_provider = random.choices(providers, weights=provider_weights, k=1)[0] | |
| # Randomly select a model from the selected provider | |
| selected_model = random.choice(providers_models[selected_provider][0]) | |
| result = client.predict( | |
| message, # str in 'Message' Textbox component | |
| api_name="/chat" | |
| ) | |
| print(f"Sending message: {message} \n\t Result {result}") | |
| def main(): | |
| prompt_injection_dataset = load_dataset("deepset/prompt-injections") | |
| toxicity_prompts = load_dataset("allenai/real-toxicity-prompts") | |
| nats_url = os.environ.get('CHAT_APP_URL', "http://localhost:7860/") | |
| client = Client(nats_url) | |
| while True: | |
| # Collect messages | |
| prompt_injection_sample = prompt_injection_dataset["train"].shuffle(seed=42).select(range(10)) | |
| toxicity_sample = toxicity_prompts["train"].shuffle(seed=42).select(range(10)) | |
| plain_messages = ["this is test conversation" for _ in range(10)] | |
| # Combine all messages into a single list | |
| all_messages = [msg["text"] for msg in prompt_injection_sample] + \ | |
| [msg["prompt"]["text"] for msg in toxicity_sample] + \ | |
| plain_messages | |
| # Shuffle the combined list to mix message types | |
| random.shuffle(all_messages) | |
| # Send each message | |
| for message in all_messages: | |
| send_message(client, message) | |
| sleep(random.uniform(0.5, 2)) # Random sleep between 0.5 to 2 seconds | |
| if __name__ == "__main__": | |
| main() | |