Change suggestion endpoint for play word endpoint

now the endpoint receives a word and make the AI play
also refactored play module

a3c/eval.py

@@ -7,29 +7,23 @@ from .utils import v_wrap
 
 
 def evaluate_checkpoints(dir, env):
-    n_s = env.observation_space.shape[0]
-    n_a = env.action_space.n
-    words_list = env.words
-    word_width = len(env.words[0])
-    net = GreedyNet(n_s, n_a, words_list, word_width)
     results = {}
     for checkpoint in os.listdir(dir):
-        checkpoint_path = os.path.join(dir, checkpoint)
-        if os.path.isfile(checkpoint_path):
-            net.load_state_dict(torch.load(checkpoint_path))
-            wins, guesses = evaluate(net, env)
+        pretrained_model_path = os.path.join(dir, checkpoint)
+        if os.path.isfile(pretrained_model_path):
+            wins, guesses = evaluate(env, pretrained_model_path)
             results[checkpoint] = wins, guesses
     return dict(sorted(results.items(), key=lambda x: (x[1][0], -x[1][1]), reverse=True))
 
 
-def evaluate(net, env):
+def evaluate(env, pretrained_model_path):
     n_wins = 0
     n_guesses = 0
     n_win_guesses = 0
     env = env.unwrapped
     N = env.allowable_words
     for goal_word in env.words[:N]:
-        win, outcomes = play(net, env)
+        win, outcomes = play(env, pretrained_model_path, goal_word)
         if win:
             n_wins += 1
             n_win_guesses += len(outcomes)

a3c/play.py

@@ -13,6 +13,21 @@ def get_play_model_path():
     return os.path.join(model_checkpoint_dir, model_name)
 
 
+def get_net(env, pretrained_model_path):
+    n_s = env.observation_space.shape[0]
+    n_a = env.action_space.n
+    words_list = env.words
+    word_width = len(env.words[0])
+    net = GreedyNet(n_s, n_a, words_list, word_width)
+    net.load_state_dict(torch.load(pretrained_model_path))
+    return net
+
+
+def get_initial_state(env):
+    state = env.reset()
+    return state
+
+
 def suggest(
         env,
         words,
@@ -27,14 +42,9 @@ def suggest(
     :param sequence: History of moves and outcomes until now
     :return:
     """
-    n_s = env.observation_space.shape[0]
-    n_a = env.action_space.n
     env = env.unwrapped
-    state = env.reset()
-    words_list = env.words
-    word_width = len(env.words[0])
-    net = GreedyNet(n_s, n_a, words_list, word_width)
-    net.load_state_dict(torch.load(pretrained_model_path))
+    net = get_net(env, pretrained_model_path)
+    state = get_initial_state(env)
     for word, mask in zip(words, states):
         word = word.upper()
         mask = list(map(int, mask))
@@ -42,16 +52,20 @@ def suggest(
     return env.words[net.choose_action(v_wrap(state[None, :]))]
 
 
-def play(net, env):
-    state = env.reset()
+def play(env, pretrained_model_path, goal_word = None):
+    env = env.unwrapped
+    net = get_net(env, pretrained_model_path)
+    state = get_initial_state(env)
+    if goal_word:
+        env.set_goal_word(goal_word)
     outcomes = []
     win = False
     for i in range(env.max_turns):
         action = net.choose_action(v_wrap(state[None, :]))
         state, reward, done, _ = env.step(action)
-        outcomes.append((env.words[action], reward))
+        outcomes.append(env.words[action])
         if done:
-            if reward >= 0:
+            if reward > 0:
                 win = True
             break
     return win, outcomes

api_rest/api.py

@@ -1,38 +1,36 @@
-from a3c.play import get_play_model_path, suggest
+from a3c.play import get_play_model_path, play
 from flask import Flask, request, jsonify
+from flask_cors import cross_origin
 from wordle_env.words import target_vocabulary
 from wordle_env.wordle import get_env
 
 app = Flask(__name__)
 
 
-def validate_params(words, states):
-    # Check if the input lists are valid (i.e. all elements have length 5 and numbers are between 0 and 2 inclusive)
-    if not all(len(w) == 5 and w in target_vocabulary for w in words):
-        return True, 'Invalid input, words must be 5 characters long and must be an eligible word'
-
-    if not all(len(n) == 5 and all(c.isdigit() and 0 <= int(c) <= 2 for c in n) for n in states):
-        return True, 'Invalid input, states must be 5 characters long and the numbers between 0 and 2 inclusive'
-
+def validate_goal_word(word):
+    if not word:
+        return True, 'Goal word not provided'
+    if word.upper() not in target_vocabulary:
+        return True, 'Goal word not in vocabulary'
     return False, ''
 
 
-@app.route('/suggest', methods=['GET'])
-def get_suggestion():
+@app.route('/play_word', methods=['GET'])
+@cross_origin(origin='*', headers=['Content-Type', 'Authorization'])
+def get_play():
     # Get the list of words and list of number strings from the request
-    words = [word.strip().upper()
-             for word in request.args.get('words').split(',')]
-    states = [state.strip() for state in request.args.get('states').split(',')]
-    print(states)
-    error, msge = validate_params(words, states)
+    word = request.args.get('goal_word')
+
+    error, msge = validate_goal_word(word)
     if error:
         return jsonify({'error': msge}), 400
 
+    word = word.upper()
     env = get_env()
     model_path = get_play_model_path()
     # Call the suggest function with the input lists and return the result
-    suggestion = suggest(env, words, states, model_path)
-    return jsonify({'suggestion': suggestion})
+    won, attempts = play(env, model_path, word)
+    return jsonify({'attempts': attempts, 'won': won})
 
 
 if __name__ == '__main__':