



import argparse
import collections
from pathlib import Path

import torch
from torch.serialization import default_restore_location

from transformers import BertConfig, DPRConfig, DPRContextEncoder, DPRQuestionEncoder, DPRReader


UpperCamelCase__			=    collections.namedtuple(
    "CheckpointState", ["model_dict", "optimizer_dict", "scheduler_dict", "offset", "epoch", "encoder_params"]
)






def 			UpperCAmelCase__     (     lowercase__ )			->   CheckpointState:
						print(F"Reading saved model from {model_file}" )
						__lowercase											=					torch.load(lowercase__    ,							map_location=lambda lowercase__    ,							lowercase__ : default_restore_location(lowercase__    ,							"""cpu""" ) )
						return CheckpointState(**lowercase__ )
class       _lowerCAmelCase :






			"""simple docstring"""





			def __init__(					self						:							List[str]      ,  lowercase						:							Path )    ->       Dict:
									"""simple docstring"""







									__lowercase											=					src_file





			def 		snake_case__      (					self						:							List[str] )    ->       List[Any]:
									"""simple docstring"""







									raise NotImplementedError







			@staticmethod
			def 		snake_case__      (					lowercase						:							str      ,  *lowercase						:							str      ,  **lowercase						:							Optional[Any] )    ->       "DPRState":
									"""simple docstring"""







									if comp_type.startswith("""c""" ):
															return DPRContextEncoderState(*lowercase      ,  **lowercase )
									if comp_type.startswith("""q""" ):
															return DPRQuestionEncoderState(*lowercase      ,  **lowercase )
									if comp_type.startswith("""r""" ):
															return DPRReaderState(*lowercase      ,  **lowercase )
									else:
															raise ValueError("""Component type must be either 'ctx_encoder', 'question_encoder' or 'reader'.""" )





class       _lowerCAmelCase (			_UpperCAmelCase      ):






			"""simple docstring"""





			def 		snake_case__      (					self						:							Optional[Any] )    ->       str:
									"""simple docstring"""







									__lowercase											=					DPRContextEncoder(DPRConfig(**BertConfig.get_config_dict("""bert-base-uncased""" )[0] ) )
									print(F"Loading DPR biencoder from {self.src_file}" )
									__lowercase											=					load_states_from_checkpoint(self.src_file )
									__lowercase						,      __lowercase											=					model.ctx_encoder, """ctx_model."""
									# Fix changes from https://github.com/huggingface/transformers/commit/614fef1691edb806de976756d4948ecbcd0c0ca3
									__lowercase											=					{"""bert_model.embeddings.position_ids""": model.ctx_encoder.bert_model.embeddings.position_ids}
									for key, value in saved_state.model_dict.items():
															if key.startswith(lowercase ):
																					__lowercase											=					key[len(lowercase ) :]
																					if not key.startswith("""encode_proj.""" ):
																											__lowercase											=					"""bert_model.""" + key
																					__lowercase											=					value
									encoder.load_state_dict(lowercase )
									return model





class       _lowerCAmelCase (			_UpperCAmelCase      ):






			"""simple docstring"""





			def 		snake_case__      (					self						:							str )    ->       Optional[Any]:
									"""simple docstring"""







									__lowercase											=					DPRQuestionEncoder(DPRConfig(**BertConfig.get_config_dict("""bert-base-uncased""" )[0] ) )
									print(F"Loading DPR biencoder from {self.src_file}" )
									__lowercase											=					load_states_from_checkpoint(self.src_file )
									__lowercase						,      __lowercase											=					model.question_encoder, """question_model."""
									# Fix changes from https://github.com/huggingface/transformers/commit/614fef1691edb806de976756d4948ecbcd0c0ca3
									__lowercase											=					{"""bert_model.embeddings.position_ids""": model.question_encoder.bert_model.embeddings.position_ids}
									for key, value in saved_state.model_dict.items():
															if key.startswith(lowercase ):
																					__lowercase											=					key[len(lowercase ) :]
																					if not key.startswith("""encode_proj.""" ):
																											__lowercase											=					"""bert_model.""" + key
																					__lowercase											=					value
									encoder.load_state_dict(lowercase )
									return model





class       _lowerCAmelCase (			_UpperCAmelCase      ):






			"""simple docstring"""





			def 		snake_case__      (					self						:							Tuple )    ->       List[Any]:
									"""simple docstring"""







									__lowercase											=					DPRReader(DPRConfig(**BertConfig.get_config_dict("""bert-base-uncased""" )[0] ) )
									print(F"Loading DPR reader from {self.src_file}" )
									__lowercase											=					load_states_from_checkpoint(self.src_file )
									# Fix changes from https://github.com/huggingface/transformers/commit/614fef1691edb806de976756d4948ecbcd0c0ca3
									__lowercase											=					{
									    """encoder.bert_model.embeddings.position_ids""": model.span_predictor.encoder.bert_model.embeddings.position_ids
									}
									for key, value in saved_state.model_dict.items():
															if key.startswith("""encoder.""" ) and not key.startswith("""encoder.encode_proj""" ):
																					__lowercase											=					"""encoder.bert_model.""" + key[len("""encoder.""" ) :]
															__lowercase											=					value
									model.span_predictor.load_state_dict(lowercase )
									return model






def 			UpperCAmelCase__     (     lowercase__    ,							lowercase__    ,							lowercase__ )			->   List[Any]:
						__lowercase											=					Path(lowercase__ )
						dest_dir.mkdir(exist_ok=lowercase__ )

						__lowercase											=					DPRState.from_type(lowercase__    ,							src_file=lowercase__ )
						__lowercase											=					dpr_state.load_dpr_model()
						model.save_pretrained(lowercase__ )
						model.from_pretrained(lowercase__ )  # sanity check


if __name__ == "__main__":
			UpperCamelCase__			=    argparse.ArgumentParser()
			# Required parameters
			parser.add_argument(
			    "--type", type=str, help="Type of the component to convert: 'ctx_encoder', 'question_encoder' or 'reader'."
			)
			parser.add_argument(
			    "--src",
			    type=str,
			    help=(
			        "Path to the dpr checkpoint file. They can be downloaded from the official DPR repo"
			        " https://github.com/facebookresearch/DPR. Note that in the official repo, both encoders are stored in the"
			        " 'retriever' checkpoints."
			    ),
			)
			parser.add_argument("--dest", type=str, default=None, help="Path to the output PyTorch model directory.")
			UpperCamelCase__			=    parser.parse_args()

			UpperCamelCase__			=    Path(args.src)
			UpperCamelCase__			=    F"""converted-{src_file.name}""" if args.dest is None else args.dest
			UpperCamelCase__			=    Path(dest_dir)
			assert src_file.exists()
			assert (
			    args.type is not None
			), "Please specify the component type of the DPR model to convert: 'ctx_encoder', 'question_encoder' or 'reader'."
			convert(args.type, src_file, dest_dir)










# 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.
from typing import TYPE_CHECKING

from ..utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


UpperCamelCase__			=    {
    "agents": ["Agent", "AzureOpenAiAgent", "HfAgent", "LocalAgent", "OpenAiAgent"],
    "base": ["PipelineTool", "RemoteTool", "Tool", "launch_gradio_demo", "load_tool"],
}

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			UpperCamelCase__			=    ["DocumentQuestionAnsweringTool"]
			UpperCamelCase__			=    ["ImageCaptioningTool"]
			UpperCamelCase__			=    ["ImageQuestionAnsweringTool"]
			UpperCamelCase__			=    ["ImageSegmentationTool"]
			UpperCamelCase__			=    ["SpeechToTextTool"]
			UpperCamelCase__			=    ["TextClassificationTool"]
			UpperCamelCase__			=    ["TextQuestionAnsweringTool"]
			UpperCamelCase__			=    ["TextSummarizationTool"]
			UpperCamelCase__			=    ["TextToSpeechTool"]
			UpperCamelCase__			=    ["TranslationTool"]

if TYPE_CHECKING:
			from .agents import Agent, AzureOpenAiAgent, HfAgent, LocalAgent, OpenAiAgent
			from .base import PipelineTool, RemoteTool, Tool, launch_gradio_demo, load_tool

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .document_question_answering import DocumentQuestionAnsweringTool
						from .image_captioning import ImageCaptioningTool
						from .image_question_answering import ImageQuestionAnsweringTool
						from .image_segmentation import ImageSegmentationTool
						from .speech_to_text import SpeechToTextTool
						from .text_classification import TextClassificationTool
						from .text_question_answering import TextQuestionAnsweringTool
						from .text_summarization import TextSummarizationTool
						from .text_to_speech import TextToSpeechTool
						from .translation import TranslationTool
else:
			import sys

			UpperCamelCase__			=    _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)





