| 57 | return self._vocab_size |
| 58 | |
| 59 | class TranslationTokenizer: |
| 60 | def __init__( |
| 61 | self, |
| 62 | tokenizer: PreTrainedTokenizerBase, |
| 63 | num_bins: Dict, |
| 64 | bin_policy: Optional[Dict] = None, |
| 65 | use_spherical: bool = True, |
| 66 | ): |
| 67 | self.tokenizer = tokenizer |
| 68 | self.num_theta_bins = num_bins["theta_bins"] |
| 69 | self.num_phi_bins = num_bins["phi_bins"] |
| 70 | self.num_r_bins = num_bins["r_bins"] |
| 71 | self.use_spherical = use_spherical |
| 72 | |
| 73 | # for indexing |
| 74 | self.NP = self.num_phi_bins * self.num_r_bins |
| 75 | |
| 76 | # add special action tokens to language tokenizer |
| 77 | self._vocab_size = self.num_theta_bins * self.num_phi_bins * self.num_r_bins |
| 78 | token_list = [ACTION_TOKEN.format(i) for i in range(self._vocab_size)] |
| 79 | self.token_array = np.array(token_list) |
| 80 | |
| 81 | num_new_tokens = self.tokenizer.add_tokens(token_list, special_tokens=True) |
| 82 | print(f"Add {num_new_tokens} TRANSLATION TOKENS, tokenizer vocab size {self.tokenizer.vocab_size} / {len(tokenizer)}") |
| 83 | |
| 84 | self.token_start_idx = self.tokenizer.convert_tokens_to_ids(self.token_array[0]) |
| 85 | self.token_end_idx = self.tokenizer.convert_tokens_to_ids(self.token_array[-1]) |
| 86 | self.set_bins(bin_policy) |
| 87 | |
| 88 | def set_bins(self, bin_policy): |
| 89 | self.theta_bins = np.array(bin_policy["theta_bins"]) |
| 90 | self.phi_bins = np.array(bin_policy["phi_bins"]) |
| 91 | self.r_bins = np.array(bin_policy["r_bins"]) |
| 92 | |
| 93 | def cartesian_to_spherical(self, x, y, z): |
| 94 | theta = np.arctan2(np.sqrt(x**2 + y**2), z) # polar angle |
| 95 | phi = np.arctan2(y, x) # azimuthal angle |
| 96 | r = np.sqrt(x**2 + y**2 + z**2) |
| 97 | return theta, phi, r |
| 98 | |
| 99 | def spherical_to_cartesian(self, theta, phi, r): |
| 100 | x = r * np.sin(theta) * np.cos(phi) |
| 101 | y = r * np.sin(theta) * np.sin(phi) |
| 102 | z = r * np.cos(theta) |
| 103 | return x, y, z |
| 104 | |
| 105 | def __call__(self, action: np.ndarray) -> List[str]: |
| 106 | """Discretize continuous actions to tokens. |
| 107 | action: np.ndarray, (n, 3), continuous actions in Cartesian or Spherical coordinates. |
| 108 | return: np.ndarray, (n,), tokens. |
| 109 | """ |
| 110 | if self.use_spherical: |
| 111 | theta, phi, r = self.cartesian_to_spherical(action[:, 0], action[:, 1], action[:, 2]) |
| 112 | else: |
| 113 | theta, phi, r = action[:, 0], action[:, 1], action[:, 2] |
| 114 | |
| 115 | disc_theta = np.digitize(theta, self.theta_bins[1:-1]) # b |
| 116 | disc_phi = np.digitize(phi, self.phi_bins[1:-1]) |