Find the labels used by a given model. Args: model_class (`type`): The class of the model.
(model_class)
| 563 | |
| 564 | |
| 565 | def find_labels(model_class): |
| 566 | """ |
| 567 | Find the labels used by a given model. |
| 568 | |
| 569 | Args: |
| 570 | model_class (`type`): The class of the model. |
| 571 | """ |
| 572 | model_name = model_class.__name__ |
| 573 | framework = infer_framework(model_class) |
| 574 | if framework == "tf": |
| 575 | signature = inspect.signature(model_class.call) # TensorFlow models |
| 576 | elif framework == "pt": |
| 577 | signature = inspect.signature(model_class.forward) # PyTorch models |
| 578 | else: |
| 579 | signature = inspect.signature(model_class.__call__) # Flax models |
| 580 | |
| 581 | if "QuestionAnswering" in model_name: |
| 582 | return [p for p in signature.parameters if "label" in p or p in ("start_positions", "end_positions")] |
| 583 | else: |
| 584 | return [p for p in signature.parameters if "label" in p] |
| 585 | |
| 586 | |
| 587 | def flatten_dict(d: MutableMapping, parent_key: str = "", delimiter: str = "."): |