|
|
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) |
|
|
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, |
|
|
) |
|
|
|
