-
Notifications
You must be signed in to change notification settings - Fork 0
/
Run.py
60 lines (50 loc) · 1.5 KB
/
Run.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
import string
from Lib import *
# The trans 5 model had the best F1 score
model = NeuralNetwork()
model.load_state_dict(torch.load("models/model_trans5.pth"))
model.eval()
# Establish knowledge vectors
k_vecs = [[UNKNOWN] * 26 for _ in range(5)]
# Mark letters as incorrect
incorrect_letters = ""
for i in range(5):
for c in incorrect_letters:
c_idx = string.ascii_lowercase.index(c)
k_vecs[i][c_idx] = INCORRECT
# Mark letters that we know are in the word, but not where
maybe_idxs = {}
# Example:
# idx = string.ascii_lowercase.index('i')
# maybe_idxs[idx] = [0]
for letter, poses in maybe_idxs.items():
for i in range(0, 5):
k_vecs[i][letter] = MAYBE
for pos in poses:
k_vecs[pos][letter] = INCORRECT
# Mark correct letters
correct_letters = {} # Example: {'e': 4}
for c, i in correct_letters.items():
c_idx = string.ascii_lowercase.index(c)
k_vecs[i][c_idx] = CORRECT
# Concat all vectors
k_vec = []
for v in k_vecs:
k_vec += v
# Generate predictions
k_vec = torch.IntTensor(k_vec)
pred = model(k_vec).sigmoid()
pred = pred.reshape((5, 26))
print(pred)
# Identify the word that best matches the predicted probabilities
best = None
with open("words.txt", "r") as file:
for word in file:
word = word.strip()
p = 1.0
for i, c in enumerate(word):
char_idx = string.ascii_lowercase.index(c)
p *= pred[i][char_idx]
if best == None or best[1] < p:
best = (word, p)
print(best)