INNER CODE UNIT · Python
generate_and_update
KodCode-AI/kodcode · pipeline/completion_open_model.py:172
def generate_and_update(dataset, checkpoint_file, llm=None, params=None, tokenizer=None):
processed_dataset = copy.deepcopy(dataset)
# Initialize tokenizer
if tokenizer is not None:
if tokenizer.pad_token_id is None:
tokenizer.pad_token = tokenizer.eos_token
if "gemma-2" in args.model_path.lower():
tokenizer.padding_side = "right"
# Intialize the dataset with the checkpoint file (if it exists)
if os.path.exists(checkpoint_file):
last_checkpoint_idx = len(load_dataset_from_file(checkpoint_file))
print(f"Checkpoint file found. Resuming from last checkpoint with index {last_checkpoint_idx}.")
processed_dataset[:last_checkpoint_idx] = load_dataset_from_file(checkpoint_file)
# Calculate total number of batches
num_batches = (len(processed_dataset) - last_checkpoint_idx + BATCH_SIZE - 1) // BATCH_SIZE