can not generate with mode: Fill-in-the-middle
my code as below:
pip install -q transformers
from transformers import AutoModelForCausalLM, AutoTokenizer
import os
checkpoint = "bigcode/starcoder"
device = "cuda" # for GPU usage or "cpu" for CPU usage
tokenizer = AutoTokenizer.from_pretrained(checkpoint,use_auth_token=True)
model = AutoModelForCausalLM.from_pretrained(checkpoint, trust_remote_code=True,load_in_8bit=True,device_map={"": 0})
input_text = "def print_hello_world():\n \n print('Hello world!')"
inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
outputs = model.generate(inputs)
print(tokenizer.decode(outputs[0]))
output:
Does anyone know what is the reason for this?
I run into the same issues and have not been able to resolve it.
In your case, increasing the length of the generated tokens may help.
You can run FIM using the following code:
from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("bigcode/starcoder", truncation_side="left")
model = AutoModelForCausalLM.from_pretrained("bigcode/starcoder", torch_dtype=torch.bfloat16).cuda()
input_text = "<fim_prefix>def fib(n):<fim_suffix> else:\n return fib(n - 2) + fib(n - 1)<fim_middle>"
inputs = tokenizer(input_text, return_tensors="pt").to("cuda")
outputs = model.generate(**inputs, max_new_tokens=25)
generation = [tokenizer_fim.decode(tensor, skip_special_tokens=False) for tensor in outputs]
print(generation[0])
<fim_prefix>def fib(n):<fim_suffix> else:
return fib(n - 2) + fib(n - 1)<fim_middle>
if n < 2:
return n
<|endoftext|>