github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Bug fixes and improvements to llmadapter

Browse Source
This commit is contained in:
Jaret Burkett
2025-02-15 07:18:07 -07:00
parent bd8d7dc081
commit 87e557cf1e
3 changed files with 28 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
toolkit/models/llm_adapter.py
View File

@@ -147,13 +147,14 @@ class LLMAdapter(torch.nn.Module):
prompt_attention_mask = text_inputs.attention_mask.to(device)
# remove the system prompt from the input and attention mask
text_input_ids = text_input_ids[:, self.system_prompt_length:]
prompt_attention_mask = prompt_attention_mask[:, self.system_prompt_length:]
prompt_embeds = text_encoder(
text_input_ids, attention_mask=prompt_attention_mask, output_hidden_states=True
)
prompt_embeds = prompt_embeds.hidden_states[-2]
prompt_embeds = prompt_embeds[:, self.system_prompt_length:]
prompt_attention_mask = prompt_attention_mask[:, self.system_prompt_length:]
dtype = text_encoder.dtype