add related words

app.py

@@ -38,7 +38,7 @@ def get_models(llama=False):
 model, tokenizer = get_models()
 
 def return_top_k(sentence, k=10):
-  
+
   if sentence[-1] != ".":
     sentence = sentence + "."
 
@@ -69,9 +69,34 @@ def return_top_k(sentence, k=10):
       if (len(pred) < 2) | (pred in sentence.split()):
         predictions.pop(predictions.index(pred))  
      
-  return predictions[:10]
+  return predictions[:10] 
+
+# JS 
+def get_related_words(word, num=5):
+  model.eval()
+  with torch.no_grad():
+      sentence = [f"Descripton : It is related to {word} but not {word}. Word : "]
+      #inputs = ["Description: It is something to cut stuff with. Word: "]
+      print(sentence)
+      inputs = tokenizer(sentence, padding=True, truncation=True, return_tensors="pt",)
+
+      device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+      model.to(device)
+
+      batch = {k: v.to(device) for k, v in inputs.items()}
+      beam_outputs = model.generate(
+          input_ids=batch['input_ids'], max_new_tokens=10, num_beams=num+2, num_return_sequences=num+2, early_stopping=True
+      )
+      
+      #beam_preds = [tokenizer.decode(beam_output.detach().cpu().numpy(), skip_special_tokens=True) for beam_output in beam_outputs if ]
+      beam_preds = []
+      for beam_output in beam_outputs:
+          prediction = tokenizer.decode(beam_output.detach().cpu().numpy(), skip_special_tokens=True).strip()
+          if prediction not in " ".join(sentence):
+            beam_preds.append(prediction)
+      
+      return ", ".join(beam_preds[:num])
 
-  
 if 'messages' not in st.session_state:
   st.session_state.messages = []
   
@@ -267,10 +292,9 @@ if st.session_state.actions[-1] == 'cue':
     target = st.session_state.results["results"][word_count]
     if get_available_cues(target):
       avail_cues = get_available_cues(target)
-      cues_buttons = {}
-      for cue_type, cues in avail_cues.items():
-        #st.button(cue_type)
-        cues_buttons[cue_type] = st.button(cue_type)
+      cues_buttons = {cue_type: st.button(cue_type) for cue_type in avail_cues}
+      #for cue_type in avail_cues:
+       # cues_buttons[cue_type] = st.button(cue_type)
 
     with col3:
       b3 = st.button("All words", key="3")
@@ -279,6 +303,8 @@ if st.session_state.actions[-1] == 'cue':
     with col5:
       b5 = st.button("Exit", key="5", type='primary')
 
+    b6 = st.button("Related words")
+
     if b1:
       st.session_state.counters["letter_count"] += 1
       word_count = st.session_state.counters["word_count"]
@@ -307,6 +333,12 @@ if st.session_state.actions[-1] == 'cue':
     elif b3:
       write_bot(f"Here are all my guesses about your word: {st.session_state.results['results_print']}")
 
+    elif b6:
+      sent = f"It is related to '{target}' but not '{target}'."
+      rels = return_top_k(target)
+
+      write_bot(f'Here are words that are related to your word: {", ".join(rels)}')
+
     elif b4:
       write_bot("Yay! I am happy I could be of help!")
       new = st.button('Play again', key=63)