



import inspect
import unittest

import torch
import torch.nn as nn

from accelerate.hooks import (
    AlignDevicesHook,
    ModelHook,
    SequentialHook,
    add_hook_to_module,
    attach_align_device_hook,
    remove_hook_from_module,
    remove_hook_from_submodules,
)
from accelerate.test_utils import require_multi_gpu






class      UpperCamelCase    (						nn.Module     ):



	def __init__(       self		):
							super().__init__()
							A__							=       nn.Linear(3			,	4		)
							A__							=       nn.BatchNormad(4		)
							A__							=       nn.Linear(4			,	5		)



	def 			__A							(       self			,	UpperCAmelCase__		):
							return self.lineara(self.batchnorm(self.lineara(UpperCAmelCase__		)		)		)






class      UpperCamelCase    (						_UpperCAmelCase     ):



	def 			__A							(       self			,	UpperCAmelCase__			,	*UpperCAmelCase__			,	**UpperCAmelCase__		):
							return (args[0] + 1,) + args[1:], kwargs






class      UpperCamelCase    (						_UpperCAmelCase     ):



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__		):
							return output + 1






class      UpperCamelCase    (						unittest.TestCase     ):



	def 			__A							(       self		):
							A__							=       ModelForTest()
							A__							=       ModelHook()

							add_hook_to_module(UpperCAmelCase__			,	UpperCAmelCase__		)
							self.assertEqual(test_model._hf_hook			,	UpperCAmelCase__		)
							self.assertTrue(hasattr(UpperCAmelCase__			,	"_old_forward"		)		)

							# Check adding the hook did not change the name or the signature
							self.assertEqual(test_model.forward.__name__			,	"forward"		)
							self.assertListEqual(list(inspect.signature(test_model.forward		).parameters		)			,	["x"]		)

							remove_hook_from_module(UpperCAmelCase__		)
							self.assertFalse(hasattr(UpperCAmelCase__			,	"_hf_hook"		)		)
							self.assertFalse(hasattr(UpperCAmelCase__			,	"_old_forward"		)		)



	def 			__A							(       self		):
							A__							=       ModelForTest()
							A__							=       ModelHook()

							add_hook_to_module(UpperCAmelCase__			,	UpperCAmelCase__		)
							add_hook_to_module(UpperCAmelCase__			,	UpperCAmelCase__			,	append=UpperCAmelCase__		)

							self.assertEqual(isinstance(test_model._hf_hook			,	UpperCAmelCase__		)			,	UpperCAmelCase__		)
							self.assertEqual(len(test_model._hf_hook.hooks		)			,	2		)
							self.assertTrue(hasattr(UpperCAmelCase__			,	"_old_forward"		)		)

							# Check adding the hook did not change the name or the signature
							self.assertEqual(test_model.forward.__name__			,	"forward"		)
							self.assertListEqual(list(inspect.signature(test_model.forward		).parameters		)			,	["x"]		)

							remove_hook_from_module(UpperCAmelCase__		)
							self.assertFalse(hasattr(UpperCAmelCase__			,	"_hf_hook"		)		)
							self.assertFalse(hasattr(UpperCAmelCase__			,	"_old_forward"		)		)



	def 			__A							(       self		):
							A__							=       ModelForTest()
							A__							=       torch.randn(2			,	3		)
							A__							=       test_model(x + 1		)
							A__							=       test_model(x + 2		)

							A__							=       PreForwardHook()
							add_hook_to_module(UpperCAmelCase__			,	UpperCAmelCase__		)
							A__							=       test_model(UpperCAmelCase__		)
							self.assertTrue(torch.allclose(UpperCAmelCase__			,	UpperCAmelCase__			,	atol=1e-5		)		)

							# Attaching a hook to a model when it already has one replaces, does not chain
							A__							=       PreForwardHook()
							add_hook_to_module(UpperCAmelCase__			,	UpperCAmelCase__		)
							A__							=       test_model(UpperCAmelCase__		)
							self.assertTrue(torch.allclose(UpperCAmelCase__			,	UpperCAmelCase__			,	atol=1e-5		)		)

							# You need to use the sequential hook to chain two or more hooks
							A__							=       SequentialHook(PreForwardHook()			,	PreForwardHook()		)
							add_hook_to_module(UpperCAmelCase__			,	UpperCAmelCase__		)

							A__							=       test_model(UpperCAmelCase__		)
							assert torch.allclose(UpperCAmelCase__			,	UpperCAmelCase__			,	atol=1e-5		)



	def 			__A							(       self		):
							A__							=       ModelForTest()
							A__							=       torch.randn(2			,	3		)
							A__							=       test_model(UpperCAmelCase__		)

							A__							=       PostForwardHook()
							add_hook_to_module(UpperCAmelCase__			,	UpperCAmelCase__		)
							A__							=       test_model(UpperCAmelCase__		)
							self.assertTrue(torch.allclose(UpperCAmelCase__			,	output + 1			,	atol=1e-5		)		)

							# Attaching a hook to a model when it already has one replaces, does not chain
							A__							=       PostForwardHook()
							add_hook_to_module(UpperCAmelCase__			,	UpperCAmelCase__		)
							A__							=       test_model(UpperCAmelCase__		)
							self.assertTrue(torch.allclose(UpperCAmelCase__			,	output + 1			,	atol=1e-5		)		)

							# You need to use the sequential hook to chain two or more hooks
							A__							=       SequentialHook(PostForwardHook()			,	PostForwardHook()		)
							add_hook_to_module(UpperCAmelCase__			,	UpperCAmelCase__		)

							A__							=       test_model(UpperCAmelCase__		)
							assert torch.allclose(UpperCAmelCase__			,	output + 2			,	atol=1e-5		)



	def 			__A							(       self		):
							A__							=       ModelForTest()
							A__							=       torch.randn(2			,	3		)
							A__							=       test_model(UpperCAmelCase__		)

							A__							=       PostForwardHook()
							add_hook_to_module(UpperCAmelCase__			,	UpperCAmelCase__		)
							A__							=       test_model(UpperCAmelCase__		)
							self.assertTrue(torch.allclose(UpperCAmelCase__			,	output + 1		)		)
							self.assertTrue(outputa.requires_grad		)

							A__							=       True
							A__							=       test_model(UpperCAmelCase__		)
							self.assertFalse(outputa.requires_grad		)



	@require_multi_gpu
	def 			__A							(       self		):
							A__							=       ModelForTest()
							# Everything is on CPU
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)

							# This will move each submodule on different devices
							add_hook_to_module(model.lineara			,	AlignDevicesHook(execution_device=0		)		)
							add_hook_to_module(model.batchnorm			,	AlignDevicesHook(execution_device=0		)		)
							add_hook_to_module(model.lineara			,	AlignDevicesHook(execution_device=1		)		)

							self.assertEqual(model.lineara.weight.device			,	torch.device(0		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device(0		)		)
							self.assertEqual(model.batchnorm.running_mean.device			,	torch.device(0		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device(1		)		)

							# We can still make a forward pass. The input does not need to be on any particular device
							A__							=       torch.randn(2			,	3		)
							A__							=       model(UpperCAmelCase__		)
							self.assertEqual(output.device			,	torch.device(1		)		)

							# We can add a general hook to put back output on same device as input.
							add_hook_to_module(UpperCAmelCase__			,	AlignDevicesHook(io_same_device=UpperCAmelCase__		)		)
							A__							=       torch.randn(2			,	3		).to(0		)
							A__							=       model(UpperCAmelCase__		)
							self.assertEqual(output.device			,	torch.device(0		)		)



	def 			__A							(       self		):
							A__							=       ModelForTest()

							# Everything is on CPU
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)

							# This will move each submodule on different devices
							A__							=       {"execution_device": 0 if torch.cuda.is_available() else "cpu", "offload": True}

							add_hook_to_module(model.lineara			,	AlignDevicesHook(**UpperCAmelCase__		)		)
							add_hook_to_module(model.batchnorm			,	AlignDevicesHook(**UpperCAmelCase__		)		)
							add_hook_to_module(model.lineara			,	AlignDevicesHook(**UpperCAmelCase__		)		)

							# Parameters have been offloaded, so on the meta device
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							# Buffers are not included in the offload by default, so are on the execution device
							A__							=       torch.device(hook_kwargs["execution_device"]		)
							self.assertEqual(model.batchnorm.running_mean.device			,	UpperCAmelCase__		)

							A__							=       torch.randn(2			,	3		)
							A__							=       model(UpperCAmelCase__		)
							self.assertEqual(output.device			,	UpperCAmelCase__		)

							# Removing hooks loads back the weights in the model.
							remove_hook_from_module(model.lineara		)
							remove_hook_from_module(model.batchnorm		)
							remove_hook_from_module(model.lineara		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)

							# Now test with buffers included in the offload
							A__							=       {
							    "execution_device": 0 if torch.cuda.is_available() else "cpu",
							    "offload": True,
							    "offload_buffers": True,
							}

							add_hook_to_module(model.lineara			,	AlignDevicesHook(**UpperCAmelCase__		)		)
							add_hook_to_module(model.batchnorm			,	AlignDevicesHook(**UpperCAmelCase__		)		)
							add_hook_to_module(model.lineara			,	AlignDevicesHook(**UpperCAmelCase__		)		)

							# Parameters have been offloaded, so on the meta device, buffers included
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.batchnorm.running_mean.device			,	torch.device("meta"		)		)

							A__							=       torch.randn(2			,	3		)
							A__							=       model(UpperCAmelCase__		)
							self.assertEqual(output.device			,	UpperCAmelCase__		)

							# Removing hooks loads back the weights in the model.
							remove_hook_from_module(model.lineara		)
							remove_hook_from_module(model.batchnorm		)
							remove_hook_from_module(model.lineara		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)



	def 			__A							(       self		):
							A__							=       ModelForTest()

							# Everything is on CPU
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)

							# This will move each submodule on different devices
							A__							=       0 if torch.cuda.is_available() else "cpu"
							attach_align_device_hook(UpperCAmelCase__			,	execution_device=UpperCAmelCase__			,	offload=UpperCAmelCase__		)

							# Parameters have been offloaded, so on the meta device
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							# Buffers are not included in the offload by default, so are on the execution device
							A__							=       torch.device(UpperCAmelCase__		)
							self.assertEqual(model.batchnorm.running_mean.device			,	UpperCAmelCase__		)

							A__							=       torch.randn(2			,	3		)
							A__							=       model(UpperCAmelCase__		)
							self.assertEqual(output.device			,	UpperCAmelCase__		)

							# Removing hooks loads back the weights in the model.
							remove_hook_from_submodules(UpperCAmelCase__		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)

							# Now test with buffers included in the offload
							attach_align_device_hook(UpperCAmelCase__			,	execution_device=UpperCAmelCase__			,	offload=UpperCAmelCase__			,	offload_buffers=UpperCAmelCase__		)

							# Parameters have been offloaded, so on the meta device, buffers included
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.batchnorm.running_mean.device			,	torch.device("meta"		)		)

							A__							=       torch.randn(2			,	3		)
							A__							=       model(UpperCAmelCase__		)
							self.assertEqual(output.device			,	UpperCAmelCase__		)

							# Removing hooks loads back the weights in the model.
							remove_hook_from_submodules(UpperCAmelCase__		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)



	def 			__A							(       self		):
							A__							=       ModelForTest()

							# Everything is on CPU
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)

							# This will move each submodule on different devices
							A__							=       0 if torch.cuda.is_available() else "cpu"
							attach_align_device_hook(
							    UpperCAmelCase__			,	execution_device=UpperCAmelCase__			,	offload=UpperCAmelCase__			,	weights_map=model.state_dict()		)

							# Parameters have been offloaded, so on the meta device
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							# Buffers are not included in the offload by default, so are on the execution device
							A__							=       torch.device(UpperCAmelCase__		)
							self.assertEqual(model.batchnorm.running_mean.device			,	UpperCAmelCase__		)

							A__							=       torch.randn(2			,	3		)
							A__							=       model(UpperCAmelCase__		)
							self.assertEqual(output.device			,	UpperCAmelCase__		)

							# Removing hooks loads back the weights in the model.
							remove_hook_from_submodules(UpperCAmelCase__		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)

							# Now test with buffers included in the offload
							attach_align_device_hook(
							    UpperCAmelCase__			,	execution_device=UpperCAmelCase__			,	offload=UpperCAmelCase__			,	weights_map=model.state_dict()			,	offload_buffers=UpperCAmelCase__			,	)

							# Parameters have been offloaded, so on the meta device, buffers included
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.batchnorm.running_mean.device			,	torch.device("meta"		)		)

							A__							=       torch.randn(2			,	3		)
							A__							=       model(UpperCAmelCase__		)
							self.assertEqual(output.device			,	UpperCAmelCase__		)

							# Removing hooks loads back the weights in the model.
							remove_hook_from_submodules(UpperCAmelCase__		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)










import datasets

from .evaluate import evaluate


UpperCAmelCase_							:      List[Any]      =  "\\n@inproceedings{Rajpurkar2016SQuAD10,\n  title={SQuAD: 100, 000+ Questions for Machine Comprehension of Text},\n  author={Pranav Rajpurkar and Jian Zhang and Konstantin Lopyrev and Percy Liang},\n  booktitle={EMNLP},\n  year={2016}\n}\n"

UpperCAmelCase_							:      Any      =  "\nThis metric wrap the official scoring script for version 1 of the Stanford Question Answering Dataset (SQuAD).\n\nStanford Question Answering Dataset (SQuAD) is a reading comprehension dataset, consisting of questions posed by\ncrowdworkers on a set of Wikipedia articles, where the answer to every question is a segment of text, or span,\nfrom the corresponding reading passage, or the question might be unanswerable.\n"

UpperCAmelCase_							:      Tuple      =  "\nComputes SQuAD scores (F1 and EM).\nArgs:\n    predictions: List of question-answers dictionaries with the following key-values:\n        - 'id': id of the question-answer pair as given in the references (see below)\n        - 'prediction_text': the text of the answer\n    references: List of question-answers dictionaries with the following key-values:\n        - 'id': id of the question-answer pair (see above),\n        - 'answers': a Dict in the SQuAD dataset format\n            {\n                'text': list of possible texts for the answer, as a list of strings\n                'answer_start': list of start positions for the answer, as a list of ints\n            }\n            Note that answer_start values are not taken into account to compute the metric.\nReturns:\n    'exact_match': Exact match (the normalized answer exactly match the gold answer)\n    'f1': The F-score of predicted tokens versus the gold answer\nExamples:\n\n    >>> predictions = [{'prediction_text': '1976', 'id': '56e10a3be3433e1400422b22'}]\n    >>> references = [{'answers': {'answer_start': [97], 'text': ['1976']}, 'id': '56e10a3be3433e1400422b22'}]\n    >>> squad_metric = datasets.load_metric(\"squad\")\n    >>> results = squad_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {'exact_match': 100.0, 'f1': 100.0}\n"






@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION       ,    _KWARGS_DESCRIPTION     )
class      UpperCamelCase    (						datasets.Metric     ):



	def 			__A							(       self		):
							return datasets.MetricInfo(
							    description=_DESCRIPTION			,	citation=_CITATION			,	inputs_description=_KWARGS_DESCRIPTION			,	features=datasets.Features(
							        {
							            "predictions": {"id": datasets.Value("string"		), "prediction_text": datasets.Value("string"		)},
							            "references": {
							                "id": datasets.Value("string"		),
							                "answers": datasets.features.Sequence(
							                    {
							                        "text": datasets.Value("string"		),
							                        "answer_start": datasets.Value("int32"		),
							                    }		),
							            },
							        }		)			,	codebase_urls=["https://rajpurkar.github.io/SQuAD-explorer/"]			,	reference_urls=["https://rajpurkar.github.io/SQuAD-explorer/"]			,	)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__		):
							A__							=       {prediction["id"]: prediction["prediction_text"] for prediction in predictions}
							A__							=       [
							    {
							        "paragraphs": [
							            {
							                "qas": [
							                    {
							                        "answers": [{"text": answer_text} for answer_text in ref["answers"]["text"]],
							                        "id": ref["id"],
							                    }
							                    for ref in references
							                ]
							            }
							        ]
							    }
							]
							A__							=       evaluate(dataset=UpperCAmelCase__			,	predictions=UpperCAmelCase__		)
							return score





