3 ms·
I wrote the program between <BEF> and <MSG>, and it generated the following output: <NME> diff_model.py <BEF> import argparse import torch import tran
by mortehu 4y ago
I wrote the program between <BEF> and <MSG>, and it generated the following output:
<NME> diff_model.py
<BEF> import argparse
import torch
import transformers
def main():
argparser = argparse.ArgumentParser()
argparser.add_argument('--checkpoint', default='CarperAI/diff-codegen-2b-v2', choices=['CarperAI/diff-codegen-6b-v2', 'CarperAI/diff-codegen-2b-v2', 'CarperAI/diff-codegen-350m-v2'], help='Model to use')
args = argparser.parse_args()
model = transformers.AutoModelForCausalLM.from_pretrained(args.checkpoint)
tokenizer = transformers.AutoTokenizer.from_pretrained(args.checkpoint)
# Use CUDA.
model = model.cuda()
text = '<NME> diff_model.py\n<BEF> '
text += open("diff_model.py").read()
text += '<MSG> Print a message after loading the model\n<DFF>\n'
completion = model.generate(tokenizer.encode(text, return_tensors="pt").cuda(), max_new_tokens=400)
completion = tokenizer.decode(completion[0], skip_special_tokens=True)
print(completion)
if __name__ == '__main__':
main()
<MSG> Print a message after loading the model
<DFF>
<DFF> @@ -1,4 +1,5 @@
'''
+import sys
import argparse
import torch
@@ -8,6 +9,7 @@ def main():
tokenizer = transformers.AutoTokenizer.from_pretrained(args.checkpoint)
# Use CUDA.
+ print("CUDA:", torch.cuda.is_available())
model = model.cuda()
text = '<NME> diff_model.py\n<BEF> '
@@ -10,3 +12,4 @@ def main():
completion = tokenizer.decode(completion[0], skip_special_tokens=True)
print(completion)
+ sys.exit(0)
\ No newline at end of file
[\ No newline at end of file repeated many times]
It takes about 150 seconds to run on a 3090 Ti when the model is already on disk.