diff --git a/src/transformers/decision/modeling_utils.py b/src/transformers/decision/modeling_utils.py index eaa313d530a2..543e6a0054f2 100644 --- a/src/transformers/decision/modeling_utils.py +++ b/src/transformers/decision/modeling_utils.py @@ -19,6 +19,7 @@ from pathlib import Path import torch +from torch.nn.utils.rnn import pad_sequence from ..configuration_utils import PreTrainedConfig from ..utils import cached_file @@ -26,10 +27,41 @@ from .template_utils import Question -def process_scores(decision_config, input_ids, marks, questions, scores) -> list[torch.Tensor]: - """The logits processors of the decision config on the option scores of each question of a prompt.""" - tail = input_ids[0, int(marks[-1]) :] # the tokens after the last mark, where the templates write the state - return [decision_config.logits_processors[q.type](tail, s.float()) for q, s in zip(questions, scores)] +def batch_prompts(prompts: list[tuple[dict, list[Question]]]) -> dict[str, torch.Tensor]: + """ + The prompts of a request as one right-padded batch, with `logits_to_keep` the mask of the positions read: the + token before each mark, so the read positions of a batch are its questions in order. + """ + inputs = [prompt for prompt, _ in prompts] + reads = [] + for prompt in inputs: + read = torch.zeros(prompt["input_ids"].shape[-1], dtype=torch.bool, device=prompt["input_ids"].device) + read[prompt["marks"] - 1] = True + reads.append(read) + batch = { + "attention_mask": pad_sequence([torch.ones_like(read, dtype=torch.long) for read in reads], batch_first=True), + "logits_to_keep": pad_sequence(reads, batch_first=True), + } + for key in inputs[0].keys() - {"marks", "attention_mask"}: + values = [prompt[key] for prompt in inputs] + if values[0].shape[-1] == inputs[0]["input_ids"].shape[-1]: # per token: `input_ids`, `mm_token_type_ids`... + batch[key] = pad_sequence([value[0] for value in values], batch_first=True) + else: + batch[key] = torch.cat(values) + return batch + + +def process_scores(decision_config, batch, questions, scores) -> list[list[torch.Tensor]]: + """ + The logits processors of the decision config on the option scores of each question of each prompt of a batch from + [`batch_prompts`], `scores` being those of all its questions in order. + """ + scores, processed = iter(scores), [] + lengths = batch["attention_mask"].sum(-1).tolist() + for input_ids, read, length, group in zip(batch["input_ids"], batch["logits_to_keep"], lengths, questions): + tail = input_ids[int(read.nonzero()[-1]) + 1 : length] # the tokens after the last mark: the state + processed.append([decision_config.logits_processors[q.type](tail, next(scores).float()) for q in group]) + return processed def _hub_kwargs(kwargs: dict) -> dict: @@ -53,6 +85,8 @@ def from_pretrained(cls, pretrained_model_name_or_path, *args, config=None, deci head_config, decision_config = cls.checkpoint_config(pretrained_model_name_or_path, **kwargs) if head_config is None and config is not None: raise ValueError(f"{pretrained_model_name_or_path} has no {cls.__name__} head next to its weights.") + if head_config is not None and config is not None: # the auto classes set `attn_implementation` on it + kwargs.setdefault("attn_implementation", config._attn_implementation_internal) config = head_config or config model = super().from_pretrained(pretrained_model_name_or_path, *args, config=config, **kwargs) if decision_config is None: # the decision config of the repo, read now like `generation_config` @@ -100,29 +134,28 @@ def save_pretrained(self, save_directory, *args, **kwargs): def decide( self, input_ids: torch.LongTensor, - marks: torch.LongTensor, - decision_questions: list[Question], + logits_to_keep: torch.BoolTensor, + decision_questions: list[list[Question]], decision_config: DecisionConfig | None = None, **kwargs, - ) -> list[torch.Tensor]: + ) -> list[list[torch.Tensor]]: """ - Scores the options of each question of a prompt from `decision_config.prepare_inputs`: one forward pass, - then the logits processors of the decision config. Returns the scores of each question, which + Scores the options of each question of a batch of prompts from [`batch_prompts`]: one forward pass, then the + logits processors of the decision config. Returns the scores of each question of each prompt, which `Question.answer` turns into answers. """ decision_config = decision_config or self.decision_config - inputs = {"input_ids": input_ids, **kwargs} - num_options = torch.tensor([[len(question.options) for question in decision_questions]]) - question_types = [Question.questions_types.index(question.type) for question in decision_questions] + batch = {"input_ids": input_ids, "logits_to_keep": logits_to_keep, **kwargs} + batch = {key: value.to(self.device) if hasattr(value, "to") else value for key, value in batch.items()} + questions = [question for group in decision_questions for question in group] + num_options = [len(question.options) for question in questions] + question_types = [Question.questions_types.index(question.type) for question in questions] scores = self( - marks=marks[None].to(self.device), - num_options=num_options.to(self.device), - question_type_ids=torch.tensor([question_types], device=self.device), - **{key: value.to(self.device) if hasattr(value, "to") else value for key, value in inputs.items()}, - ) - return process_scores( - decision_config, inputs["input_ids"], marks, decision_questions, scores.split(num_options[0].tolist()) + num_options=torch.tensor(num_options, device=self.device), + question_type_ids=torch.tensor(question_types, device=self.device), + **batch, ) + return process_scores(decision_config, batch, decision_questions, scores.split(num_options)) __all__ = ["DecisionMixin"] diff --git a/src/transformers/models/jev/modeling_jev.py b/src/transformers/models/jev/modeling_jev.py index 1ebff6d1111d..896e95473589 100644 --- a/src/transformers/models/jev/modeling_jev.py +++ b/src/transformers/models/jev/modeling_jev.py @@ -16,7 +16,7 @@ from ...decision import Question from ...decision.modeling_utils import DecisionMixin -from ...modeling_utils import PreTrainedModel +from ...modeling_utils import PreTrainedModel, slice_logits_to_keep from ...utils import auto_docstring from ..auto import AutoModel from .configuration_jev import JevConfig @@ -46,15 +46,15 @@ def __init__(self, config: JevConfig): def forward( self, input_ids: torch.LongTensor, - marks: torch.LongTensor, + logits_to_keep: torch.BoolTensor, num_options: torch.LongTensor, question_type_ids: torch.LongTensor, **kwargs, ) -> torch.Tensor: - # A batch of one prompt - marks, num_options, question_type_ids = marks[0], num_options[0], question_type_ids[0] - hidden_states = self.model(input_ids=input_ids, use_cache=False, **kwargs).last_hidden_state[0] - logits = self.proj(hidden_states[marks - 1].to(self.proj.weight.dtype)) + # The hidden states before the marks of the batch: one per question, in order + hidden_states = self.model(input_ids=input_ids, use_cache=False, **kwargs).last_hidden_state + hidden_states = slice_logits_to_keep(hidden_states, logits_to_keep) + logits = self.proj(hidden_states.to(self.proj.weight.dtype)) if self.config.softcap is not None: logits = self.config.softcap * torch.tanh(logits / self.config.softcap) # The options of a question are scored by the rows from the offset of its type on, in order diff --git a/src/transformers/pipelines/__init__.py b/src/transformers/pipelines/__init__.py index 6f4c6552dfbf..235d3a684dbd 100755 --- a/src/transformers/pipelines/__init__.py +++ b/src/transformers/pipelines/__init__.py @@ -1052,8 +1052,8 @@ def pipeline( ) task = get_task(model, token) - # Retrieve the task - if task in custom_tasks: + # Retrieve the task: the repo's own pipeline, unless remote code is refused for a task we implement + if task in custom_tasks and not (trust_remote_code is False and task in SUPPORTED_TASKS): targeted_task, task_options = clean_custom_task(custom_tasks[task]) if pipeline_class is None: if not trust_remote_code: diff --git a/src/transformers/pipelines/decision.py b/src/transformers/pipelines/decision.py index 59bbc8dd8855..478d051e6add 100644 --- a/src/transformers/pipelines/decision.py +++ b/src/transformers/pipelines/decision.py @@ -8,7 +8,7 @@ if is_torch_available(): - from ..decision.modeling_utils import DecisionMixin, process_scores + from ..decision.modeling_utils import DecisionMixin, batch_prompts, process_scores @add_end_docstrings(build_pipeline_init_args(has_tokenizer=True)) @@ -81,21 +81,15 @@ def preprocess(self, request: dict[str, Any] | str): return {"prompts": self.decision_config.prepare_inputs(request, self.processor or self.tokenizer)} def _forward(self, model_inputs): - # TODO: batch / continious batching model, decision_config, prompts = self.model, self.decision_config, model_inputs["prompts"] + batch, questions = batch_prompts(prompts), [questions for _, questions in prompts] if isinstance(model, DecisionMixin): - scores = [ - model.decide(**inputs, decision_questions=q, decision_config=decision_config) for inputs, q in prompts - ] + scores = model.decide(**batch, decision_questions=questions, decision_config=decision_config) return {"prompts": prompts, "scores": scores} - scores = [] - for inputs, questions in prompts: # the logits of the labels before each mark - inputs = dict(inputs) - marks = inputs.pop("marks") - logits = model(**inputs, use_cache=False, logits_to_keep=marks - 1).logits[0] - labels = [logits[i, self.label_token_ids(tuple(q.labels))] for i, q in enumerate(questions)] - scores.append(process_scores(decision_config, inputs["input_ids"], marks, questions, labels)) - return {"prompts": prompts, "scores": scores} + logits = model(**batch, use_cache=False).logits # the logits before each mark: one row per question + flat = [question for group in questions for question in group] + labels = [logits[i, self.label_token_ids(tuple(q.labels))] for i, q in enumerate(flat)] + return {"prompts": prompts, "scores": process_scores(decision_config, batch, questions, labels)} def postprocess(self, model_outputs): answers = {} diff --git a/src/transformers/utils/auto_docstring.py b/src/transformers/utils/auto_docstring.py index a6283f9c6a17..b8095fd5eb57 100644 --- a/src/transformers/utils/auto_docstring.py +++ b/src/transformers/utils/auto_docstring.py @@ -2232,23 +2232,16 @@ class ModelArgs: "shape": None, } - marks = { - "description": """ - Positions recorded by the `<|mark|>` of a decision template: the number of tokens of the prompt before each of them. - """, - "shape": "of shape `(num_marks,)`", - } - num_options = { "description": """ - Number of options of each question of the prompt. + Number of options of each question of the batch, in order. """, "shape": "of shape `(num_questions,)`", } question_type_ids = { "description": """ - Type of each question of the prompt, as its index in `Question.questions_types`. + Type of each question of the batch, as its index in `Question.questions_types`. """, "shape": "of shape `(num_questions,)`", } diff --git a/tests/models/jev/test_modeling_jev.py b/tests/models/jev/test_modeling_jev.py index a399a633aef1..2467df029c0d 100644 --- a/tests/models/jev/test_modeling_jev.py +++ b/tests/models/jev/test_modeling_jev.py @@ -30,7 +30,7 @@ if is_torch_available(): import torch - from transformers import JevForDecision + from transformers import AutoModelForDecision, JevForDecision @require_torch @@ -70,6 +70,7 @@ def test_from_a_vllm_decision_checkpoint(self): Path(tmp, "decision_config.json").write_text(json.dumps(settings), encoding="utf-8") save_file({"weight": readout}, f"{tmp}/readout.safetensors", metadata={"format": "pt"}) model = JevForDecision.from_pretrained(tmp) + eager = AutoModelForDecision.from_pretrained(tmp, attn_implementation="eager") decide = pipeline("decision", model=tmp) with tempfile.TemporaryDirectory() as tmp: # saved as a decision model, with its decision config model.save_pretrained(tmp) @@ -78,6 +79,7 @@ def test_from_a_vllm_decision_checkpoint(self): torch.testing.assert_close(reloaded.proj.weight, readout) self.assertEqual(model.config.num_verbalizers, 3) self.assertFalse(model.config.text_config.is_causal) + self.assertEqual(eager.model.config._attn_implementation, "eager") # not lost to the config of its head torch.testing.assert_close(model.proj.weight, readout) self.assertIsInstance(decide.model, JevForDecision) self.assertEqual(decide.decision_config.temperature, dict.fromkeys(["noul", "choice", "score"], 2.0)) @@ -127,9 +129,9 @@ def test_config_json_cannot_add_extra_files(self): class JevIntegrationTest(unittest.TestCase): def test_jev_9b(self): expected = { - "intent": {"refund": 0.9999, "track": 0.0000, "other": 0.0001}, - "urgent": 0.5419, - "anger": {"0": 0.0165, "1": 0.1113, "2": 0.5163, "3": 0.3126, "4": 0.0419, "5": 0.0013}, + "intent": {"refund": 0.9998, "track": 0.0000, "other": 0.0002}, + "urgent": 0.5292, + "anger": {"0": 0.0216, "1": 0.1246, "2": 0.5156, "3": 0.2958, "4": 0.0410, "5": 0.0014}, } check_decision( self, "hf-internal-testing/JEV-9B", "hf-internal-testing/JEV-9B", expected, atol=1e-2, dtype="bfloat16" @@ -138,9 +140,9 @@ def test_jev_9b(self): def test_jev_35b(self): # System 1 merged into the weights: its head is rows of the `lm_head` expected = { - "intent": {"refund": 0.9981, "track": 0.0001, "other": 0.0018}, - "urgent": 0.7539, - "anger": {"0": 0.0155, "1": 0.0859, "2": 0.4147, "3": 0.3896, "4": 0.0872, "5": 0.007}, + "intent": {"refund": 0.9969, "track": 0.0001, "other": 0.0030}, + "urgent": 0.7632, + "anger": {"0": 0.0145, "1": 0.0825, "2": 0.4165, "3": 0.3914, "4": 0.0878, "5": 0.0073}, } check_decision(self, "hf-internal-testing/JEV-35B", "hf-internal-testing/JEV-35B", expected, atol=1e-2) diff --git a/tests/test_decision_head_common.py b/tests/test_decision_head_common.py index 18709c85a10f..3e6119cc29be 100644 --- a/tests/test_decision_head_common.py +++ b/tests/test_decision_head_common.py @@ -20,13 +20,13 @@ from huggingface_hub import hf_hub_download from transformers import AutoConfig, AutoTokenizer, DecisionConfig, is_torch_available, pipeline -from transformers.testing_utils import torch_device if is_torch_available(): import torch from transformers import AutoModelForDecision, Qwen2Config, Qwen2ForCausalLM + from transformers.decision.modeling_utils import batch_prompts def render(processor, request, config=None): @@ -89,14 +89,13 @@ def test_chat_template_reproduces_the_model_code(self): sequences = [{"input_ids": prompt.input_ids, "marks": prompt.marks} for prompt in self.prompts] self.assertEqual(sequences, self.test["sequences"]) - def scores(self, model): + def scores(self, model, prompts=None): + """The scores of each question of each prompt, from one batch of `prompts` (all of them by default).""" + prompts = prompts or self.prompts + batch = batch_prompts([(prompt.inputs, prompt.questions) for prompt in prompts]) + questions = [prompt.questions for prompt in prompts] with torch.no_grad(): - return [ - model.decide( - **prompt.inputs, decision_questions=prompt.questions, decision_config=self.decision_config - ) - for prompt in self.prompts - ] + return model.decide(**batch, decision_questions=questions, decision_config=self.decision_config) def test_scores_and_save_load(self): for model_class in self.all_model_classes: @@ -108,6 +107,9 @@ def test_scores_and_save_load(self): for prompt, scores, loaded_scores in zip(self.prompts, self.scores(model), self.scores(loaded)): torch.testing.assert_close(loaded_scores, scores) self.assertEqual([len(s) for s in scores], [len(question.options) for question in prompt.questions]) + # The padded batch scores each prompt like it alone + one_by_one = [scores for prompt in self.prompts for scores in self.scores(model, [prompt])] + torch.testing.assert_close(self.scores(model), one_by_one) def check_decision( @@ -116,8 +118,15 @@ def check_decision( """The answers to the request of `decision_test.json` (probabilities in order, or `noul`) match `expected`.""" test = json.loads(Path(hf_hub_download(decision_repo, "decision_test.json")).read_bytes()) # By default the original checkpoint, which the decision files are converted from + # Some checkpoints ship their own remote-code decision pipeline: this one is ours decide = pipeline( - "decision", model=model or test["model"], decision_config=decision_repo, device=torch_device, **pipeline_kwargs + "decision", + model=model or test["model"], + decision_config=decision_repo, + # 27-35B checkpoints: sharded over the 24 GB CI GPUs + device_map="auto", + trust_remote_code=False, + **pipeline_kwargs, ) answers = decide(test["request"]) for name, value in expected.items():