(text: str, eos_emb, threshold, add_begin)
| 80 | |
| 81 | @torch.no_grad() |
| 82 | def predict_next_text(text: str, eos_emb, threshold, add_begin): |
| 83 | sents_original = sent_tokenize(text) |
| 84 | if add_begin: |
| 85 | sents_original = ["Begin of text."] + sents_original |
| 86 | sents = [s.strip() for s in sents_original if s.strip()] |
| 87 | embeddings = [] |
| 88 | |
| 89 | while len(sents) <= 32: |
| 90 | next_sentence, embeddings, stop = predict_next_sentence( |
| 91 | sents, embeddings, eos_emb, threshold |
| 92 | ) |
| 93 | if stop: |
| 94 | break |
| 95 | sents_original.append(next_sentence) |
| 96 | sents.append(next_sentence.strip()) |
| 97 | |
| 98 | if add_begin: |
| 99 | sents_original = sents_original[1:] |
| 100 | return " ".join(sents_original) |
| 101 | |
| 102 | |
| 103 | if __name__ == "__main__": |
no test coverage detected