


'''simple docstring'''


# 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 		lowerCAmelCase__	(	UpperCAmelCase__		):
			lowerCAmelCase       :		str				      =   "naver-clova-ix/donut-base-finetuned-docvqa"
			lowerCAmelCase       :		List[Any]				      =   (
			    "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."
			)
			lowerCAmelCase       :		Any				      =   "document_qa"
			lowerCAmelCase       :		Union[str, Any]				      =   AutoProcessor
			lowerCAmelCase       :		List[str]				      =   VisionEncoderDecoderModel

			lowerCAmelCase       :		int				      =   ["image", "text"]
			lowerCAmelCase       :		List[Any]				      =   ["text"]




			def __init__(						self     :   Union[str, Any]		,     *lowerCamelCase__     :   List[Any]		,     **lowerCamelCase__     :   Optional[Any]       )      ->List[Any]:



					'''simple docstring'''







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

					super().__init__(*lowerCamelCase__		,     **lowerCamelCase__       )




			def 							lowerCAmelCase__			(						self     :   Any		,     lowerCamelCase__     :   "Image"		,     lowerCamelCase__     :   str       )      ->str:



					'''simple docstring'''







					_UpperCAmelCase						:				List[Any]	    =					"<s_docvqa><s_question>{user_input}</s_question><s_answer>"
					_UpperCAmelCase						:				int	    =					task_prompt.replace("{user_input}"		,     lowerCamelCase__       )
					_UpperCAmelCase						:				List[Any]	    =					self.pre_processor.tokenizer(
					    lowerCamelCase__		,     add_special_tokens=lowerCamelCase__		,     return_tensors="pt"       ).input_ids
					_UpperCAmelCase						:				int	    =					self.pre_processor(lowerCamelCase__		,     return_tensors="pt"       ).pixel_values

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




			def 							lowerCAmelCase__			(						self     :   Dict		,     lowerCamelCase__     :   Tuple       )      ->Union[str, Any]:



					'''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=lowerCamelCase__		,     pad_token_id=self.pre_processor.tokenizer.pad_token_id		,     eos_token_id=self.pre_processor.tokenizer.eos_token_id		,     use_cache=lowerCamelCase__		,     num_beams=1		,     bad_words_ids=[[self.pre_processor.tokenizer.unk_token_id]]		,     return_dict_in_generate=lowerCamelCase__		,     ).sequences






			def 							lowerCAmelCase__			(						self     :   Tuple		,     lowerCamelCase__     :   Optional[int]       )      ->int:



					'''simple docstring'''







					_UpperCAmelCase						:				Optional[int]	    =					self.pre_processor.batch_decode(lowerCamelCase__       )[0]
					_UpperCAmelCase						:				List[Any]	    =					sequence.replace(self.pre_processor.tokenizer.eos_token		,     ""       )
					_UpperCAmelCase						:				Dict	    =					sequence.replace(self.pre_processor.tokenizer.pad_token		,     ""       )
					_UpperCAmelCase						:				Any	    =					re.sub(R"<.*?>"		,     ""		,     lowerCamelCase__		,     count=1       ).strip()  # remove first task start token
					_UpperCAmelCase						:				List[str]	    =					self.pre_processor.tokenajson(lowerCamelCase__       )

					return sequence["answer"]







'''simple docstring'''


#
# This a `torch.distributed` diagnostics script that checks that all GPUs in the cluster (one or
# many nodes) can talk to each other via nccl and allocate gpu memory.
#
# To run first adjust the number of processes and nodes:
#
# python -m torch.distributed.run --nproc_per_node 2 --nnodes 1 torch-distributed-gpu-test.py
#
# You may need to add --master_addr $MASTER_ADDR --master_port $MASTER_PORT if using a custom addr:port
#
# You can also use the rdzv API: --rdzv_endpoint $MASTER_ADDR:$MASTER_PORT --rdzv_backend c10d
#
# use torch.distributed.launch instead of torch.distributed.run for torch < 1.9
#
# If you get a hanging in `barrier` calls you have some network issues, you may try to debug this with:
#
# NCCL_DEBUG=INFO python -m torch.distributed.run --nproc_per_node 2 --nnodes 1 torch-distributed-gpu-test.py
#
# which should tell you what's going on behind the scenes.
#
#
# This script can be run via `srun` in the SLURM environment as well. Here is a SLURM script that
# runs on 2 nodes of 4 gpus per node:
#
# #SBATCH --job-name=test-nodes        # name
# #SBATCH --nodes=2                    # nodes
# #SBATCH --ntasks-per-node=1          # crucial - only 1 task per dist per node!
# #SBATCH --cpus-per-task=10           # number of cores per tasks
# #SBATCH --gres=gpu:4                 # number of gpus
# #SBATCH --time 0:05:00               # maximum execution time (HH:MM:SS)
# #SBATCH --output=%x-%j.out           # output file name
#
# GPUS_PER_NODE=4
# MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
# MASTER_PORT=6000
#
# srun --jobid $SLURM_JOBID bash -c 'python -m torch.distributed.run \
# --nproc_per_node $GPUS_PER_NODE --nnodes $SLURM_NNODES --node_rank $SLURM_PROCID \
# --master_addr $MASTER_ADDR --master_port $MASTER_PORT \
# torch-distributed-gpu-test.py'
#

import fcntl
import os
import socket

import torch
import torch.distributed as dist






def        __lowerCAmelCase       (*__lowerCAmelCase      ):
		with open(__lowerCAmelCase					,       "r"      ) as fh:
				fcntl.flock(__lowerCAmelCase					,       fcntl.LOCK_EX      )
				try:
						print(*__lowerCAmelCase      )
				finally:
						fcntl.flock(__lowerCAmelCase					,       fcntl.LOCK_UN      )


lowerCamelCase__					=   int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(local_rank)
lowerCamelCase__					=   torch.device('cuda', local_rank)
lowerCamelCase__					=   socket.gethostname()

lowerCamelCase__					=   F'''[{hostname}-{local_rank}]'''

try:
						# test distributed
						dist.init_process_group('nccl')
						dist.all_reduce(torch.ones(1).to(device), op=dist.ReduceOp.SUM)
						dist.barrier()

						# test cuda is available and can allocate memory
						torch.cuda.is_available()
						torch.ones(1).cuda(local_rank)

						# global rank
						lowerCamelCase__					=   dist.get_rank()
						lowerCamelCase__					=   dist.get_world_size()

						printflock(F'''{gpu} is OK (global rank: {rank}/{world_size})''')

						dist.barrier()
						if rank == 0:
												printflock(F'''pt={torch.__version__}, cuda={torch.version.cuda}, nccl={torch.cuda.nccl.version()}''')

except Exception:
						printflock(F'''{gpu} is broken''')
						raise



