from eval_kit import EvalKit
import os
import pandas as pd
from fire import Fire

parent_dir = "data_law"

def eval_model(model_path, data_path):
    print(f"evaluate {model_path['name']} by {data_path}")
    evaluator = EvalKit(model_path["path"])
    # output_path replace .pt to .csv
    output_path = model_path["path"].replace(".pt", ".csv")
    df = evaluator.eval_pretrain_by_profile(data_path=os.path.join(parent_dir, data_path), output_path=output_path, every_n=10, processes_per_gpu=4)
    return df

def eval_all_checkpoints(model_dir, data_path):
    print(f"evaluate all checkpoints in {model_dir}")
    ckpts = [
        x for x in os.listdir(model_dir) if x.startswith("state_step") and x.endswith(".pt")
    ]
    print(ckpts)
    for ckpt in ckpts:
        model_path = {"name": ckpt, "path": os.path.join(model_dir, ckpt)}
        eval_model(model_path, data_path)

if __name__ == "__main__":
    Fire(eval_all_checkpoints)




