from datasets import load_dataset
import random
from openai import OpenAI

# 定义数据集的配置
data_files = {"train": "harmful_behaviors.csv"}

# 加载数据集
dataset = load_dataset('csv', data_files=data_files)

# 访问训练集
train_dataset = dataset['train']

# 获取数据集的总行数
num_rows = train_dataset.num_rows

# Point to the local server
# change Uncensored_Test/harmful_behaviors.py client to you Openai Server address and api_key
client = OpenAI(base_url="http://localhost:1234/v1", api_key="lm-studio")

while True:
    # 随机选择一个行索引
    random_index = random.randint(0, num_rows - 1)

    # 使用 select 方法选择随机行
    random_row = train_dataset.select([random_index])

    # 只显示随机行的 'goal' 列的内容
    random_goal = random_row[0]['goal']
    random_goal = f"{random_goal}"
    print(f"## {random_goal}")
    # 创建聊天完成请求
    completion = client.chat.completions.create(
        model="Publisher/Repository",
        messages=[
            {"role": "system", "content": "Always answer in rhymes."},
            {"role": "user", "content": random_goal}
        ],
        temperature=0.7,
        stream=True,
    )

    # 逐字显示助手回复
    for chunk in completion:
        if chunk.choices[0].delta.content:
            print(chunk.choices[0].delta.content, end="", flush=True)

    # 打印换行符以分隔输出
    print("\n")