import pickle

import os
from tqdm import tqdm, trange

# Get the path to the generations folder
generations_folder = "generations_llama"

# Get a list of all .pkl files in the generations folder
FILENAMES = [os.path.join(generations_folder, f) for f in os.listdir(generations_folder) if f.endswith(".pkl")]
print('file names:', FILENAMES)
combined = []
# open every file
for filename in FILENAMES:
    with open(filename, "rb") as f:
        data = pickle.load(f)
        combined.extend(data)

print(f"Combined {len(combined)} generations")
print(f"combined[0] length: {len(combined[0])}")
with open("combined_generations.pkl", "wb") as f:
    pickle.dump(combined, f)
