import wandb
import matplotlib.pyplot as plt

# 初始化 wandb API
api = wandb.Api()

# 假设您想要访问的项目名为 'my_project'，并且您的 wandb 用户名为 'my_username'
project_name = "aicolab/RWKV-5-Test"

# 获取项目中的runs
runs = api.runs(project_name)

# 假设您只关心第一个run
run = runs[0]

# 获取特定run的所有数据（例如，损失和准确度）
# 这里我们获取 'loss' 和 'accuracy' 作为示例
history = run.history()

# 绘制图表
plt.figure(figsize=(10, 4))

# 绘制损失图
plt.subplot(1, 2, 1)
plt.plot(history["loss"])
plt.title("Loss over Time")
plt.xlabel("Epoch")
plt.ylabel("Loss")

# 绘制准确度图
plt.subplot(1, 2, 2)
plt.plot(history["lr"])
plt.title("lr over Time")
plt.xlabel("Epoch")
plt.ylabel("lr")

plt.tight_layout()
plt.show()