49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
|
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")
|