| |
| 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)) |
|
|
