
'''simple docstring'''
import unittest

from transformers.testing_utils import is_torch_available, require_torch, tooslow

from ...generation.test_utils import torch_device
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import (
						    CpmAntConfig,
						    CpmAntForCausalLM,
						    CpmAntModel,
						    CpmAntTokenizer,
						)


@require_torch
class        SCREAMING_SNAKE_CASE__	:







				def __init__(    self		,							A_		,							A_=2		,							A_=8		,							A_=True		,							A_=False		,							A_=False		,							A_=False		,							A_=False		,							A_=99		,							A_=32		,							A_=3		,							A_=4		,							A_=37		,							A_=32		,							A_=128		,							A_=8		,							A_=8		,							A_=8		,							A_=1.0		,							A_=True		,							)->						str:







											'''simple docstring'''






											UpperCamelCase					       =			parent
											UpperCamelCase					       =			batch_size
											UpperCamelCase					       =			seq_length
											UpperCamelCase					       =			is_training
											UpperCamelCase					       =			use_token_type_ids
											UpperCamelCase					       =			use_input_mask
											UpperCamelCase					       =			use_labels
											UpperCamelCase					       =			use_mc_token_ids
											UpperCamelCase					       =			vocab_size
											UpperCamelCase					       =			hidden_size
											UpperCamelCase					       =			num_hidden_layers
											UpperCamelCase					       =			num_attention_heads
											UpperCamelCase					       =			intermediate_size
											UpperCamelCase					       =			num_buckets
											UpperCamelCase					       =			max_distance
											UpperCamelCase					       =			prompt_length
											UpperCamelCase					       =			prompt_types
											UpperCamelCase					       =			segment_types
											UpperCamelCase					       =			init_std
											UpperCamelCase					       =			return_dict







				def 		UpperCAmelCase_  (    self			)->						Dict:







											'''simple docstring'''






											UpperCamelCase					       =			{}
											UpperCamelCase					       =			ids_tensor([self.batch_size, self.seq_length]		,							self.vocab_size			).type(torch.intaa			)
											UpperCamelCase					       =			False

											UpperCamelCase					       =			self.get_config()

											return (config, input_ids)







				def 		UpperCAmelCase_  (    self			)->						List[Any]:







											'''simple docstring'''






											return CpmAntConfig(
											    vocab_size=self.vocab_size		,							hidden_size=self.hidden_size		,							num_hidden_layers=self.num_hidden_layers		,							num_attention_heads=self.num_attention_heads		,							dim_ff=self.intermediate_size		,							position_bias_num_buckets=self.num_buckets		,							position_bias_max_distance=self.max_distance		,							prompt_types=self.prompt_types		,							prompt_length=self.prompt_length		,							segment_types=self.segment_types		,							use_cache=A_		,							init_std=self.init_std		,							return_dict=self.return_dict		,							)







				def 		UpperCAmelCase_  (    self		,							A_		,							A_		,							*A_			)->						Union[str, Any]:







											'''simple docstring'''






											UpperCamelCase					       =			CpmAntModel(config=A_			)
											model.to(A_			)
											model.eval()

											UpperCamelCase					       =			model(**A_			).last_hidden_state

											self.parent.assertEqual(hidden_states.shape		,							(self.batch_size, self.seq_length, config.hidden_size)			)







				def 		UpperCAmelCase_  (    self		,							A_		,							A_		,							*A_			)->						str:







											'''simple docstring'''






											UpperCamelCase					       =			CpmAntForCausalLM(A_			)
											model.to(A_			)
											UpperCamelCase					       =			input_ids['input_ids'].to(A_			)
											model.eval()

											UpperCamelCase					       =			model(**A_			)
											self.parent.assertEqual(
											    model_output.logits.shape		,							(self.batch_size, self.seq_length, config.vocab_size + config.prompt_types * config.prompt_length)		,							)







				def 		UpperCAmelCase_  (    self			)->						Dict:







											'''simple docstring'''






											UpperCamelCase       ,     UpperCamelCase					       =			self.prepare_config_and_inputs()
											return config, inputs_dict


@require_torch
class        SCREAMING_SNAKE_CASE__	(			snake_case_				,    snake_case_				,    unittest.TestCase):
				lowerCAmelCase_			       =				(CpmAntModel, CpmAntForCausalLM) if is_torch_available() else ()
				lowerCAmelCase_			       =				(
				    {"""feature-extraction""": CpmAntModel, """text-generation""": CpmAntForCausalLM} if is_torch_available() else {}
				)

				lowerCAmelCase_			       =				False
				lowerCAmelCase_			       =				False
				lowerCAmelCase_			       =				False
				lowerCAmelCase_			       =				False
				lowerCAmelCase_			       =				False







				def 		UpperCAmelCase_  (    self			)->						Optional[Any]:







											'''simple docstring'''






											UpperCamelCase					       =			CpmAntModelTester(self			)
											UpperCamelCase					       =			ConfigTester(self		,							config_class=A_			)







				def 		UpperCAmelCase_  (    self			)->						Union[str, Any]:







											'''simple docstring'''






											self.config_tester.create_and_test_config_common_properties()
											self.config_tester.create_and_test_config_to_json_string()
											self.config_tester.create_and_test_config_to_json_file()
											self.config_tester.create_and_test_config_from_and_save_pretrained()
											self.config_tester.check_config_can_be_init_without_params()
											self.config_tester.check_config_arguments_init()







				def 		UpperCAmelCase_  (    self			)->						Tuple:







											'''simple docstring'''






											unittest.skip('CPMAnt doesn\'t support input_embeds.'			)(self.test_inputs_embeds			)







				def 		UpperCAmelCase_  (    self			)->						Any:







											'''simple docstring'''






											unittest.skip(
											    'CPMAnt doesn\'t support retain grad in hidden_states or attentions, because prompt management will peel off the output.hidden_states from graph.\
                 So is attentions. We strongly recommand you use loss to tune model.'			)(self.test_retain_grad_hidden_states_attentions			)







				def 		UpperCAmelCase_  (    self			)->						int:







											'''simple docstring'''






											UpperCamelCase       ,     UpperCamelCase					       =			self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_cpmant_model(A_		,							A_			)







				def 		UpperCAmelCase_  (    self			)->						List[Any]:







											'''simple docstring'''






											UpperCamelCase       ,     UpperCamelCase					       =			self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_lm_head_model(A_		,							A_			)


@require_torch
class        SCREAMING_SNAKE_CASE__	(			unittest.TestCase):







				@tooslow
				def 		UpperCAmelCase_  (    self			)->						List[str]:







											'''simple docstring'''






											UpperCamelCase					       =			'今天天气真好！'
											UpperCamelCase					       =			'openbmb/cpm-ant-10b'
											UpperCamelCase					       =			CpmAntModel.from_pretrained(A_			)
											UpperCamelCase					       =			CpmAntTokenizer.from_pretrained(A_			)
											UpperCamelCase					       =			tokenizer(A_		,							return_tensors='pt'			)
											UpperCamelCase					       =			model(**A_			).last_hidden_state

											UpperCamelCase					       =			torch.tensor(
											    [[[6.1_708, 5.9_244, 1.0_835], [6.5_207, 6.2_893, -11.3_324], [-1.0_107, -0.0_576, -5.9_577]]]		,							)
											self.assertTrue(torch.allclose(hidden_states[:, :3, :3]		,							A_		,							atol=1e-2			)			)


@require_torch
class        SCREAMING_SNAKE_CASE__	(			unittest.TestCase):







				@tooslow
				def 		UpperCAmelCase_  (    self			)->						Tuple:







											'''simple docstring'''






											UpperCamelCase					       =			'今天天气真好！'
											UpperCamelCase					       =			'openbmb/cpm-ant-10b'
											UpperCamelCase					       =			CpmAntForCausalLM.from_pretrained(A_			)
											UpperCamelCase					       =			CpmAntTokenizer.from_pretrained(A_			)
											UpperCamelCase					       =			tokenizer(A_		,							return_tensors='pt'			)
											UpperCamelCase					       =			model(**A_			).logits

											UpperCamelCase					       =			torch.tensor(
											    [[[-6.4_267, -6.4_083, -6.3_958], [-5.8_802, -5.9_447, -5.7_811], [-5.3_896, -5.4_820, -5.4_295]]]		,							)
											self.assertTrue(torch.allclose(hidden_states[:, :3, :3]		,							A_		,							atol=1e-2			)			)







				@tooslow
				def 		UpperCAmelCase_  (    self			)->						Union[str, Any]:







											'''simple docstring'''






											UpperCamelCase					       =			'openbmb/cpm-ant-10b'
											UpperCamelCase					       =			CpmAntForCausalLM.from_pretrained(A_			)
											UpperCamelCase					       =			CpmAntTokenizer.from_pretrained(A_			)
											UpperCamelCase					       =			'今天天气不错，'
											UpperCamelCase					       =			'今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的'
											UpperCamelCase					       =			tokenizer(A_		,							return_tensors='pt'			)
											UpperCamelCase					       =			model.generate(**A_			)
											UpperCamelCase					       =			tokenizer.batch_decode(A_			)
											self.assertEqual(A_		,							A_			)







				@tooslow
				def 		UpperCAmelCase_  (    self			)->						Optional[Any]:







											'''simple docstring'''






											UpperCamelCase					       =			'openbmb/cpm-ant-10b'
											UpperCamelCase					       =			CpmAntForCausalLM.from_pretrained(A_			)
											UpperCamelCase					       =			CpmAntTokenizer.from_pretrained(A_			)
											UpperCamelCase					       =			['今天天气不错，', '新年快乐，万事如意！']
											UpperCamelCase					       =			[
											    '今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的',
											    '新年快乐，万事如意！在这辞旧迎新的美好时刻，我谨代表《农村新技术》杂志社全体同仁，向一直以来关心、支持《农村新技术》杂志发展的各级领导、各界朋友和广大读者致以最诚挚的',
											]
											UpperCamelCase					       =			tokenizer(A_		,							return_tensors='pt'		,							padding=A_			)
											UpperCamelCase					       =			model.generate(**A_			)
											UpperCamelCase					       =			tokenizer.batch_decode(A_			)
											self.assertEqual(A_		,							A_			)






'''simple docstring'''
from typing import SupportsIndex

import numpy as np
from scipy.ndimage import convolve







def 						A_(			A		:						np.ndarray						,			A		:						np.ndarray						,			A		:						np.ndarray):
							UpperCamelCase					       =			np.stack((horizontal_flow, vertical_flow)						,			2)

							# Create a grid of all pixel coordinates and subtract the flow to get the
							# target pixels coordinates
							UpperCamelCase					       =			np.stack(
							    np.meshgrid(np.arange(0						,			image.shape[1])						,			np.arange(0						,			image.shape[0]))						,			2)
							UpperCamelCase					       =			np.round(grid - flow).astype(np.intaa)

							# Find the locations outside of the original image
							UpperCamelCase					       =			(grid < 0) | (grid >= np.array([image.shape[1], image.shape[0]]))
							UpperCamelCase					       =			0

							UpperCamelCase					       =			image[grid[:, :, 1], grid[:, :, 0]]

							# Set pixels at invalid locations to 0
							UpperCamelCase					       =			0

							return warped







def 						A_(			A		:						np.ndarray						,			A		:						np.ndarray						,			A		:						SupportsIndex						,			A		:						float | None = None						,			):
							if alpha is None:
														UpperCamelCase					       =			0.1

							# Initialize flow
							UpperCamelCase					       =			np.zeros_like(A)
							UpperCamelCase					       =			np.zeros_like(A)

							# Prepare kernels for the calculation of the derivatives and the average velocity
							UpperCamelCase					       =			np.array([[-1, 1], [-1, 1]]) * 0.25
							UpperCamelCase					       =			np.array([[-1, -1], [1, 1]]) * 0.25
							UpperCamelCase					       =			np.array([[1, 1], [1, 1]]) * 0.25
							UpperCamelCase					       =			np.array(
							    [[1 / 12, 1 / 6, 1 / 12], [1 / 6, 0, 1 / 6], [1 / 12, 1 / 6, 1 / 12]])

							# Iteratively refine the flow
							for _ in range(A):
														UpperCamelCase					       =			warp(A						,			A						,			A)
														UpperCamelCase					       =			convolve(A						,			A) + convolve(A						,			A)
														UpperCamelCase					       =			convolve(A						,			A) + convolve(A						,			A)
														UpperCamelCase					       =			convolve(A						,			A) + convolve(A						,			-kernel_t)

														UpperCamelCase					       =			convolve(A						,			A)
														UpperCamelCase					       =			convolve(A						,			A)

														# This updates the flow as proposed in the paper (Step 12)
														UpperCamelCase					       =			(
														    derivative_x * avg_horizontal_velocity
														    + derivative_y * avg_vertical_velocity
														    + derivative_t
														)
														UpperCamelCase					       =			update / (alpha**2 + derivative_x**2 + derivative_y**2)

														UpperCamelCase					       =			avg_horizontal_velocity - derivative_x * update
														UpperCamelCase					       =			avg_vertical_velocity - derivative_y * update

							return horizontal_flow, vertical_flow


if __name__ == "__main__":
						import doctest

						doctest.testmod()




