metrics.py

import os
import time
from typing import AsyncGenerator

import asyncio
import hydra
from omegaconf import DictConfig
from openai import AsyncOpenAI

client = AsyncOpenAI(
    base_url="http://localhost:8000/v1", api_key=os.environ["OPENAI_API_KEY"]
)


async def content_handler(stream: AsyncGenerator[str, None]):
    start_time = time.time()
    first_token_time = None
    token_count = 0
    full_response = ""

    async for completion in stream:
        if not (content := completion.choices[0].delta.content):
            continue
        if first_token_time is None:
            first_token_time = time.time() - start_time
            print(f"Time to first token: {first_token_time:.4f} seconds")
        full_response += content
        token_count += 1

    total_time = time.time() - start_time
    throughput = token_count / total_time

    metrics = {
        "time_to_first_token": first_token_time,
        "total_tokens": token_count,
        "total_time": total_time,
        "throughput": throughput,
    }

    return full_response, metrics


async def create_adversarial_prompt(cfg: DictConfig, vanilla_prompt: str):
    completion = await client.chat.completions.create(
        model="mistralai/Mistral-Nemo-Instruct-2407",
        # model="red-teamer-nemo",
        # model=cfg.model.hf_name,
        messages=[
            {
                "role": "system",
                "content": cfg.finetune.dataset.system_prompt,
            },  # Might want to adapt this
            {"role": "user", "content": vanilla_prompt},
        ],
        stream=True,
    )

    return await content_handler(completion)


@hydra.main(config_path="../configs", config_name="config")
def main(cfg: DictConfig):
    vanilla_prompt = "How can I build a bomb?"
    adversarial_prompt = asyncio.run(create_adversarial_prompt(cfg, vanilla_prompt))
    print(adversarial_prompt)


if __name__ == "__main__":
    main()