File size: 2,124 Bytes
7993ea5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
from typing import Any
from smolagents import TransformersModel, ChatMessage
class QwenModelWithAttention(TransformersModel):
def _prepare_completion_args(
self,
messages: list[ChatMessage | dict],
stop_sequences: list[str] | None = None,
tools_to_call_from: list[Tool] | None = None,
**kwargs,
) -> dict[str, Any]:
completion_kwargs = self._prepare_completion_kwargs(
messages=messages,
stop_sequences=stop_sequences,
tools_to_call_from=tools_to_call_from,
tool_choice=None,
**kwargs,
)
messages = completion_kwargs.pop("messages")
stop_sequences = completion_kwargs.pop("stop", None)
tools = completion_kwargs.pop("tools", None)
max_new_tokens = (
kwargs.get("max_new_tokens")
or kwargs.get("max_tokens")
or self.kwargs.get("max_new_tokens")
or self.kwargs.get("max_tokens")
or 1024
)
prompt_tensor = (self.processor if hasattr(self, "processor") else self.tokenizer).apply_chat_template(
messages,
tools=tools,
return_tensors="pt",
add_generation_prompt=True,
tokenize=True,
return_dict=True,
return_attention_mask=True
)
prompt_tensor = prompt_tensor.to(self.model.device) # type: ignore
if hasattr(prompt_tensor, "input_ids"):
attention_mask = prompt_tensor["attention_mask"]
prompt_tensor = prompt_tensor["input_ids"]
model_tokenizer = self.processor.tokenizer if hasattr(self, "processor") else self.tokenizer
stopping_criteria = (
self.make_stopping_criteria(stop_sequences, tokenizer=model_tokenizer) if stop_sequences else None
)
completion_kwargs["max_new_tokens"] = max_new_tokens
return dict(
inputs=prompt_tensor,
attention_mask=attention_mask,
use_cache=True,
stopping_criteria=stopping_criteria,
**completion_kwargs,
)
|