Extract a single tensor from model output, which might be a tuple, dict, or ModelOutput-like object.
(output: Any)
| 646 | # --------------------------------------------------------------------------- |
| 647 | |
| 648 | def extract_tensor(output: Any) -> torch.Tensor: |
| 649 | """ |
| 650 | Extract a single tensor from model output, which might be a tuple, dict, |
| 651 | or ModelOutput-like object. |
| 652 | """ |
| 653 | if isinstance(output, torch.Tensor): |
| 654 | return output |
| 655 | |
| 656 | # HuggingFace ModelOutput or similar dataclass-like object |
| 657 | if hasattr(output, "logits"): |
| 658 | return output.logits |
| 659 | if hasattr(output, "last_hidden_state"): |
| 660 | return output.last_hidden_state |
| 661 | |
| 662 | # Tuple/list: return first tensor element |
| 663 | if isinstance(output, (tuple, list)): |
| 664 | for item in output: |
| 665 | if isinstance(item, torch.Tensor): |
| 666 | return item |
| 667 | # Recurse into first element |
| 668 | if len(output) > 0: |
| 669 | return extract_tensor(output[0]) |
| 670 | |
| 671 | # Dict: try common keys |
| 672 | if isinstance(output, dict): |
| 673 | for key in ["logits", "last_hidden_state", "output", "hidden_states"]: |
| 674 | if key in output and isinstance(output[key], torch.Tensor): |
| 675 | return output[key] |
| 676 | # Return first tensor value |
| 677 | for v in output.values(): |
| 678 | if isinstance(v, torch.Tensor): |
| 679 | return v |
| 680 | |
| 681 | raise ValueError( |
| 682 | f"Cannot extract tensor from output of type {type(output)}. " |
| 683 | f"Consider adding support for this output format." |
| 684 | ) |
| 685 | |
| 686 | |
| 687 | def compare_outputs( |
no outgoing calls
no test coverage detected