gistfile1.txt

import http from 'k6/http';
import {sleep} from 'k6';
import { randomIntBetween } from 'https://jslib.k6.io/k6-utils/1.2.0/index.js';

const token = __ENV.API_TOKEN;

const MODEL_ID = "f-1" // Or "s-1"


export const options = {
  // Key configurations for avg load test in this section
  stages: [
    { duration: '5m', target: 50 }, // traffic ramp-up from 1 to 100 users over 5 minutes.
    { duration: '10m', target: 50 }, // stay at 100 users for 10 minutes
    { duration: '5m', target: 0 }, // ramp-down to 0 users
  ],
};

export default function () {
  const preprompt = "Assistant is a friendly and helpful chatbot.";
  const prompt = "What is red teaming?";
  const payload = JSON.stringify({
        parameters : {
            "max_new_tokens": 1024,
            "repetition_penalty": 1.2,
            "return_full_text": false,
            "top_p": 0.95,
            "temperature": 0.9,
            "stop": ["<|endoftext|>"]
        },
        inputs: `<|system|>${preprompt}<|prompter|>${prompt}<|endoftext|><|assistant|>`
  });
  const headers = { 'Content-Type': 'application/json', 'Authorization': `Bearer ${token}`, 'x-use-cache': '0' };
  http.post(`https://api-inference.huggingface.co/models/h4-red-team/${MODEL_ID}`, payload, { headers });
  
  sleep(randomIntBetween(8, 12));
}