Snapshot

completions_test.py

@@ -0,0 +1,27 @@
+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