Fix bug in state update with mask

Fix import in wordle game simulator

wordle_env/state.py

@@ -62,7 +62,6 @@ def update_from_mask(state: WordleState, word: str, mask: List[int]) -> WordleSt
     for i, c in enumerate(word):
         cint = ord(c) - ord(WORDLE_CHARS[0])
         offset = 1 + cint * WORDLE_N * 3
-        state[1 + cint] = 1
         if mask[i] == YES:
             prior_yes.append(c)
             # char at position i = yes, all other chars at position i == no
@@ -79,9 +78,9 @@ def update_from_mask(state: WordleState, word: str, mask: List[int]) -> WordleSt
             prior_maybe.append(c)
             # Char at position i = no, and in other positions maybe except it had a value before, other chars stay as they are
             for char_idx in range(0, WORDLE_N * 3, 3):
-                    char_offset = offset + char_idx
-                    if tuple(state[char_offset: char_offset + 3]) == (0, 0, 0):
-                        state[char_offset: char_offset + 3] = [0, 1, 0]
+                char_offset = offset + char_idx
+                if tuple(state[char_offset: char_offset + 3]) == (0, 0, 0):
+                    state[char_offset: char_offset + 3] = [0, 1, 0]
             state[offset + 3 * i:offset + 3 * i + 3] = [1, 0, 0]
         elif mask[i] == NO:
             # Need to check this first in case there's prior maybe + yes
@@ -96,7 +95,7 @@ def update_from_mask(state: WordleState, word: str, mask: List[int]) -> WordleSt
                         state[offset + 3 * j:offset + 3 * j + 3] = [1, 0, 0]
             else:
                 # Just straight up no
-                state[offset:offset+3*WORDLE_N] = [1, 0, 0]*WORDLE_N
+                state[offset:offset + 3 * WORDLE_N] = [1, 0, 0] * WORDLE_N
     return state
 
 

wordle_game.py

@@ -1,7 +1,7 @@
 from rich.prompt import Prompt
 from rich.console import Console
 from random import choice
-from words import target_vocabulary, complete_vocabulary
+from wordle_env.words import target_vocabulary, complete_vocabulary
 
 SQUARES = {
     'correct_place': '🟩',