Generate different representations between Python and Java code

[Hkh9966]

When I executed model.generate(**encoding, max_length=128) in my script, I found that by default only python code can be generated correctly, while Java code only has completion functions.

Gen python code:

Gen java code:

Here is my script:

import torch
from datetime import datetime
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

device = "cuda"  # for GPU usage or "cpu" for CPU usage
checkpoint = "Salesforce/codet5p-2b"

tokenizer = AutoTokenizer.from_pretrained(checkpoint)

model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint, trust_remote_code=True).to(device)


while True:
    user_input = input("Input: ")
    if user_input == "exit":
        break

    encoding = tokenizer(user_input, return_tensors="pt").to(device)
    encoding['decoder_input_ids'] = encoding['input_ids'].clone()
    outputs = model.generate(**encoding, max_length=128)
    print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Did I use it wrong? Is there any way to tell the model that generate java code by default?

[yuewang-cuhk]

Hello, may I know which model do you use? For CodeT5+ 2B/6B/16B, they are further finetuned on Python code and more suitable for Python code generation/completion.