


"""simple docstring"""







import dataclasses
import re
import string
from typing import Any, Dict, Iterator, List, Mapping, Optional, Sequence, Tuple

import numpy as np

from . import residue_constants


a__				:							Optional[Any]						=     Mapping[str, np.ndarray]
a__				:							Optional[Any]						=     Mapping[str, Any]  # Is a nested dict.
a__				:							Optional[int]						=     0.01




@dataclasses.dataclass(frozen=UpperCamelCase)
class  UpperCamelCase_	:

	"""simple docstring"""





	snake_case__				:  np.ndarray  # [num_res, num_atom_type, 3]

	# Amino-acid type for each residue represented as an integer between 0 and
	# 20, where 20 is 'X'.
	snake_case__				:  np.ndarray  # [num_res]

	# Binary float mask to indicate presence of a particular atom. 1.0 if an atom
	# is present and 0.0 if not. This should be used for loss masking.
	snake_case__				:  np.ndarray  # [num_res, num_atom_type]

	# Residue index as used in PDB. It is not necessarily continuous or 0-indexed.
	snake_case__				:  np.ndarray  # [num_res]

	# B-factors, or temperature factors, of each residue (in sq. angstroms units),
	# representing the displacement of the residue from its ground truth mean
	# value.
	snake_case__				:  np.ndarray  # [num_res, num_atom_type]

	# Chain indices for multi-chain predictions
	snake_case__				:  Optional[np.ndarray]    					=			None

	# Optional remark about the protein. Included as a comment in output PDB
	# files
	snake_case__				:  Optional[str]    					=			None

	# Templates used to generate this protein (prediction-only)
	snake_case__				:  Optional[Sequence[str]]    					=			None

	# Chain corresponding to each parent
	snake_case__				:  Optional[Sequence[int]]    					=			None




def    UpperCAmelCase__       (lowerCAmelCase_   ):




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE						   =				R"(\[[A-Z]+\]\n)"
					__SCREAMING_SNAKE_CASE						   =				[tag.strip() for tag in re.split(lowerCAmelCase_  ,						lowerCAmelCase_   ) if len(lowerCAmelCase_   ) > 0]
					__SCREAMING_SNAKE_CASE						   =				zip(tags[0::2]  ,						[l.split("\n"   ) for l in tags[1::2]]   )

					__SCREAMING_SNAKE_CASE						   =				["N", "CA", "C"]
					__SCREAMING_SNAKE_CASE						   =				None
					__SCREAMING_SNAKE_CASE						   =				None
					__SCREAMING_SNAKE_CASE						   =				None
					for g in groups:
										if "[PRIMARY]" == g[0]:
															__SCREAMING_SNAKE_CASE						   =				g[1][0].strip()
															for i in range(len(lowerCAmelCase_   )   ):
																				if seq[i] not in residue_constants.restypes:
																									__SCREAMING_SNAKE_CASE						   =				"X"  # FIXME: strings are immutable
															__SCREAMING_SNAKE_CASE						   =				np.array(
															    [residue_constants.restype_order.get(lowerCAmelCase_  ,						residue_constants.restype_num   ) for res_symbol in seq]   )
										elif "[TERTIARY]" == g[0]:
															__SCREAMING_SNAKE_CASE						   =				[]
															for axis in range(3   ):
																				tertiary.append(list(map(lowerCAmelCase_  ,						g[1][axis].split()   )   )   )
															__SCREAMING_SNAKE_CASE						   =				np.array(lowerCAmelCase_   )
															__SCREAMING_SNAKE_CASE						   =				np.zeros((len(tertiary[0]   ) // 3, residue_constants.atom_type_num, 3)   ).astype(np.floataa   )
															for i, atom in enumerate(lowerCAmelCase_   ):
																				__SCREAMING_SNAKE_CASE						   =				np.transpose(tertiary_np[:, i::3]   )
															atom_positions *= PICO_TO_ANGSTROM
										elif "[MASK]" == g[0]:
															__SCREAMING_SNAKE_CASE						   =				np.array(list(map({"-": 0, "+": 1}.get  ,						g[1][0].strip()   )   )   )
															__SCREAMING_SNAKE_CASE						   =				np.zeros(
															    (
															        len(lowerCAmelCase_   ),
															        residue_constants.atom_type_num,
															    )   ).astype(np.floataa   )
															for i, atom in enumerate(lowerCAmelCase_   ):
																				__SCREAMING_SNAKE_CASE						   =				1
															atom_mask *= mask[..., None]

					assert aatype is not None

					return Protein(
					    atom_positions=lowerCAmelCase_  ,						atom_mask=lowerCAmelCase_  ,						aatype=lowerCAmelCase_  ,						residue_index=np.arange(len(lowerCAmelCase_   )   )  ,						b_factors=lowerCAmelCase_  ,						)




def    UpperCAmelCase__       (lowerCAmelCase_  ,						lowerCAmelCase_ = 0   ):




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE						   =				[]

					__SCREAMING_SNAKE_CASE						   =				prot.remark
					if remark is not None:
										pdb_headers.append(f"""REMARK {remark}"""   )

					__SCREAMING_SNAKE_CASE						   =				prot.parents
					__SCREAMING_SNAKE_CASE						   =				prot.parents_chain_index
					if parents is not None and parents_chain_index is not None:
										__SCREAMING_SNAKE_CASE						   =				[p for i, p in zip(lowerCAmelCase_  ,						lowerCAmelCase_   ) if i == chain_id]

					if parents is None or len(lowerCAmelCase_   ) == 0:
										__SCREAMING_SNAKE_CASE						   =				["N/A"]

					pdb_headers.append(f"""PARENT {' '.join(lowerCAmelCase_   )}"""   )

					return pdb_headers




def    UpperCAmelCase__       (lowerCAmelCase_  ,						lowerCAmelCase_   ):




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE						   =				[]
					__SCREAMING_SNAKE_CASE						   =				pdb_str.split("\n"   )

					__SCREAMING_SNAKE_CASE						   =				prot.remark
					if remark is not None:
										out_pdb_lines.append(f"""REMARK {remark}"""   )

					__SCREAMING_SNAKE_CASE						   =				42
					if prot.parents is not None and len(prot.parents   ) > 0:
										__SCREAMING_SNAKE_CASE						   =				[]
										if prot.parents_chain_index is not None:
															__SCREAMING_SNAKE_CASE						   =				{}
															for p, i in zip(prot.parents  ,						prot.parents_chain_index   ):
																				parent_dict.setdefault(str(lowerCAmelCase_   )  ,						[]   )
																				parent_dict[str(lowerCAmelCase_   )].append(lowerCAmelCase_   )

															__SCREAMING_SNAKE_CASE						   =				max([int(lowerCAmelCase_   ) for chain_idx in parent_dict]   )
															for i in range(max_idx + 1   ):
																				__SCREAMING_SNAKE_CASE						   =				parent_dict.get(str(lowerCAmelCase_   )  ,						["N/A"]   )
																				parents_per_chain.append(lowerCAmelCase_   )
										else:
															parents_per_chain.append(list(prot.parents   )   )
					else:
										__SCREAMING_SNAKE_CASE						   =				[["N/A"]]

					def make_parent_line(lowerCAmelCase_   ) -> str:
										return f"""PARENT {' '.join(lowerCAmelCase_   )}"""

					out_pdb_lines.append(make_parent_line(parents_per_chain[0]   )   )

					__SCREAMING_SNAKE_CASE						   =				0
					for i, l in enumerate(lowerCAmelCase_   ):
										if "PARENT" not in l and "REMARK" not in l:
															out_pdb_lines.append(lowerCAmelCase_   )
										if "TER" in l and "END" not in lines[i + 1]:
															chain_counter += 1
															if not chain_counter >= len(lowerCAmelCase_   ):
																				__SCREAMING_SNAKE_CASE						   =				parents_per_chain[chain_counter]
															else:
																				__SCREAMING_SNAKE_CASE						   =				["N/A"]

															out_pdb_lines.append(make_parent_line(lowerCAmelCase_   )   )

					return "\n".join(lowerCAmelCase_   )




def    UpperCAmelCase__       (lowerCAmelCase_   ):




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE						   =				residue_constants.restypes + ["X"]

					def res_atoa(lowerCAmelCase_   ) -> str:
										return residue_constants.restype_atoa.get(restypes[r]  ,						"UNK"   )

					__SCREAMING_SNAKE_CASE						   =				residue_constants.atom_types

					__SCREAMING_SNAKE_CASE						   =				[]

					__SCREAMING_SNAKE_CASE						   =				prot.atom_mask
					__SCREAMING_SNAKE_CASE						   =				prot.aatype
					__SCREAMING_SNAKE_CASE						   =				prot.atom_positions
					__SCREAMING_SNAKE_CASE						   =				prot.residue_index.astype(np.intaa   )
					__SCREAMING_SNAKE_CASE						   =				prot.b_factors
					__SCREAMING_SNAKE_CASE						   =				prot.chain_index

					if np.any(aatype > residue_constants.restype_num   ):
										raise ValueError("Invalid aatypes."   )

					__SCREAMING_SNAKE_CASE						   =				get_pdb_headers(lowerCAmelCase_   )
					if len(lowerCAmelCase_   ) > 0:
										pdb_lines.extend(lowerCAmelCase_   )

					__SCREAMING_SNAKE_CASE						   =				aatype.shape[0]
					__SCREAMING_SNAKE_CASE						   =				1
					__SCREAMING_SNAKE_CASE						   =				0
					__SCREAMING_SNAKE_CASE						   =				string.ascii_uppercase
					__SCREAMING_SNAKE_CASE						   =				None
					# Add all atom sites.
					for i in range(lowerCAmelCase_   ):
										__SCREAMING_SNAKE_CASE						   =				res_atoa(aatype[i]   )
										for atom_name, pos, mask, b_factor in zip(lowerCAmelCase_  ,						atom_positions[i]  ,						atom_mask[i]  ,						b_factors[i]   ):
															if mask < 0.5:
																				continue

															__SCREAMING_SNAKE_CASE						   =				"ATOM"
															__SCREAMING_SNAKE_CASE						   =				atom_name if len(lowerCAmelCase_   ) == 4 else f""" {atom_name}"""
															__SCREAMING_SNAKE_CASE						   =				""
															__SCREAMING_SNAKE_CASE						   =				""
															__SCREAMING_SNAKE_CASE						   =				1.00
															__SCREAMING_SNAKE_CASE						   =				atom_name[0]  # Protein supports only C, N, O, S, this works.
															__SCREAMING_SNAKE_CASE						   =				""

															__SCREAMING_SNAKE_CASE						   =				"A"
															if chain_index is not None:
																				__SCREAMING_SNAKE_CASE						   =				chain_tags[chain_index[i]]

															# PDB is a columnar format, every space matters here!
															__SCREAMING_SNAKE_CASE						   =				(
															    f"""{record_type:<6}{atom_index:>5} {name:<4}{alt_loc:>1}"""
															    f"""{res_name_a:>3} {chain_tag:>1}"""
															    f"""{residue_index[i]:>4}{insertion_code:>1}   """
															    f"""{pos[0]:>8.3f}{pos[1]:>8.3f}{pos[2]:>8.3f}"""
															    f"""{occupancy:>6.2f}{b_factor:>6.2f}          """
															    f"""{element:>2}{charge:>2}"""
															)
															pdb_lines.append(lowerCAmelCase_   )
															atom_index += 1

										__SCREAMING_SNAKE_CASE						   =				i == n - 1
										if chain_index is not None:
															if i != n - 1 and chain_index[i + 1] != prev_chain_index:
																				__SCREAMING_SNAKE_CASE						   =				True
																				__SCREAMING_SNAKE_CASE						   =				chain_index[i + 1]

										if should_terminate:
															# Close the chain.
															__SCREAMING_SNAKE_CASE						   =				"TER"
															__SCREAMING_SNAKE_CASE						   =				(
															    f"""{chain_end:<6}{atom_index:>5}      {res_atoa(aatype[i]   ):>3} {chain_tag:>1}{residue_index[i]:>4}"""
															)
															pdb_lines.append(lowerCAmelCase_   )
															atom_index += 1

															if i != n - 1:
																				# "prev" is a misnomer here. This happens at the beginning of
																				# each new chain.
																				pdb_lines.extend(get_pdb_headers(lowerCAmelCase_  ,						lowerCAmelCase_   )   )

					pdb_lines.append("END"   )
					pdb_lines.append(""   )
					return "\n".join(lowerCAmelCase_   )




def    UpperCAmelCase__       (lowerCAmelCase_   ):




					'''simple docstring'''


					return residue_constants.STANDARD_ATOM_MASK[prot.aatype]




def    UpperCAmelCase__       (lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_ = None  ,						lowerCAmelCase_ = None  ,						lowerCAmelCase_ = None  ,						lowerCAmelCase_ = None  ,						lowerCAmelCase_ = None  ,						):




					'''simple docstring'''


					return Protein(
					    aatype=features["aatype"]  ,						atom_positions=result["final_atom_positions"]  ,						atom_mask=result["final_atom_mask"]  ,						residue_index=features["residue_index"] + 1  ,						b_factors=b_factors if b_factors is not None else np.zeros_like(result["final_atom_mask"]   )  ,						chain_index=lowerCAmelCase_  ,						remark=lowerCAmelCase_  ,						parents=lowerCAmelCase_  ,						parents_chain_index=lowerCAmelCase_  ,						)




"""simple docstring"""



import unittest

from transformers import JukeboxTokenizer
from transformers.testing_utils import require_torch







class _lowerCAmelCase				(      unittest.TestCase	):
   __UpperCAmelCase :       Union[str, Any]							  =  JukeboxTokenizer
   __UpperCAmelCase :       Union[str, Any]							  =  {
       '''artist''': '''Zac Brown Band''',
       '''genres''': '''Country''',
       '''lyrics''': '''I met a traveller from an antique land,
        Who said "Two vast and trunkless legs of stone
        Stand in the desert. . . . Near them, on the sand,
        Half sunk a shattered visage lies, whose frown,
        And wrinkled lip, and sneer of cold command,
        Tell that its sculptor well those passions read
        Which yet survive, stamped on these lifeless things,
        The hand that mocked them, and the heart that fed;
        And on the pedestal, these words appear:
        My name is Ozymandias, King of Kings;
        Look on my Works, ye Mighty, and despair!
        Nothing beside remains. Round the decay
        Of that colossal Wreck, boundless and bare
        The lone and level sands stretch far away
        ''',
   }







   @require_torch
   def  lowerCamelCase		(	self						)	->		int:

       '''simple docstring'''




       import torch

       snake_case    :      Any      					=    JukeboxTokenizer.from_pretrained("openai/jukebox-1b-lyrics"						)
       snake_case    :      Optional[Any]      					=    tokenizer(**self.metas						)["input_ids"]
       # fmt: off
       snake_case    :      Optional[int]      					=    [
           torch.tensor([[
               0, 0, 0, 7169, 507, 9, 76, 39, 31, 46, 76, 27,
               76, 46, 44, 27, 48, 31, 38, 38, 31, 44, 76, 32,
               44, 41, 39, 76, 27, 40, 76, 27, 40, 46, 35, 43,
               47, 31, 76, 38, 27, 40, 30, 64, 78, 76, 76, 76,
               76, 76, 76, 76, 76, 23, 34, 41, 76, 45, 27, 35,
               30, 76, 71, 20, 49, 41, 76, 48, 27, 45, 46, 76,
               27, 40, 30, 76, 46, 44, 47, 40, 37, 38, 31, 45,
               45, 76, 38, 31, 33, 45, 76, 41, 32, 76, 45, 46,
               41, 40, 31, 78, 76, 76, 76, 76, 76, 76, 76, 76,
               19, 46, 27, 40, 30, 76, 35, 40, 76, 46, 34, 31,
               76, 30, 31, 45, 31, 44, 46, 63, 76, 63, 76, 63,
               76, 63, 76, 14, 31, 27, 44, 76, 46, 34, 31, 39,
               64, 76, 41, 40, 76, 46, 34, 31, 76, 45, 27, 40,
               30, 64, 78, 76, 76, 76, 76, 76, 76, 76, 76, 8,
               27, 38, 32, 76, 45, 47, 40, 37, 76, 27, 76, 45,
               34, 27, 46, 46, 31, 44, 31, 30, 76, 48, 35, 45,
               27, 33, 31, 76, 38, 35, 31, 45, 64, 76, 49, 34,
               41, 45, 31, 76, 32, 44, 41, 49, 40, 64, 78, 76,
               76, 76, 76, 76, 76, 76, 76, 1, 40, 30, 76, 49,
               44, 35, 40, 37, 38, 31, 30, 76, 38, 35, 42, 64,
               76, 27, 40, 30, 76, 45, 40, 31, 31, 44, 76, 41,
               32, 76, 29, 41, 38, 30, 76, 29, 41, 39, 39, 27,
               40, 30, 64, 78, 76, 76, 76, 76, 76, 76, 76, 76,
               20, 31, 38, 38, 76, 46, 34, 27, 46, 76, 35, 46,
               45, 76, 45, 29, 47, 38, 42, 46, 41, 44, 76, 49,
               31, 38, 38, 76, 46, 34, 41, 45, 31, 76, 42, 27,
               45, 45, 35, 41, 40, 45, 76, 44, 31, 27, 30, 78,
               76, 76, 76, 76, 76, 76, 76, 76, 23, 34, 35, 29,
               34, 76, 51, 31, 46, 76, 45, 47, 44, 48, 35, 48,
               31, 64, 76, 45, 46, 27, 39, 42, 31, 30, 76, 41,
               40, 76, 46, 34, 31, 45, 31, 76, 38, 35, 32, 31,
               38, 31, 45, 45, 76, 46, 34, 35, 40, 33, 45, 64,
               78, 76, 76, 76, 76, 76, 76, 76, 76, 20, 34, 31,
               76, 34, 27, 40, 30, 76, 46, 34, 27, 46, 76, 39,
               41, 29, 37, 31, 30, 76, 46, 34, 31, 39, 64, 76,
               27, 40, 30, 76, 46, 34, 31, 76, 34, 31, 27, 44,
               46, 76, 46, 34, 27, 46, 76, 32, 31, 30, 66, 78,
               76, 76, 76, 76, 76, 76, 76, 76, 1, 40, 30, 76,
               41, 40, 76, 46, 34, 31, 76, 42, 31, 30, 31, 45,
               46, 27, 38, 64, 76, 46, 34, 31, 45, 31, 76, 49,
               41, 44, 30, 45, 76, 27, 42, 42, 31, 27, 44, 65,
               78, 76, 76, 76, 76, 76, 76, 76, 76, 13, 51, 76,
               40, 27, 39, 31, 76, 35, 45, 76, 15, 52, 51, 39,
               27, 40, 30, 35, 27, 45, 64, 76, 11, 35, 40, 33,
               76, 41, 32, 76, 11, 35, 40, 33, 45, 66, 78, 76,
               76, 76, 76, 76, 76, 76, 76, 12, 41, 41, 37, 76,
               41, 40, 76, 39, 51, 76, 23, 41, 44, 37, 45, 64,
               76, 51, 31, 76, 13, 35, 33, 34, 46, 51, 64, 76,
               27, 40, 30, 76, 30, 31, 45, 42, 27, 35, 44, 67,
               78, 76, 76, 76, 76, 76, 76, 76, 76, 14, 41, 46,
               34, 35, 40, 33, 76, 28, 31, 45, 35, 30, 31, 76,
               44, 31, 39, 27, 35, 40, 45, 63, 76, 18, 41, 47,
               40, 30, 76, 46, 34, 31, 76, 30, 31, 29, 27, 51,
               78, 76, 76, 76, 76, 76, 76, 76, 76, 15, 32, 76,
               46, 34, 27, 46, 76, 29, 41, 38, 41, 45, 45, 27,
               38, 76, 23, 44, 31, 29, 37, 64, 76, 28, 41, 47,
               40, 30, 38, 31, 45, 45, 76, 27, 40, 30, 76, 28,
               27, 44, 31, 78, 76, 76, 76, 76, 76, 76, 76, 76,
               20, 34, 31, 76, 38, 41, 40, 31, 76, 27, 40, 30,
               76, 38, 31, 48, 31, 38, 76, 45, 27, 40, 30, 45,
               76, 45, 46, 44, 31, 46, 29, 34, 76, 32, 27, 44,
               76, 27, 49, 27, 51, 78, 76, 76, 76, 76, 76, 76,
               76, 76]]						),
           torch.tensor([[0, 0, 0, 1069, 11]]						),
           torch.tensor([[0, 0, 0, 1069, 11]]						),
       ]
       # fmt: on
       self.assertTrue(torch.allclose(tokens[0]						,	EXPECTED_OUTPUT[0]						)						)
       self.assertTrue(torch.allclose(tokens[1]						,	EXPECTED_OUTPUT[1]						)						)
       self.assertTrue(torch.allclose(tokens[2]						,	EXPECTED_OUTPUT[2]						)						)







   @require_torch
   def  lowerCamelCase		(	self						)	->		Any:

       '''simple docstring'''




       import torch

       snake_case    :      Tuple      					=    JukeboxTokenizer.from_pretrained("openai/jukebox-5b-lyrics"						)
       snake_case    :      Optional[Any]      					=    tokenizer(**self.metas						)["input_ids"]
       # fmt: off
       snake_case    :      List[Any]      					=    [
           torch.tensor([[
               0, 0, 0, 1069, 11, -1, -1, -1, -1, 9, 77, 39,
               31, 46, 77, 27, 77, 46, 44, 27, 48, 31, 38, 38,
               31, 44, 77, 32, 44, 41, 39, 77, 27, 40, 77, 27,
               40, 46, 35, 43, 47, 31, 77, 38, 27, 40, 30, 64,
               79, 77, 77, 77, 77, 77, 77, 77, 77, 23, 34, 41,
               77, 45, 27, 35, 30, 77, 72, 20, 49, 41, 77, 48,
               27, 45, 46, 77, 27, 40, 30, 77, 46, 44, 47, 40,
               37, 38, 31, 45, 45, 77, 38, 31, 33, 45, 77, 41,
               32, 77, 45, 46, 41, 40, 31, 79, 77, 77, 77, 77,
               77, 77, 77, 77, 19, 46, 27, 40, 30, 77, 35, 40,
               77, 46, 34, 31, 77, 30, 31, 45, 31, 44, 46, 63,
               77, 63, 77, 63, 77, 63, 77, 14, 31, 27, 44, 77,
               46, 34, 31, 39, 64, 77, 41, 40, 77, 46, 34, 31,
               77, 45, 27, 40, 30, 64, 79, 77, 77, 77, 77, 77,
               77, 77, 77, 8, 27, 38, 32, 77, 45, 47, 40, 37,
               77, 27, 77, 45, 34, 27, 46, 46, 31, 44, 31, 30,
               77, 48, 35, 45, 27, 33, 31, 77, 38, 35, 31, 45,
               64, 77, 49, 34, 41, 45, 31, 77, 32, 44, 41, 49,
               40, 64, 79, 77, 77, 77, 77, 77, 77, 77, 77, 1,
               40, 30, 77, 49, 44, 35, 40, 37, 38, 31, 30, 77,
               38, 35, 42, 64, 77, 27, 40, 30, 77, 45, 40, 31,
               31, 44, 77, 41, 32, 77, 29, 41, 38, 30, 77, 29,
               41, 39, 39, 27, 40, 30, 64, 79, 77, 77, 77, 77,
               77, 77, 77, 77, 20, 31, 38, 38, 77, 46, 34, 27,
               46, 77, 35, 46, 45, 77, 45, 29, 47, 38, 42, 46,
               41, 44, 77, 49, 31, 38, 38, 77, 46, 34, 41, 45,
               31, 77, 42, 27, 45, 45, 35, 41, 40, 45, 77, 44,
               31, 27, 30, 79, 77, 77, 77, 77, 77, 77, 77, 77,
               23, 34, 35, 29, 34, 77, 51, 31, 46, 77, 45, 47,
               44, 48, 35, 48, 31, 64, 77, 45, 46, 27, 39, 42,
               31, 30, 77, 41, 40, 77, 46, 34, 31, 45, 31, 77,
               38, 35, 32, 31, 38, 31, 45, 45, 77, 46, 34, 35,
               40, 33, 45, 64, 79, 77, 77, 77, 77, 77, 77, 77,
               77, 20, 34, 31, 77, 34, 27, 40, 30, 77, 46, 34,
               27, 46, 77, 39, 41, 29, 37, 31, 30, 77, 46, 34,
               31, 39, 64, 77, 27, 40, 30, 77, 46, 34, 31, 77,
               34, 31, 27, 44, 46, 77, 46, 34, 27, 46, 77, 32,
               31, 30, 66, 79, 77, 77, 77, 77, 77, 77, 77, 77,
               1, 40, 30, 77, 41, 40, 77, 46, 34, 31, 77, 42,
               31, 30, 31, 45, 46, 27, 38, 64, 77, 46, 34, 31,
               45, 31, 77, 49, 41, 44, 30, 45, 77, 27, 42, 42,
               31, 27, 44, 65, 79, 77, 77, 77, 77, 77, 77, 77,
               77, 13, 51, 77, 40, 27, 39, 31, 77, 35, 45, 77,
               15, 52, 51, 39, 27, 40, 30, 35, 27, 45, 64, 77,
               11, 35, 40, 33, 77, 41, 32, 77, 11, 35, 40, 33,
               45, 66, 79, 77, 77, 77, 77, 77, 77, 77, 77, 12,
               41, 41, 37, 77, 41, 40, 77, 39, 51, 77, 23, 41,
               44, 37, 45, 64, 77, 51, 31, 77, 13, 35, 33, 34,
               46, 51, 64, 77, 27, 40, 30, 77, 30, 31, 45, 42,
               27, 35, 44, 67, 79, 77, 77, 77, 77, 77, 77, 77,
               77, 14, 41, 46, 34, 35, 40, 33, 77, 28, 31, 45,
               35, 30, 31, 77, 44, 31, 39, 27, 35, 40, 45, 63,
               77, 18, 41, 47, 40, 30, 77, 46, 34, 31, 77, 30,
               31, 29, 27, 51, 79, 77, 77, 77, 77, 77, 77, 77,
               77, 15, 32, 77, 46, 34, 27, 46, 77, 29, 41, 38,
               41, 45, 45, 27, 38, 77, 23, 44, 31, 29, 37, 64,
               77, 28, 41, 47, 40, 30, 38, 31, 45, 45, 77, 27,
               40, 30, 77, 28, 27, 44, 31, 79, 77, 77, 77, 77,
               77, 77, 77, 77, 20, 34, 31, 77, 38, 41, 40, 31,
               77, 27, 40, 30, 77, 38, 31, 48, 31, 38, 77, 45,
               27, 40, 30, 45, 77, 45, 46, 44, 31, 46, 29, 34,
               77, 32, 27, 44, 77, 27, 49, 27, 51, 79, 77, 77,
               77, 77, 77, 77, 77, 77]]						),
           torch.tensor([[0, 0, 0, 1069, 11, -1, -1, -1, -1]]						),
           torch.tensor([[0, 0, 0, 1069, 11, -1, -1, -1, -1]]						),
       ]
       # fmt: on
       self.assertTrue(torch.allclose(tokens[0]						,	EXPECTED_OUTPUT[0]						)						)
       self.assertTrue(torch.allclose(tokens[1]						,	EXPECTED_OUTPUT[1]						)						)
       self.assertTrue(torch.allclose(tokens[2]						,	EXPECTED_OUTPUT[2]						)						)





