






import tempfile

import torch

from diffusers import PNDMScheduler

from .test_schedulers import SchedulerCommonTest



class 				__UpperCAmelCase		(lowerCamelCase_   ):
			__snake_case    :				str									=							(PNDMScheduler,)
			__snake_case    :				Tuple									=							(('''num_inference_steps''', 50),)







			def 							UpperCamelCase				(     self:  Optional[int]				,     **UpperCAmelCase_:  int					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  {
									    """num_train_timesteps""": 1_000,
									    """beta_start""": 0.00_01,
									    """beta_end""": 0.02,
									    """beta_schedule""": """linear""",
									}

									config.update(**__snake_case					)
									return config







			def 							UpperCamelCase				(     self:  Optional[Any]				,     UpperCAmelCase_:  Optional[Any]=0				,     **UpperCAmelCase_:  Dict					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  dict(self.forward_default_kwargs					)
									_SCREAMING_SNAKE_CASE     				=  kwargs.pop("""num_inference_steps"""				,     __snake_case					)
									_SCREAMING_SNAKE_CASE     				=  self.dummy_sample
									_SCREAMING_SNAKE_CASE     				=  0.1 * sample
									_SCREAMING_SNAKE_CASE     				=  [residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

									for scheduler_class in self.scheduler_classes:
															_SCREAMING_SNAKE_CASE     				=  self.get_scheduler_config(**__snake_case					)
															_SCREAMING_SNAKE_CASE     				=  scheduler_class(**__snake_case					)
															scheduler.set_timesteps(__snake_case					)
															# copy over dummy past residuals
															_SCREAMING_SNAKE_CASE     				=  dummy_past_residuals[:]

															with tempfile.TemporaryDirectory() as tmpdirname:
																					scheduler.save_config(__snake_case					)
																					_SCREAMING_SNAKE_CASE     				=  scheduler_class.from_pretrained(__snake_case					)
																					new_scheduler.set_timesteps(__snake_case					)
																					# copy over dummy past residuals
																					_SCREAMING_SNAKE_CASE     				=  dummy_past_residuals[:]

															_SCREAMING_SNAKE_CASE     				=  scheduler.step_prk(__snake_case				,     __snake_case				,     __snake_case				,     **__snake_case					).prev_sample
															_SCREAMING_SNAKE_CASE     				=  new_scheduler.step_prk(__snake_case				,     __snake_case				,     __snake_case				,     **__snake_case					).prev_sample

															assert torch.sum(torch.abs(output - new_output					)					) < 1E-5, "Scheduler outputs are not identical"

															_SCREAMING_SNAKE_CASE     				=  scheduler.step_plms(__snake_case				,     __snake_case				,     __snake_case				,     **__snake_case					).prev_sample
															_SCREAMING_SNAKE_CASE     				=  new_scheduler.step_plms(__snake_case				,     __snake_case				,     __snake_case				,     **__snake_case					).prev_sample

															assert torch.sum(torch.abs(output - new_output					)					) < 1E-5, "Scheduler outputs are not identical"







			def 							UpperCamelCase				(     self:  int					):







									'''simple docstring'''




									pass







			def 							UpperCamelCase				(     self:  List[str]				,     UpperCAmelCase_:  int=0				,     **UpperCAmelCase_:  Union[str, Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  dict(self.forward_default_kwargs					)
									_SCREAMING_SNAKE_CASE     				=  kwargs.pop("""num_inference_steps"""				,     __snake_case					)
									_SCREAMING_SNAKE_CASE     				=  self.dummy_sample
									_SCREAMING_SNAKE_CASE     				=  0.1 * sample
									_SCREAMING_SNAKE_CASE     				=  [residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

									for scheduler_class in self.scheduler_classes:
															_SCREAMING_SNAKE_CASE     				=  self.get_scheduler_config()
															_SCREAMING_SNAKE_CASE     				=  scheduler_class(**__snake_case					)
															scheduler.set_timesteps(__snake_case					)

															# copy over dummy past residuals (must be after setting timesteps)
															_SCREAMING_SNAKE_CASE     				=  dummy_past_residuals[:]

															with tempfile.TemporaryDirectory() as tmpdirname:
																					scheduler.save_config(__snake_case					)
																					_SCREAMING_SNAKE_CASE     				=  scheduler_class.from_pretrained(__snake_case					)
																					# copy over dummy past residuals
																					new_scheduler.set_timesteps(__snake_case					)

																					# copy over dummy past residual (must be after setting timesteps)
																					_SCREAMING_SNAKE_CASE     				=  dummy_past_residuals[:]

															_SCREAMING_SNAKE_CASE     				=  scheduler.step_prk(__snake_case				,     __snake_case				,     __snake_case				,     **__snake_case					).prev_sample
															_SCREAMING_SNAKE_CASE     				=  new_scheduler.step_prk(__snake_case				,     __snake_case				,     __snake_case				,     **__snake_case					).prev_sample

															assert torch.sum(torch.abs(output - new_output					)					) < 1E-5, "Scheduler outputs are not identical"

															_SCREAMING_SNAKE_CASE     				=  scheduler.step_plms(__snake_case				,     __snake_case				,     __snake_case				,     **__snake_case					).prev_sample
															_SCREAMING_SNAKE_CASE     				=  new_scheduler.step_plms(__snake_case				,     __snake_case				,     __snake_case				,     **__snake_case					).prev_sample

															assert torch.sum(torch.abs(output - new_output					)					) < 1E-5, "Scheduler outputs are not identical"







			def 							UpperCamelCase				(     self:  int				,     **UpperCAmelCase_:  Tuple					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.scheduler_classes[0]
									_SCREAMING_SNAKE_CASE     				=  self.get_scheduler_config(**__snake_case					)
									_SCREAMING_SNAKE_CASE     				=  scheduler_class(**__snake_case					)

									_SCREAMING_SNAKE_CASE     				=  10
									_SCREAMING_SNAKE_CASE     				=  self.dummy_model()
									_SCREAMING_SNAKE_CASE     				=  self.dummy_sample_deter
									scheduler.set_timesteps(__snake_case					)

									for i, t in enumerate(scheduler.prk_timesteps					):
															_SCREAMING_SNAKE_CASE     				=  model(__snake_case				,     __snake_case					)
															_SCREAMING_SNAKE_CASE     				=  scheduler.step_prk(__snake_case				,     __snake_case				,     __snake_case					).prev_sample

									for i, t in enumerate(scheduler.plms_timesteps					):
															_SCREAMING_SNAKE_CASE     				=  model(__snake_case				,     __snake_case					)
															_SCREAMING_SNAKE_CASE     				=  scheduler.step_plms(__snake_case				,     __snake_case				,     __snake_case					).prev_sample

									return sample







			def 							UpperCamelCase				(     self:  Optional[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  dict(self.forward_default_kwargs					)

									_SCREAMING_SNAKE_CASE     				=  kwargs.pop("""num_inference_steps"""				,     __snake_case					)

									for scheduler_class in self.scheduler_classes:
															_SCREAMING_SNAKE_CASE     				=  self.get_scheduler_config()
															_SCREAMING_SNAKE_CASE     				=  scheduler_class(**__snake_case					)

															_SCREAMING_SNAKE_CASE     				=  self.dummy_sample
															_SCREAMING_SNAKE_CASE     				=  0.1 * sample

															if num_inference_steps is not None and hasattr(__snake_case				,     """set_timesteps"""					):
																					scheduler.set_timesteps(__snake_case					)
															elif num_inference_steps is not None and not hasattr(__snake_case				,     """set_timesteps"""					):
																					_SCREAMING_SNAKE_CASE     				=  num_inference_steps

															# copy over dummy past residuals (must be done after set_timesteps)
															_SCREAMING_SNAKE_CASE     				=  [residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]
															_SCREAMING_SNAKE_CASE     				=  dummy_past_residuals[:]

															_SCREAMING_SNAKE_CASE     				=  scheduler.step_prk(__snake_case				,     0				,     __snake_case				,     **__snake_case					).prev_sample
															_SCREAMING_SNAKE_CASE     				=  scheduler.step_prk(__snake_case				,     1				,     __snake_case				,     **__snake_case					).prev_sample

															self.assertEqual(output_a.shape				,     sample.shape					)
															self.assertEqual(output_a.shape				,     output_a.shape					)

															_SCREAMING_SNAKE_CASE     				=  scheduler.step_plms(__snake_case				,     0				,     __snake_case				,     **__snake_case					).prev_sample
															_SCREAMING_SNAKE_CASE     				=  scheduler.step_plms(__snake_case				,     1				,     __snake_case				,     **__snake_case					).prev_sample

															self.assertEqual(output_a.shape				,     sample.shape					)
															self.assertEqual(output_a.shape				,     output_a.shape					)







			def 							UpperCamelCase				(     self:  Any					):







									'''simple docstring'''




									for timesteps in [100, 1_000]:
															self.check_over_configs(num_train_timesteps=__snake_case					)







			def 							UpperCamelCase				(     self:  List[str]					):







									'''simple docstring'''




									for steps_offset in [0, 1]:
															self.check_over_configs(steps_offset=__snake_case					)

									_SCREAMING_SNAKE_CASE     				=  self.scheduler_classes[0]
									_SCREAMING_SNAKE_CASE     				=  self.get_scheduler_config(steps_offset=1					)
									_SCREAMING_SNAKE_CASE     				=  scheduler_class(**__snake_case					)
									scheduler.set_timesteps(10					)
									assert torch.equal(
									    scheduler.timesteps				,     torch.LongTensor(
									        [901, 851, 851, 801, 801, 751, 751, 701, 701, 651, 651, 601, 601, 501, 401, 301, 201, 101, 1]					)				,     )







			def 							UpperCamelCase				(     self:  List[str]					):







									'''simple docstring'''




									for beta_start, beta_end in zip([0.00_01, 0.0_01]				,     [0.0_02, 0.02]					):
															self.check_over_configs(beta_start=__snake_case				,     beta_end=__snake_case					)







			def 							UpperCamelCase				(     self:  Any					):







									'''simple docstring'''




									for schedule in ["linear", "squaredcos_cap_v2"]:
															self.check_over_configs(beta_schedule=__snake_case					)







			def 							UpperCamelCase				(     self:  List[Any]					):







									'''simple docstring'''




									for prediction_type in ["epsilon", "v_prediction"]:
															self.check_over_configs(prediction_type=__snake_case					)







			def 							UpperCamelCase				(     self:  Any					):







									'''simple docstring'''




									for t in [1, 5, 10]:
															self.check_over_forward(time_step=__snake_case					)







			def 							UpperCamelCase				(     self:  Optional[int]					):







									'''simple docstring'''




									for t, num_inference_steps in zip([1, 5, 10]				,     [10, 50, 100]					):
															self.check_over_forward(num_inference_steps=__snake_case					)







			def 							UpperCamelCase				(     self:  Optional[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  27

									for scheduler_class in self.scheduler_classes:
															_SCREAMING_SNAKE_CASE     				=  self.dummy_sample
															_SCREAMING_SNAKE_CASE     				=  0.1 * sample

															_SCREAMING_SNAKE_CASE     				=  self.get_scheduler_config()
															_SCREAMING_SNAKE_CASE     				=  scheduler_class(**__snake_case					)

															scheduler.set_timesteps(__snake_case					)

															# before power of 3 fix, would error on first step, so we only need to do two
															for i, t in enumerate(scheduler.prk_timesteps[:2]					):
																					_SCREAMING_SNAKE_CASE     				=  scheduler.step_prk(__snake_case				,     __snake_case				,     __snake_case					).prev_sample







			def 							UpperCamelCase				(     self:  Union[str, Any]					):







									'''simple docstring'''




									with self.assertRaises(__snake_case					):
															_SCREAMING_SNAKE_CASE     				=  self.scheduler_classes[0]
															_SCREAMING_SNAKE_CASE     				=  self.get_scheduler_config()
															_SCREAMING_SNAKE_CASE     				=  scheduler_class(**__snake_case					)

															scheduler.step_plms(self.dummy_sample				,     1				,     self.dummy_sample					).prev_sample







			def 							UpperCamelCase				(     self:  Any					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.full_loop()
									_SCREAMING_SNAKE_CASE     				=  torch.sum(torch.abs(__snake_case					)					)
									_SCREAMING_SNAKE_CASE     				=  torch.mean(torch.abs(__snake_case					)					)

									assert abs(result_sum.item() - 198.1_318					) < 1E-2
									assert abs(result_mean.item() - 0.25_80					) < 1E-3







			def 							UpperCamelCase				(     self:  Tuple					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.full_loop(prediction_type="""v_prediction"""					)
									_SCREAMING_SNAKE_CASE     				=  torch.sum(torch.abs(__snake_case					)					)
									_SCREAMING_SNAKE_CASE     				=  torch.mean(torch.abs(__snake_case					)					)

									assert abs(result_sum.item() - 67.39_86					) < 1E-2
									assert abs(result_mean.item() - 0.08_78					) < 1E-3







			def 							UpperCamelCase				(     self:  int					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.full_loop(set_alpha_to_one=__snake_case				,     beta_start=0.01					)
									_SCREAMING_SNAKE_CASE     				=  torch.sum(torch.abs(__snake_case					)					)
									_SCREAMING_SNAKE_CASE     				=  torch.mean(torch.abs(__snake_case					)					)

									assert abs(result_sum.item() - 230.0_399					) < 1E-2
									assert abs(result_mean.item() - 0.29_95					) < 1E-3







			def 							UpperCamelCase				(     self:  Tuple					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.full_loop(set_alpha_to_one=__snake_case				,     beta_start=0.01					)
									_SCREAMING_SNAKE_CASE     				=  torch.sum(torch.abs(__snake_case					)					)
									_SCREAMING_SNAKE_CASE     				=  torch.mean(torch.abs(__snake_case					)					)

									assert abs(result_sum.item() - 186.9_482					) < 1E-2
									assert abs(result_mean.item() - 0.24_34					) < 1E-3









# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import re

from ..models.auto import AutoProcessor
from ..models.vision_encoder_decoder import VisionEncoderDecoderModel
from ..utils import is_vision_available
from .base import PipelineTool


if is_vision_available():
		from PIL import Image



class 				__UpperCAmelCase		(_UpperCAmelCase   ):
			__snake_case    :				Dict									=							"naver-clova-ix/donut-base-finetuned-docvqa"
			__snake_case    :				Dict									=							(
			    "This is a tool that answers a question about an document (pdf). It takes an input named `document` which "
			    "should be the document containing the information, as well as a `question` that is the question about the "
			    "document. It returns a text that contains the answer to the question."
			)
			__snake_case    :				Dict									=							"document_qa"
			__snake_case    :				Any									=							AutoProcessor
			__snake_case    :				int									=							VisionEncoderDecoderModel

			__snake_case    :				Union[str, Any]									=							["image", "text"]
			__snake_case    :				Optional[int]									=							["text"]







			def __init__(     self:  List[Any]				,     *UpperCAmelCase_:  Optional[int]				,     **UpperCAmelCase_:  List[str]					):







									'''simple docstring'''




									if not is_vision_available():
															raise ValueError("""Pillow must be installed to use the DocumentQuestionAnsweringTool."""					)

									super().__init__(*UpperCAmelCase_				,     **UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  Optional[int]				,     UpperCAmelCase_:  "Image"				,     UpperCAmelCase_:  str					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  """<s_docvqa><s_question>{user_input}</s_question><s_answer>"""
									_SCREAMING_SNAKE_CASE     				=  task_prompt.replace("""{user_input}"""				,     UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  self.pre_processor.tokenizer(
									    UpperCAmelCase_				,     add_special_tokens=UpperCAmelCase_				,     return_tensors="""pt"""					).input_ids
									_SCREAMING_SNAKE_CASE     				=  self.pre_processor(UpperCAmelCase_				,     return_tensors="""pt"""					).pixel_values

									return {"decoder_input_ids": decoder_input_ids, "pixel_values": pixel_values}







			def 							UpperCamelCase				(     self:  Dict				,     UpperCAmelCase_:  str					):







									'''simple docstring'''




									return self.model.generate(
									    inputs["""pixel_values"""].to(self.device					)				,     decoder_input_ids=inputs["""decoder_input_ids"""].to(self.device					)				,     max_length=self.model.decoder.config.max_position_embeddings				,     early_stopping=UpperCAmelCase_				,     pad_token_id=self.pre_processor.tokenizer.pad_token_id				,     eos_token_id=self.pre_processor.tokenizer.eos_token_id				,     use_cache=UpperCAmelCase_				,     num_beams=1				,     bad_words_ids=[[self.pre_processor.tokenizer.unk_token_id]]				,     return_dict_in_generate=UpperCAmelCase_				,     ).sequences







			def 							UpperCamelCase				(     self:  Dict				,     UpperCAmelCase_:  Tuple					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.pre_processor.batch_decode(UpperCAmelCase_					)[0]
									_SCREAMING_SNAKE_CASE     				=  sequence.replace(self.pre_processor.tokenizer.eos_token				,     """"""					)
									_SCREAMING_SNAKE_CASE     				=  sequence.replace(self.pre_processor.tokenizer.pad_token				,     """"""					)
									_SCREAMING_SNAKE_CASE     				=  re.sub(R"""<.*?>"""				,     """"""				,     UpperCAmelCase_				,     count=1					).strip()  # remove first task start token
									_SCREAMING_SNAKE_CASE     				=  self.pre_processor.tokenajson(UpperCAmelCase_					)

									return sequence["answer"]

