| from completions import calculate_log_probabilities, load_model, tokenize, split_into_words | |
| model, tokenizer, device = load_model() | |
| def test_text_to_words(): | |
| text = """Hello | |
| world!""" | |
| token_probs = calculate_log_probabilities(model, tokenizer, tokenize(text, tokenizer, device)) | |
| words = split_into_words(token_probs, tokenizer) | |
| expected_words = ["Hello", "\n", "world", "!"] | |
| assert [w.text for w in words] == expected_words | |
| def test_multiline(): | |
| text = """// Context: C code from an image manipulation library. | |
| for (int y = 0; y < HEIGHT; y++) { | |
| for (int x = 0; x < WIDTH; x++) { | |
| buf[y * HEIGHT + x] = 0; | |
| } | |
| }""" | |
| tokenized = tokenize(text, tokenizer, device) | |
| print(tokenized) | |
| token_probs = calculate_log_probabilities(model, tokenizer, tokenized) | |
| words = split_into_words(token_probs, tokenizer) | |
| print("---", [w.text for w in words]) | |
| expected_words = ["//", " Context", ":", " C", " code", " from", " an", " image", " manipulation", " library", ".\n", | |
| "for", "(", "int", "y", "=", "0", ";", "y", "<", "HEIGHT", ";", "y", "+", "+", ")", "{", "\n", " ", "for", "(", "int", "x", "=", "0", ";", "x", "<", "WIDTH", ";", "x", "+", "+", ")", "{", "\n", " ", "buf", "[", "y", "*", "HEIGHT", "+", "x", "]", "=", "0", ";", "\n", " ", "}", "\n", "}"] | |
| assert [w.text for w in words] == expected_words | |