Generate token by token using forward function
Dear community,
I am trying implement a code where I can generate token by token and feed it into the model to generate the next token. Since this is a many-to-many model, I am struggling how to provide the target language to the decoder. I am providing the code i used so far,
tokenizer.src_lang = "en"
input_text = "I am going home"
target_language = "ta"
decoder_input_ids = [model.config.decoder_start_token_id]
predicted_ids = []
for i in range(20):
outputs = model(input_ids=input_ids, decoder_input_ids=torch.tensor([decoder_input_ids]))
logits = outputs.logits[:,i,:]
# perform argmax on the last dimension (i.e. greedy decoding)
predicted_id = logits.argmax(-1)
predicted_ids.append(predicted_id.item())
print(tokenizer.decode([predicted_id.squeeze()]))
# add predicted id to decoder_input_ids
decoder_input_ids = decoder_input_ids + [predicted_id]
if predicted_id.item() == tokenizer.eos_token_id:
break
# Decode the generated output
generated_text = tokenizer.decode(predicted_ids, skip_special_tokens=False)
print("Generated text:", generated_text)
What this code does is, it provides a decoder_start_token_id as the first decoder_input _token, so this gives perfect translation for a random language. I have no control over which language i want it to generate. The model.generate has parameter called forced_bos_token_id where we provide the lang_id of our target language, hence I know it has to be something similar to it, but i tried adding the lang_id into the decoder_input_ids but it fails to provide the intended output.
Your help is much appreciated!