o yia@s`ddlmZmZddlZddlmZddlmZddlmZm Z m Z m Z m Z GdddeZ dS))AnyDictN)Tensor)Metric) PREDS_TYPE TARGETS_TYPE_squad_compute_squad_input_check _squad_updatecseZdZUdZdZeed<dZeed<dZeed<e ed<e ed<e ed <d e ffd d Z d e de ddfddZdeee ffddZZS)SQuADa}Calculate `SQuAD Metric`_ which corresponds to the scoring script for version 1 of the Stanford Question Answering Dataset (SQuAD). As input to ``forward`` and ``update`` the metric accepts the following input: - ``preds`` (:class:`~Dict`): A Dictionary or List of Dictionary-s that map ``id`` and ``prediction_text`` to the respective values Example ``prediction``: .. code-block:: python {"prediction_text": "TorchMetrics is awesome", "id": "123"} - ``target`` (:class:`~Dict`): A Dictionary or List of Dictionary-s that contain the ``answers`` and ``id`` in the SQuAD Format. Example ``target``: .. code-block:: python { 'answers': [{'answer_start': [1], 'text': ['This is a test answer']}], 'id': '1', } Reference SQuAD Format: .. code-block:: python { 'answers': {'answer_start': [1], 'text': ['This is a test text']}, 'context': 'This is a test context.', 'id': '1', 'question': 'Is this a test?', 'title': 'train test' } As output of ``forward`` and ``compute`` the metric returns the following output: - ``squad`` (:class:`~Dict`): A dictionary containing the F1 score (key: "f1"), and Exact match score (key: "exact_match") for the batch. Args: kwargs: Additional keyword arguments, see :ref:`Metric kwargs` for more info. Example: >>> from torchmetrics import SQuAD >>> preds = [{"prediction_text": "1976", "id": "56e10a3be3433e1400422b22"}] >>> target = [{"answers": {"answer_start": [97], "text": ["1976"]}, "id": "56e10a3be3433e1400422b22"}] >>> squad = SQuAD() >>> squad(preds, target) {'exact_match': tensor(100.), 'f1': tensor(100.)} Fis_differentiableThigher_is_betterfull_state_updatef1_score exact_matchtotalkwargsc sjtjdi||jdtjdtjddd|jdtjdtjddd|jdtjdtjddddS) Nrr)dtypesum)namedefaultdist_reduce_fxrr)super__init__ add_statetorchtensorfloatint)selfr __class__rK/home/ubuntu/.local/lib/python3.10/site-packages/torchmetrics/text/squad.pyr^s zSQuAD.__init__predstargetreturnNcCsLt||\}}t||\}}}|j|7_|j|7_|j|7_dS)z*Update state with predictions and targets.N)r r rrr)r r$r% preds_dict target_dictrrrrrr#updatehs z SQuAD.updatecCst|j|j|jS)z5Aggregate the F1 Score and Exact match for the batch.)rrrr)r rrr#computepsz SQuAD.compute)__name__ __module__ __qualname____doc__r bool__annotations__r rrrrrrr)rstrr* __classcell__rrr!r#r s  8   r )typingrrrr torchmetricsr"torchmetrics.functional.text.squadrrrr r r rrrr#s