Optimize with argmax

app.py

@@ -72,14 +72,8 @@ def train(data: str, message: str):
         model.save(f"cache/{data_hash}")
     tokens = tokenizer.texts_to_sequences([message,])[0]
     prediction = model.predict(np.array([(list(tokens)+[0,]*inp_len)[:inp_len],]))[0]
-    max_o = 0
-    max_v = 0
-    for ind, i in enumerate(prediction):
-        if max_v < i:
-            max_v = i
-            max_o = ind
     keras.backend.clear_session()
-    return responses[ind]
+    return responses[np.argmax(prediction)]
 
 iface = gr.Interface(fn=train, inputs=["text", "text"], outputs="text")
 iface.launch()