Views
No views yet
1# pip install -q transformers
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4checkpoint = "bigcode/tiny_starcoder_py"
5device = "cuda" # for GPU usage or "cpu" for CPU usage
6
7tokenizer = AutoTokenizer.from_pretrained(checkpoint)
8model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
9
10inputs = tokenizer.encode("def print_hello_world():", return_tensors="pt").to(device)
11outputs = model.generate(inputs)
12print(tokenizer.decode(outputs[0]))1input_text = "<fim_prefix>def print_one_two_three():\n print('one')\n <fim_suffix>\n print('three')<fim_middle>"
2inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
3outputs = model.generate(inputs)
4print(tokenizer.decode(outputs[0]))