import matplotlib.pyplot as plt

# 数据
x_labels = ['100', '250', '500', '1000', '2000', '4000', '6000']
models = {
    'DS_QW_7B-prompt': [33.8, 43.6, 55.2, 72, 84.4, 89.6, 90],
    # 'DS_QW_7B-promptv2': [34.4, 42, 55.2, 73, 86.2, 89.2, 90.8],
    # 'SFT(with special token)': [40, 45.2, 66.4, 73, 84, 89.2, 92.8],
    'SFT+RL(with special token)': [38.4, 46.4, 68.2, 74.8, 86, 91.4, 92.2]
}

# 创建图形和轴
fig, ax = plt.subplots()

# 绘制每条折线
colors = ['blue', 'green', 'red', 'purple']
for (model, scores), color in zip(models.items(), colors):
    ax.plot(x_labels, scores, label=model, color=color, marker='o')

# 设置标题和标签
ax.set_title('MATH500 Scores by Model and Sample Size')
ax.set_xlabel('Sample Size')
ax.set_ylabel('Score')

# 显示图例
ax.legend()

# 保存图形
plt.savefig('math500_scores.png', bbox_inches='tight')

# 关闭图形，避免显示
plt.close()