File size: 1,520 Bytes
403b890 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 | # 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))
|