Bug fix

app.py

@@ -2,7 +2,7 @@ import streamlit as st
 import torch
 from tqdm import tqdm
 from peft import PeftModel, PeftConfig
-from transformers import AutoModelForSeq2SeqLM
+from transformers import AutoModelForSeq2SeqLM, AutoModelForCausalLM
 from transformers import AutoTokenizer
 import numpy as np
 import time
@@ -11,9 +11,10 @@ import time
 import nltk
 nltk.download('wordnet')
 from nltk.corpus import wordnet as wn
+from nltk.tokenize import word_tokenize
 
 @st.cache_resource
-def get_models():
+def get_models(llama=False):
   st.write('Loading the model...')
   config = PeftConfig.from_pretrained("NursNurs/T5ForReverseDictionary")
   model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-large")
@@ -21,8 +22,16 @@ def get_models():
 
   tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large")
 
+  # JS
+  if llama:
+    model_name = 'meta-llama/Llama-2-7b-chat-hf'
+    access_token = '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'
+    llama_tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=access_token, use_fast=True)#, use_fast=True)
+    llama_model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=access_token, device_map={'':0})#, load_in_4bit=True)
+
   st.write("The assistant is loaded and ready to use!")
-  return model, tokenizer
+
+  return model, tokenizer, llama_model, llama_tokenizer if llama else model, tokenizer
 
 model, tokenizer = get_models()
 
@@ -124,6 +133,7 @@ def ask_if_helped():
     write_bot("I am happy to help!")
   elif n:
     st.session_state.actions.append('cue')
+    st.session_state.is_helpful['ask'] = False
     #cue_generation()
   elif new:
     write_bot("Please describe your word!")
@@ -132,6 +142,20 @@ def ask_if_helped():
 ## removed: if st.session_state.actions[-1] == "result":
 
 # JS
+def get_related_words_llama(relation, target, device, num=5):
+    prompt = f"Provide {num} {relation}s for the word '{target}'. Your answer consists of these {num} words only. Do not include the word '{target}' itself in your answer"
+
+    inputs = tokenizer([prompt], return_tensors='pt').to(device)
+    output = model.generate(
+        **inputs, max_new_tokens=40, temperature=.75, early_stopping=True, 
+        )
+    chatbot_response = tokenizer.decode(output[:, inputs['input_ids'].shape[-1]:][0], skip_special_tokens=True).strip()
+
+    postproc = [word for word in word_tokenize(chatbot_response) if len(word)>=3]
+
+    return postproc[-num:] if len(postproc)>=num else postproc
+
+
 def postproc_wn(related_words, syns=False):
     if syns:
         related_words = [word.split('.')[0] if word[0] != "." else word.split('.')[1] for word in related_words]
@@ -174,73 +198,9 @@ def get_available_cues(target):
     else:
       return None 
 
-def cue_generation():
-  if st.session_state.actions[-1] == 'cue':
-    guessed = False
-    write_bot('What do you want to see?', remember=False, blink=False)
-
-    while guessed == False:
-      b1 = st.button("Next letter", key="1")
-      b2 = st.button("Next word", key="2")
-
-      # JS
-      word_count = st.session_state.counters["word_count"]
-      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)
-
-      b3 = st.button("All words", key="3")
-      b4 = st.button("I remembered the word!", key="4", type='primary')
-      b5 = st.button("Exit", key="5", type='primary')
-
-      if b1:
-        st.session_state.counters["letter_count"] += 1
-        word_count = st.session_state.counters["word_count"]
-        letter_count = st.session_state.counters["letter_count"]
-        write_bot(f'The word starts with {st.session_state.results["results"][word_count][:letter_count]}', remember=False)
-
-      elif b2:
-        st.session_state.counters["letter_count"] = 1
-        letter_count = st.session_state.counters["letter_count"]
-        st.session_state.counters["word_count"] += 1
-        word_count = st.session_state.counters["word_count"]
-        write_bot(f'The next word starts with {st.session_state.results["results"][word_count][:letter_count]}', remember=False)
-
-      elif "Synonyms" in cues_buttons and cues_buttons['Synonyms']:
-        write_bot(f'Here are synonyms for the current word: {", ".join(avail_cues["Synonyms"])}', remember=False)
-
-      elif "Hypernyms" in cues_buttons and cues_buttons['Hypernyms']:
-        write_bot(f'Here are hypernyms for the current word: {", ".join(avail_cues["Hypernyms"])}', remember=False)
-
-      elif "Hyponyms" in cues_buttons and cues_buttons['Hyponyms']:
-        write_bot(f'Here are hyponyms for the current word: {", ".join(avail_cues["Hyponyms"])}', remember=False)
-
-      elif "Examples" in cues_buttons and cues_buttons['Examples']:   
-        write_bot(f'Here are example contexts for the current word: {", ".join(avail_cues["Examples"])}', remember=False)
-
-      elif b3:
-        write_bot(f"Here are all my guesses about your word: {st.session_state.results['results_print']}")
-
-      elif b4:
-        write_bot("Yay! I am happy I could be of help!")
-        new = st.button('Play again', key=63)
-        if new:
-          write_bot("Please describe your word!")
-        guessed = True
-      
-        break
-      
-      elif b5:
-        write_bot("I am sorry I couldn't help you this time. See you soon!")
-      st.session_state.actions.append('cue')
-      new = st.button('Play again', key=64)
-      if new:
-        write_bot("Please describe your word!")
-      break
+# JS: moved the cue generation further down
+#def cue_generation():
+ # if st.session_state.actions[-1] == 'cue':
   
 #display user message in chat message container
 prompt = get_text() 
@@ -320,16 +280,16 @@ if st.session_state.actions[-1] == 'cue':
       word_count = st.session_state.counters["word_count"]
       write_bot(f'The next word starts with {st.session_state.results["results"][word_count][:letter_count]}', remember=False)
 
-    elif "Synonyms" in cues_buttons and cues_buttons['Synonyms']:
+    elif get_available_cues(target) and "Synonyms" in cues_buttons and cues_buttons['Synonyms']:
       write_bot(f'Here are synonyms for the current word: {", ".join(avail_cues["Synonyms"])}', remember=False)
 
-    elif "Hypernyms" in cues_buttons and cues_buttons['Hypernyms']:
+    elif get_available_cues(target) and "Hypernyms" in cues_buttons and cues_buttons['Hypernyms']:
       write_bot(f'Here are hypernyms for the current word: {", ".join(avail_cues["Hypernyms"])}', remember=False)
 
-    elif "Hyponyms" in cues_buttons and cues_buttons['Hyponyms']:
+    elif get_available_cues(target) and "Hyponyms" in cues_buttons and cues_buttons['Hyponyms']:
       write_bot(f'Here are hyponyms for the current word: {", ".join(avail_cues["Hyponyms"])}', remember=False)
 
-    elif "Examples" in cues_buttons and cues_buttons['Examples']:   
+    elif get_available_cues(target) and "Examples" in cues_buttons and cues_buttons['Examples']:   
       write_bot(f'Here are example contexts for the current word: {", ".join(avail_cues["Examples"])}', remember=False)
 
     elif b3: