I have a decent amount of experience with Tensorflow and I thought I’d teach myself how to design an LLM using this tutorial of theirs. https://www.tensorflow.org/text/tutorials/transformer

Most of the project makes sense to me, but I don’t think I fully understand what shape the inputs and outputs are supposed to take. They provide an infographic, but I’m still not very sure. I have inputs and outputs that look like a set of prompts which each have their own set of potential answers.

Infographic stating inputs and outputs to transformer model

Right now because each prompt has several possible answers, I expand the inputs to have the same prompt matching every instance of a possible answer.

prompts = ["hey man what's up", "hey dude"]
responses = [["nothing much", "the sky"], ["hey man", "how's it going?"]]

x_train = []
y_train = []

for i in range(len(prompts)):
    x_train.extend([prompts[i]] * len(responses[i]))
    y_train.extend(responses[i])

# Sets now look like this, although real ones are tokenized and padded at the end
x_train = ["hey man what's up", "hey man what's up", "hey dude", "hey dude"]
y_train = ["nothing much",      "the sky",           "hey man",   "how's it going?"]

The tutorial shows that I need to feed the data to the model as something like (prompt, response), next_response where next_response is the response string minus the first token and add the next token in the sequence, but specifically it looks like it only gets rid of the first token and adds the end token.

Infographic depicting both possible approaches to the question

Right now I am doing the thing on the left and I’m getting terrible results during training. Should I try to implement the solution on the right? And is my way of padding the data to allow for multiple responses per prompt actually poisoning the model?

Any help is appreciated.