


import unittest

from transformers import load_tool

from .test_tools_common import ToolTesterMixin


__lowerCamelCase   	=  '''
Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf originally as a company that developed a chatbot app targeted at teenagers.[2] After open-sourcing the model behind the chatbot, the company pivoted to focus on being a platform for machine learning.

In March 2021, Hugging Face raised $40 million in a Series B funding round.[3]

On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model.[4] In 2022, the workshop concluded with the announcement of BLOOM, a multilingual large language model with 176 billion parameters.[5]
'''


class 	a__						(    unittest.TestCase    ,							lowerCAmelCase_     ):

       def     UpperCAmelCase( self   :  Optional[int]     ):
            a_     :  List[Any]				      =     load_tool("""summarization"""     )
            self.tool.setup()
            a_     :  Any				      =     load_tool("""summarization"""   ,				remote=lowerCamelCase_     )

       def     UpperCAmelCase( self   :  Any     ):
            a_     :  int				      =     self.tool(lowerCamelCase_     )
            self.assertEqual(
                lowerCamelCase_   ,				"""Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM."""   ,				)

       def     UpperCAmelCase( self   :  Union[str, Any]     ):
            a_     :  str				      =     self.remote_tool(lowerCamelCase_     )
            self.assertEqual(
                lowerCamelCase_   ,				"""Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM."""   ,				)

       def     UpperCAmelCase( self   :  Tuple     ):
            a_     :  Any				      =     self.tool(text=lowerCamelCase_     )
            self.assertEqual(
                lowerCamelCase_   ,				"""Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM."""   ,				)

       def     UpperCAmelCase( self   :  Optional[Any]     ):
            a_     :  Optional[Any]				      =     self.remote_tool(text=lowerCamelCase_     )
            self.assertEqual(
                lowerCamelCase_   ,				"""Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM."""   ,				)




"""simple docstring"""


from typing import Dict, Tuple, overload

import torch
import torch.types
from torch import nn

from . import residue_constants as rc
from .rigid_utils import Rigid, Rotation
from .tensor_utils import batched_gather







@overload
def    UpperCAmelCase       (						UpperCAmelCase		,   UpperCAmelCase		,   UpperCAmelCase   )							-> torch.Tensor:
     ...







@overload
def    UpperCAmelCase       (						UpperCAmelCase		,   UpperCAmelCase		,   UpperCAmelCase   )							-> Tuple[torch.Tensor, torch.Tensor]:
     ...







def    UpperCAmelCase       (						UpperCAmelCase		,   UpperCAmelCase		,   UpperCAmelCase   )							-> Dict:
     snake_case_   			=   aatype == rc.restype_order['G']
     snake_case_   			=   rc.atom_order['CA']
     snake_case_   			=   rc.atom_order['CB']
     snake_case_   			=   torch.where(
         is_gly[..., None].expand(*((-1,) * len(is_gly.shape   ))		,   3   )		,   all_atom_positions[..., ca_idx, :]		,   all_atom_positions[..., cb_idx, :]		,   )

     if all_atom_masks is not None:
          snake_case_   			=   torch.where(
              UpperCAmelCase		,   all_atom_masks[..., ca_idx]		,   all_atom_masks[..., cb_idx]		,   )
          return pseudo_beta, pseudo_beta_mask
     else:
          return pseudo_beta







def    UpperCAmelCase       (						UpperCAmelCase		,   UpperCAmelCase   )							-> torch.Tensor:
     snake_case_   			=   batched_gather(
         UpperCAmelCase		,   batch['residx_atom37_to_atom14']		,   dim=-2		,   no_batch_dims=len(atomaa.shape[:-2]   )		,   )

     snake_case_   			=   atomaa_data * batch['atom37_atom_exists'][..., None]

     return atomaa_data







def    UpperCAmelCase       (						UpperCAmelCase   )							-> torch.Tensor:
     snake_case_   			=   template_feats['template_aatype']
     snake_case_   			=   template_feats['template_torsion_angles_sin_cos']
     snake_case_   			=   template_feats['template_alt_torsion_angles_sin_cos']
     snake_case_   			=   template_feats['template_torsion_angles_mask']
     snake_case_   			=   torch.cat(
         [
             nn.functional.one_hot(UpperCAmelCase		,   22   ),
             torsion_angles_sin_cos.reshape(*torsion_angles_sin_cos.shape[:-2]		,   14   ),
             alt_torsion_angles_sin_cos.reshape(*alt_torsion_angles_sin_cos.shape[:-2]		,   14   ),
             torsion_angles_mask,
         ]		,   dim=-1		,   )

     return template_angle_feat







def    UpperCAmelCase       (						UpperCAmelCase		,   UpperCAmelCase		,   UpperCAmelCase		,   UpperCAmelCase		,   UpperCAmelCase = False		,   UpperCAmelCase = 1e-20		,   UpperCAmelCase = 1e8		,   )							-> torch.Tensor:
     snake_case_   			=   batch['template_pseudo_beta_mask']
     snake_case_   			=   template_mask[..., None] * template_mask[..., None, :]

     # Compute distogram (this seems to differ slightly from Alg. 5)
     snake_case_   			=   batch['template_pseudo_beta']
     snake_case_   			=   torch.sum((tpb[..., None, :] - tpb[..., None, :, :]) ** 2		,   dim=-1		,   keepdim=UpperCAmelCase   )
     snake_case_   			=   torch.linspace(UpperCAmelCase		,   UpperCAmelCase		,   UpperCAmelCase		,   device=tpb.device   ) ** 2
     snake_case_   			=   torch.cat([lower[1:], lower.new_tensor([inf]   )]		,   dim=-1   )
     snake_case_   			=   ((dgram > lower) * (dgram < upper)).type(dgram.dtype   )

     snake_case_   			=   [dgram, template_mask_ad[..., None]]

     snake_case_   			=   nn.functional.one_hot(
         batch['template_aatype']		,   rc.restype_num + 2		,   )

     snake_case_   			=   batch['template_aatype'].shape[-1]
     to_concat.append(aatype_one_hot[..., None, :, :].expand(*aatype_one_hot.shape[:-2]		,   UpperCAmelCase		,   -1		,   -1   )   )
     to_concat.append(aatype_one_hot[..., None, :].expand(*aatype_one_hot.shape[:-2]		,   -1		,   UpperCAmelCase		,   -1   )   )

     snake_case_			,      snake_case_			,      snake_case_   			=   [rc.atom_order[a] for a in ['N', 'CA', 'C']]
     snake_case_   			=   Rigid.make_transform_from_reference(
         n_xyz=batch['template_all_atom_positions'][..., n, :]		,   ca_xyz=batch['template_all_atom_positions'][..., ca, :]		,   c_xyz=batch['template_all_atom_positions'][..., c, :]		,   eps=UpperCAmelCase		,   )
     snake_case_   			=   rigids.get_trans()[..., None, :, :]
     snake_case_   			=   rigids[..., None].invert_apply(UpperCAmelCase   )

     snake_case_   			=   torch.rsqrt(eps + torch.sum(rigid_vec**2		,   dim=-1   )   )

     snake_case_   			=   batch['template_all_atom_mask']
     snake_case_   			=   t_aa_masks[..., n] * t_aa_masks[..., ca] * t_aa_masks[..., c]
     snake_case_   			=   template_mask[..., None] * template_mask[..., None, :]

     snake_case_   			=   inv_distance_scalar * template_mask_ad
     snake_case_   			=   rigid_vec * inv_distance_scalar[..., None]

     if not use_unit_vector:
          snake_case_   			=   unit_vector * 0.0

     to_concat.extend(torch.unbind(unit_vector[..., None, :]		,   dim=-1   )   )
     to_concat.append(template_mask_ad[..., None]   )

     snake_case_   			=   torch.cat(UpperCAmelCase		,   dim=-1   )
     snake_case_   			=   act * template_mask_ad[..., None]

     return act







def    UpperCAmelCase       (						UpperCAmelCase   )							-> torch.Tensor:
     snake_case_   			=   nn.functional.one_hot(batch['extra_msa']		,   23   )
     snake_case_   			=   [
         msa_ahot,
         batch['extra_has_deletion'].unsqueeze(-1   ),
         batch['extra_deletion_value'].unsqueeze(-1   ),
     ]
     return torch.cat(UpperCAmelCase		,   dim=-1   )







def    UpperCAmelCase       (						UpperCAmelCase		,   UpperCAmelCase		,   UpperCAmelCase		,   UpperCAmelCase		,   )							-> Rigid:
     # [*, N, 8, 4, 4]
     snake_case_   			=   rrgdf[aatype, ...]

     # [*, N, 8] transformations, i.e.
     #   One [*, N, 8, 3, 3] rotation matrix and
     #   One [*, N, 8, 3]    translation matrix
     snake_case_   			=   r.from_tensor_axa(UpperCAmelCase   )

     snake_case_   			=   alpha.new_zeros((*((1,) * len(alpha.shape[:-1]   )), 2)   )
     snake_case_   			=   1

     # [*, N, 8, 2]
     snake_case_   			=   torch.cat([bb_rot.expand(*alpha.shape[:-2]		,   -1		,   -1   ), alpha]		,   dim=-2   )

     # [*, N, 8, 3, 3]
     # Produces rotation matrices of the form:
     # [
     #   [1, 0  , 0  ],
     #   [0, a_2,-a_1],
     #   [0, a_1, a_2]
     # ]
     # This follows the original code rather than the supplement, which uses
     # different indices.

     snake_case_   			=   alpha.new_zeros(default_r.get_rots().get_rot_mats().shape   )
     snake_case_   			=   1
     snake_case_   			=   alpha[..., 1]
     snake_case_   			=   -alpha[..., 0]
     snake_case_   			=   alpha

     snake_case_   			=   default_r.compose(Rigid(Rotation(rot_mats=UpperCAmelCase   )		,   UpperCAmelCase   )   )

     snake_case_   			=   all_frames[..., 5]
     snake_case_   			=   all_frames[..., 6]
     snake_case_   			=   all_frames[..., 7]

     snake_case_   			=   all_frames[..., 4]
     snake_case_   			=   chia_frame_to_bb.compose(UpperCAmelCase   )
     snake_case_   			=   chia_frame_to_bb.compose(UpperCAmelCase   )
     snake_case_   			=   chia_frame_to_bb.compose(UpperCAmelCase   )

     snake_case_   			=   Rigid.cat(
         [
             all_frames[..., :5],
             chia_frame_to_bb.unsqueeze(-1   ),
             chia_frame_to_bb.unsqueeze(-1   ),
             chia_frame_to_bb.unsqueeze(-1   ),
         ]		,   dim=-1		,   )

     snake_case_   			=   r[..., None].compose(UpperCAmelCase   )

     return all_frames_to_global







def    UpperCAmelCase       (						UpperCAmelCase		,   UpperCAmelCase		,   UpperCAmelCase		,   UpperCAmelCase		,   UpperCAmelCase		,   UpperCAmelCase		,   )							-> torch.Tensor:
     # [*, N, 14]
     snake_case_   			=   group_idx[aatype, ...]

     # [*, N, 14, 8]
     snake_case_   			=   nn.functional.one_hot(
         UpperCAmelCase		,   num_classes=default_frames.shape[-3]		,   )

     # [*, N, 14, 8]
     snake_case_   			=   r[..., None, :] * group_mask_one_hot

     # [*, N, 14]
     snake_case_   			=   t_atoms_to_global.map_tensor_fn(lambda UpperCAmelCase   : torch.sum(UpperCAmelCase		,   dim=-1   )   )

     # [*, N, 14, 1]
     snake_case_   			=   atom_mask[aatype, ...].unsqueeze(-1   )

     # [*, N, 14, 3]
     snake_case_   			=   lit_positions[aatype, ...]
     snake_case_   			=   t_atoms_to_global.apply(UpperCAmelCase   )
     snake_case_   			=   pred_positions * atom_mask

     return pred_positions

