从 llm.c 到 llm.rs（三）

char *tiny_stories_train = "data/TinyStories_train.bin";
char *tiny_stories_val = "data/TinyStories_val.bin";
char *tiny_shakespeare_train = "data/tiny_shakespeare_train.bin";
char *tiny_shakespeare_val = "data/tiny_shakespeare_val.bin";
char *train_tokens = access(tiny_shakespeare_train, F_OK) != -1 ? tiny_shakespeare_train : tiny_stories_train;
char *val_tokens = access(tiny_shakespeare_val, F_OK) != -1 ? tiny_shakespeare_val : tiny_stories_val;

int B = 4;
int T = 64;
DataLoader train_loader;
dataloader_init(&train_loader, train_tokens, B, T);
printf("train dataset num_batches: %d\n", train_loader.num_batches);
DataLoader val_loader;
dataloader_init(&val_loader, val_tokens, B, T);
printf("val dataset num_batches: %d\n", val_loader.num_batches);
int val_num_batches = 10;

void dataloader_init(DataLoader *loader, char *filename, int B, int T)
{
    loader->B = B;
    loader->T = T;

    // open the input file for reading
    loader->tokens_file = fopen(filename, "rb");
    if (loader->tokens_file == NULL)
    {
        printf("Error opening tokens file\n");
        exit(1);
    }

    // determine the file size
    fseek(loader->tokens_file, 0, SEEK_END);
    loader->file_size = ftell(loader->tokens_file);
    fseek(loader->tokens_file, 0, SEEK_SET);
    if (loader->file_size < (B * T + 1) * sizeof(int))
    {
        printf("Error: file size is too small for the batch size and sequence length\n");
        exit(1);
    }
    loader->current_position = 0; // start at the beginning

    // allocate space for B*T + 1 integers to store the inputs and targets
    loader->batch = (int *)malloc((B * T + 1) * sizeof(int));
    loader->inputs = loader->batch;
    loader->targets = loader->batch + 1; // targets are shifted by one
    loader->num_batches = loader->file_size / (B * T * sizeof(int));
}

impl DataLoader {
    fn new(tokens_filename: &str, micro_batch_len: usize, sequence_len: usize) -> Self {
        let tokens_file_content = fs::read(tokens_filename).unwrap();
        let tokens_file_size = tokens_file_content.len();
        let batch_token_size = micro_batch_len * sequence_len * 4;
        Self {
            tokens_file_content,
            tokens_file_size,
            micro_batch_len,
            sequence_len,
            batch_content: vec![0; batch_token_size + 4],
            inputs_index: 0,
            targets_index: 4,
            current_position: 0,
            num_batches: tokens_file_size / batch_token_size,
        }
    }
}

for (int step = 0; step <= 40; step++) {
  // once in a while estimate the validation loss
  if (step % 10 == 0) {
	  ...
  }
  // once in a while do model inference to print generated text
  if (step > 0 && step % 20 == 0) {
    ...
  }
  // do a training step
  ...
}

// free
dataloader_free(&train_loader);
dataloader_free(&val_loader);
gpt2_free(&model);
return 0;