| 447 | |
| 448 | class QueryEmbeddingPipe(QueryEmbedding): |
| 449 | def forward(self, inputs, **kwargs): |
| 450 | if not hasattr(self, "_args"): |
| 451 | self._args = get_args() |
| 452 | |
| 453 | position_ids = inputs[0] |
| 454 | if hasattr(self._args, "attn_mask"): |
| 455 | attention_mask = None |
| 456 | else: |
| 457 | attention_mask = inputs[1] |
| 458 | |
| 459 | if len(inputs) == 3: |
| 460 | tokentype_ids = inputs[2] |
| 461 | else: |
| 462 | tokentype_ids = None |
| 463 | |
| 464 | embeddings = super().forward( |
| 465 | position_ids, tokentype_ids=tokentype_ids, |
| 466 | ) |
| 467 | |
| 468 | # If cmd args has attn_mask, we don't forward it as an activation. |
| 469 | if hasattr(self._args, "attn_mask"): |
| 470 | return embeddings |
| 471 | else: |
| 472 | assert False |
| 473 | return embeddings, attention_mask |
| 474 | |
| 475 | @property |
| 476 | def word_embeddings_weight(self): |