Update modeling_rwkv5.py

modeling_rwkv5.py

@@ -789,6 +789,9 @@ class Rwkv5ForCausalLM(Rwkv5PreTrainedModel):
  # only last token for inputs_ids if the state is passed along.
  if state is not None:
  input_ids = input_ids[:, -1].unsqueeze(-1)
+ else:
+ # add in \n at the beginning
+ input_ids = torch.cat([torch.full([1,1],11,device=input_ids.device,dtype=input_ids.dtype), input_ids])
 
  # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
  if inputs_embeds is not None and state is None: