# Offline CodeT5-based Python comment-to-code/complete/FIM inference.
import argparse
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

def prompt(task, instruction=None, prefix=None, suffix=None):
    if task == 'comment_to_code':
        return 'task: comment_to_code\nlanguage: python\ninstruction:\n' + instruction
    if task == 'complete_code':
        return 'task: complete_code\nlanguage: python\nprefix:\n' + prefix
    if task == 'fill_middle':
        return 'task: fill_middle\nlanguage: python\nprefix:\n' + prefix + '\nsuffix:\n' + suffix
    raise ValueError(task)

p = argparse.ArgumentParser()
p.add_argument('--model-dir', default='.')
p.add_argument('--task', choices=['comment_to_code','complete_code','fill_middle'], default='comment_to_code')
p.add_argument('--instruction', default='Create a function that sorts integers in ascending order.')
p.add_argument('--prefix', default='def add(a, b):\n')
p.add_argument('--suffix', default='\nprint(add(2, 3))')
a = p.parse_args()
tok = AutoTokenizer.from_pretrained(a.model_dir, local_files_only=True)
model = AutoModelForSeq2SeqLM.from_pretrained(a.model_dir, local_files_only=True)
dev='cuda' if torch.cuda.is_available() else 'cpu'
model.to(dev).eval()
src = prompt(a.task,a.instruction,a.prefix,a.suffix)
enc = tok(src,return_tensors='pt',truncation=True,max_length=384).to(dev)
with torch.inference_mode():
    out=model.generate(**enc, num_beams=1, max_new_tokens=256)
print(tok.decode(out[0],skip_special_tokens=True))