import requests
import pandas as pd
# 构造要发送的 JSON 数据
results = []
out_path = "/home/data/result/test.parquet"
df = pd.read_parquet("/home/data/raw/test/4201_2355_full_label_1000-8192.parquet")
# 向服务器发送请求
url = "http://localhost:5000/get_reward"
total=0
correct = 0

for idx, row in df.iterrows():
    # 拼接 query
    q1 = str(row["chosen_prompt"]) + str(row["chosen"])
    q2 = str(row["chosen_prompt"]) + str(row["reject"])
    print(q1)
    print("\n")
    print(q2)
    payload = {"query": [q1, q2]}

    try:
    # print(q1)
        response = requests.post(url, json=payload)
        scores= response.json().get("rewards", [])
        score1, score2 = scores[0], scores[1]
        right = 1 if score1 > score2 else 0
        total += 1
        if score1 > score2:
            correct += 1
        scores= response.json().get("rewards", [])
        acc = correct / total * 100
        print(f"Row {idx}: score1={score1:.4f}, score2={score2:.4f}, "
              f"Correct={score1 > score2}, RunningAcc={acc:.2f}%")
        results.append({
                "q1": q1,
                "q2": q2,
                "chosen_score": score1,
                "reject_score": score2,
                "right": right,
                "chosen_label": row["chosen_label"],
                "chosen_violations": row["chosen_violations"],
                "reject_label": row["reject_label"],
                "reject_violations": row["reject_violations"]
            })

    except Exception as e:
        print(f"Row {idx} 出错:", e)
    if total >= 2:
        break
results_df = pd.DataFrame(results)
results_df.to_parquet(out_path)