INNER CODE UNIT · Python
process_batch
KodCode-AI/kodcode · pipeline/completion_open_model.py:130
def process_batch(batch, llm, params, tokenizer=None):
user_instructions = [item['messages'][0]['content'] for item in batch]
prompts = []
for instruction in user_instructions:
chat = [{"role": "user", "content": instruction}]
template = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
prompts.append(template)
if args.engine == "vllm":
outputs = llm.generate(prompts, params)
elif args.engine == "hf":
inputs = tokenizer(prompts, return_tensors="pt", padding=True, truncation=True).to(torch.cuda.current_device())
gen_do_sample = False if args.temperature == 0 else True
outputs = llm.generate(**inputs,
tokenizer=tokenizer,
do_sample=gen_do_sample,
temperature=args.temperature if gen_do_sample else None, # To avoid temperature` (=0) has to be a strictly positive float
top_p=args.top_p,
repetition_penalty=args.repetition_penalty,