from locust import HttpUser, task, between
from datasets import load_dataset
import random 
datCW = load_dataset("surrey_nlp/PLOD_CW")


def pick_random_payload():
    global datCW
    idx = random.randint(0, len(datCW["test"])-1)
    tokens = datCW["test"][idx]["tokens"]
    payload = ' '.join(tokens)
    print(payload)
    return payload

class ApiUser(HttpUser):
    wait_time = between(1, 5)

    @task
    def ner_endpoint(self):
        payload = pick_random_payload()
        payload = {
            "text": payload
        }
        self.client.post("/ner", json=payload)

if __name__ == "__main__":
    import os
    os.system("locust -f locustfile.py")