Divide sample into a and b segments.
(sample, np_rng)
| 79 | |
| 80 | |
| 81 | def get_a_and_b_segments(sample, np_rng): |
| 82 | """Divide sample into a and b segments.""" |
| 83 | |
| 84 | # Number of sentences in the sample. |
| 85 | n_sentences = len(sample) |
| 86 | # Make sure we always have two sentences. |
| 87 | assert n_sentences > 1, "make sure each sample has at least two sentences." |
| 88 | |
| 89 | # First part: |
| 90 | # `a_end` is how many sentences go into the `A`. |
| 91 | a_end = 1 |
| 92 | if n_sentences >= 3: |
| 93 | # Note that randin in numpy is exclusive. |
| 94 | a_end = np_rng.randint(1, n_sentences) |
| 95 | tokens_a = [] |
| 96 | for j in range(a_end): |
| 97 | tokens_a.extend(sample[j]) |
| 98 | |
| 99 | # Second part: |
| 100 | tokens_b = [] |
| 101 | for j in range(a_end, n_sentences): |
| 102 | tokens_b.extend(sample[j]) |
| 103 | |
| 104 | # Random next: |
| 105 | is_next_random = False |
| 106 | if np_rng.random() < 0.5: |
| 107 | is_next_random = True |
| 108 | tokens_a, tokens_b = tokens_b, tokens_a |
| 109 | |
| 110 | return tokens_a, tokens_b, is_next_random |
| 111 | |
| 112 | |
| 113 | def truncate_segments(tokens_a, tokens_b, len_a, len_b, max_num_tokens, np_rng): |
nothing calls this directly
no outgoing calls
no test coverage detected