microsoft / CodeBERT

CodeBERT
MIT License
2.19k stars 450 forks source link

Code completion with >=2 masks #305

Open Requalvia opened 9 months ago

Requalvia commented 9 months ago

Hi. I am using code completion task in UniXCoder, and I am wondering if it is possible to contain >=2 masks in one code. For example: `

context = """

  def write_json(data,file_path):
      data = json.<mask0>(data)
      with open(file_path, 'w') as f:
          <mask1>.write(data)
"""
tokens_ids = model.tokenize([context],max_length=512,mode="<encoder-decoder>")
source_ids = torch.tensor(tokens_ids).to(device)
prediction_ids = model.generate(source_ids, decoder_only=False, beam_size=3, max_length=128)
predictions = model.decode(prediction_ids)
print([x.replace("<mask0>","").strip() for x in predictions[0]])

` But len(prediction) is 1, so is there a method to get the prediction of \<mask1> ?

Thanks.