apunkt/doc-to-lora
1
0
Fork 0
mirror of https://github.com/SakanaAI/doc-to-lora.git synced 2026-07-23 17:01:04 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Support BatchEncoding in apply_chat_template output during QA inference

Browse source
This commit is contained in:
Sarvesh 2026-06-09 16:48:47 +05:30
parent 073d909468
commit 4ac36c1210
1 changed files with 12 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

15
demo.ipynb
View file

@ -633,9 +633,18 @@
" return_tensors=\"pt\"\n",
" ).to(device)\n",
"\n",
" # Handle cases where apply_chat_template returns a BatchEncoding dict instead of a raw Tensor\n",
" if isinstance(prompt_ids, dict):\n",
" input_ids = prompt_ids[\"input_ids\"]\n",
" attention_mask = prompt_ids.get(\"attention_mask\", torch.ones_like(input_ids))\n",
" else:\n",
" input_ids = prompt_ids\n",
" attention_mask = torch.ones_like(input_ids)\n",
"\n",
"\n",
" generated_ids = model.base_model.generate(\n",
" input_ids=prompt_ids,\n",
" attention_mask=torch.ones_like(prompt_ids),\n",
" input_ids=input_ids,\n",
" attention_mask=attention_mask,\n",
" max_new_tokens=train_args.generation_max_new_tokens,\n",
" do_sample=False,\n",
" pad_token_id=tokenizer.pad_token_id,\n",
@ -643,7 +652,7 @@
" )\n",
"\n",
" video2lora_prediction = tokenizer.decode(\n",
" generated_ids[0][prompt_ids.shape[1]:],\n",
" generated_ids[0][input_ids.shape[1]:],\n",
" skip_special_tokens=True\n",
" ).strip()\n",
"\n",