encode_instructions.py

import json
from tokenizers import Tokenizer
import numpy as np
import h5py

INSTRUCTION_JSON = "instruction-data.json"
TOKENIZER_PATH = "./tokenizer/tokenizer.json"
CONTEXT_LENGTH = 256
DATA_PATH = "./data/instructions_encoded.h5"

nprng = np.random.default_rng(1234)

def padEncodedInstructions(tokenizer, instructionData):
    finalInstructions = []
    for instruction in instructionData:
        #Okay, now, we add the bos at the start, and eos at the end
        instruction = [tokenizer.token_to_id("<bos>")] + instruction.ids
        instruction.append(tokenizer.token_to_id("<eos>"))
        paddingTokens = (CONTEXT_LENGTH - len(instruction)) + 1
        instruction += [tokenizer.token_to_id("<pad>"),] * paddingTokens
        finalInstructions.append(instruction)
    return finalInstructions

def concatInstructionAndInput(instructionData):
    instructionResponses = []
    for instructionDict in instructionData:
        totalInstruction = (
            "### INSTRUCTION:\n" +
            instructionDict["instruction"]+" "+instructionDict["input"] +
            "\n### RESPONSE:\n" + instructionDict["output"]
        )
        instructionResponses.append(totalInstruction)
    return instructionResponses

def main():
    with open(INSTRUCTION_JSON) as theFile:
        instructionData = json.load(theFile)
    # Step one, Concatenate the instructions and responses together
    totalData = concatInstructionAndInput(instructionData)
    #Now we need to tokenize everything
    tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
    tokenizedInstructions = tokenizer.encode_batch(totalData)
    #Now we need to pad all the tokenized examples
    paddedInstructions = padEncodedInstructions(tokenizer, tokenizedInstructions)
    #Okay, now we split the padded instructions up into test and train sets
    #Let's numpy-fy it, and do a random choice of train and test splits
    paddedInstructions = np.array(paddedInstructions)
    print(paddedInstructions.shape)
    train_indices = nprng.choice(np.arange(len(paddedInstructions)), size = round(0.8*len(paddedInstructions)), replace=False)
    train_examples = paddedInstructions[train_indices]
    test_indices = np.arange(len(paddedInstructions))
    test_indices = [x for x in test_indices if x not in train_indices]
    test_examples = paddedInstructions[test_indices]
    print(train_examples.shape)
    print(test_examples.shape)

    with h5py.File(DATA_PATH, "w") as theFile:
        theFile.create_dataset("train", data=train_examples)
        theFile.create_dataset("test", data=test_examples)
    print("Done!")

if __name__ == "__main__":
    main()