app.py

import streamlit as st
import torch
from tqdm import tqdm
from peft import PeftModel, PeftConfig
from transformers import AutoModelForSeq2SeqLM, AutoModelForCausalLM
from transformers import AutoTokenizer
import numpy as np
import time
import string

# JS
import nltk
nltk.download('wordnet')
from nltk.corpus import wordnet as wn
from nltk.tokenize import word_tokenize

@st.cache_resource
def get_models(llama=False):
  st.write('Loading the model...')
  config = PeftConfig.from_pretrained("NursNurs/T5ForReverseDictionary")
  model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-large")
  model = PeftModel.from_pretrained(model, "NursNurs/T5ForReverseDictionary")

  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, llama_model, llama_tokenizer 
  
  else:
    st.write("_The assistant is loaded and ready to use! :tada:_")
    return model, tokenizer

model, tokenizer = get_models()

def remove_punctuation(word):
    # Create a translation table that maps all punctuation characters to None
    translator = str.maketrans('', '', string.punctuation)
    
    # Use the translate method to remove punctuation from the word
    word_without_punctuation = word.translate(translator)
    
    return word_without_punctuation

def return_top_k(sentence, k=10, word=None, rels=False):

  if sentence[-1] != ".":
    sentence = sentence + "."

  if rels:
    inputs = [f"Description : It is related to '{word}' but not '{word}'. Word : "]
  else:
    inputs = [f"Description : {sentence} Word : "]

  inputs = tokenizer(
      inputs,
      padding=True, truncation=True,
      return_tensors="pt",
  )
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  model.to(device)

  with torch.no_grad():
    inputs = {k: v.to(device) for k, v in inputs.items()}
    output_sequences = model.generate(input_ids=inputs["input_ids"], max_new_tokens=10, num_beams=k+5, num_return_sequences=k+5, #max_length=3,
                                              top_p = 50, output_scores=True, return_dict_in_generate=True) #repetition_penalty=10000.0

    logits = output_sequences['sequences_scores'].clone().detach()
    decoded_probabilities = torch.softmax(logits, dim=0)


    #all word predictions
    predictions = [tokenizer.decode(tokens, skip_special_tokens=True) for tokens in output_sequences['sequences']]
    probabilities = [round(float(prob), 2) for prob in decoded_probabilities]

    stripped_sent = [remove_punctuation(word.lower()) for word in sentence.split()]
    for pred in predictions:
      if (len(pred) < 2) | (pred in stripped_sent):
        predictions.pop(predictions.index(pred))  
     
  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:

def get_text():
  input_text = st.chat_input()
  return input_text

def write_bot(input, remember=True, blink=True):
    with st.chat_message('assistant'):
      message_placeholder = st.empty()
      full_response = input
      if blink == True:
        response = ''
        for chunk in full_response.split():
          response += chunk + " "
          time.sleep(0.05)
          # Add a blinking cursor to simulate typing
          message_placeholder.markdown(response + "▌") 
        time.sleep(0.5)
      message_placeholder.markdown(full_response)
    if remember == True:
      st.session_state.messages.append({'role': 'assistant', 'content': full_response})
    
def ask_if_helped():
  y = st.button('Yes!', key=60)
  n = st.button('No...', key=61)
  new = st.button('I have a new word', key=62)
  if y:
    write_bot("I am happy to help!")
    st.session_state.is_helpful['ask'] = False
  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!")
    st.session_state.is_helpful['ask'] = False

## 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]
    else:
        related_words = [word.name().split('.')[0] if word.name()[0] != "." else word.name().split('.')[1] for word in related_words]
    related_words = [word.replace("_", " ") for word in related_words]

    return related_words

# JS
def get_available_cues(target):
    wn_nouns = [word.name() for word in wn.all_synsets(pos='n')]
    wn_nouns = [word.split('.')[0] if word[0] != "." else word.split('.')[1] for word in wn_nouns]

    if target in wn_nouns:
      available_cues = {}
      synset_target = wn.synsets(target, pos=wn.NOUN)[0]

      #if wn.synonyms(target)[0]:
       #   available_cues['Synonyms'] = postproc_wn(wn.synonyms(target)[0], syns=True)
      
      #if synset_target.hypernyms():
       #   available_cues['Hypernyms'] = postproc_wn(synset_target.hypernyms())


      #if synset_target.hyponyms():
        #  available_cues['Hyponyms'] = postproc_wn(synset_target.hyponyms())

      if synset_target.examples():
          examples = []
          
          for example in synset_target.examples():
              examples.append(example.replace(target, "..."))
                  
          available_cues['Examples'] = examples

      return available_cues

    else:
      return None 

# JS: moved the cue generation further down
#def cue_generation():
 # if st.session_state.actions[-1] == 'cue':

if 'messages' not in st.session_state:
  st.session_state.messages = []
  
if 'results' not in st.session_state:
  st.session_state.results = {'results': False, 'results_print': False}
  
if 'actions' not in st.session_state:
  st.session_state.actions = [""]
  
if 'counters' not in st.session_state:
  st.session_state.counters = {"letter_count": 0, "word_count": 0}
  
if 'is_helpful' not in st.session_state:
  st.session_state.is_helpful = {'ask':False}
  
if 'descriptions' not in st.session_state:
  st.session_state.descriptions = []
  
st.title("You name it! 🗣")

# JS: would remove Simon by some neutral avatar
with st.chat_message('user'):
  st.write("Hey assistant!")
  
bot = st.chat_message('assistant')
bot.write("Hello human! Wanna practice naming some words?")

#for showing history of messages
for message in st.session_state.messages:
  if message['role'] == 'user':
    with st.chat_message(message['role']):
      st.markdown(message['content'])
  else:
    with st.chat_message(message['role']):
      st.markdown(message['content'])
  
#display user message in chat message container
prompt = get_text() 
if prompt:
  #JS: would replace Simon by some neutral character 
  with st.chat_message('user'):
    st.markdown(prompt)
  #add to history
  st.session_state.messages.append({'role': 'user', 'content': prompt})
  #TODO: replace it with zero-shot classifier
  yes = ['yes', 'again', 'Yes', 'sure', 'new word', 'yes!', 'yep', 'yeah']
  if prompt in yes:
    write_bot("Please describe your word!")
  elif prompt == 'It is similar to the best place on earth':
    write_bot("Great! Let me think what it could be...")
    time.sleep(3)
    write_bot("Do you mean Saarland?")
  #if previously we asked to give a prompt
  elif (st.session_state.messages[-2]['content'] == "Please describe your word!") & (st.session_state.messages[-1]['content'] != "no"):
    write_bot("Great! Let me think what it could be...")
    st.session_state.descriptions.append(prompt)
    st.session_state.results['results'] = return_top_k(st.session_state.descriptions[-1])
    st.session_state.results['results_print'] = dict(zip(range(1, 11), st.session_state.results['results']))
    write_bot("I think I have some ideas. Do you want to see my guesses or do you want a cue?")
    st.session_state.actions.append("result")

if st.session_state.actions[-1] == "result":
  col1, col2, col3, col4, col5 = st.columns(5)
  with col1:
    a1 = st.button('Results', key=10)
  with col2:
    a2 = st.button('Cue', key=11)
  if a1:
    write_bot("Here are my guesses about your word:")
    st.write(st.session_state.results['results_print'])
    time.sleep(1)
    write_bot('Does it help you remember the word?', remember=False)
    st.session_state.is_helpful['ask'] = True
  elif a2:
    #write_bot(f'The first letter is {st.session_state.results["results"][0][0]}.')
    #time.sleep(1)
    st.session_state.actions.append('cue')
    #cue_generation()
    #write_bot('Does it help you remember the word?', remember=False)
    #st.session_state.is_helpful['ask'] = True
    
if st.session_state.is_helpful['ask']:
    ask_if_helped()

if st.session_state.actions[-1] == 'cue':
  guessed = False
  write_bot('What do you want to see?', remember=False, blink=False)

  while guessed == False:
    # JS
    word_count = st.session_state.counters["word_count"]
    target = st.session_state.results["results"][word_count]
    
    col1, col2, col3, col4, col5 = st.columns(5)
  

    with col1:
      b1 = st.button("Next letter", key="1")
    with col2:
      b2 = st.button("Related words")
    with col3:
      b3 = st.button("Next word", key="2")
    with col4:
      b4 = st.button("All words", key="3")

    # JS
    #if get_available_cues(target):
     # avail_cues = get_available_cues(target)
      #cues_buttons = {cue_type: st.button(cue_type) for cue_type in avail_cues}

    b5 = st.button("I remembered the word!", key="4", type='primary')
    b6 = st.button("Exit", key="5", type='primary')
    new = st.button('Play again', key=64, 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"]
      if letter_count < len(target):
        write_bot(f'The word starts with {st.session_state.results["results"][word_count][:letter_count]}. \n Does this help you remember the word?', remember=False)
        #ask_if_helped()
        st.session_state.is_helpful['ask'] = True
      else:
        write_bot(f'This is my predicted word: "{target}". Does this match your query?')
        #ask_if_helped()
        st.session_state.is_helpful['ask'] = True

    elif b2:
      rels = return_top_k(st.session_state.descriptions[-1], word=target, rels=True)
      write_bot(f'Here are words that are related to your word: {", ".join(rels)}. \n Does this help you remember the word?', remember=False)
      #ask_if_helped()
      st.session_state.is_helpful['ask'] = True

    elif b3:
      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)
      if letter_count < len(target):
        write_bot(f'The next word starts with {st.session_state.results["results"][word_count][:letter_count]}. \n Does this help you remember the word?', remember=False)
        #ask_if_helped()
        st.session_state.is_helpful['ask'] = True
      else:
        write_bot(f'This is my predicted word: "{target}". Does this match your query?')
        #ask_if_helped()
        st.session_state.is_helpful['ask'] = True

    #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 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 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 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 b4:
      write_bot(f"Here are all my guesses about your word: {st.session_state.results['results_print']}")

    elif b5:
      write_bot("Yay! I am happy I could be of help!")
      st.session_state.counters["word_count"] = 0
      st.session_state.counters["letter_count"] = 0
      new = st.button('Play again', key=63)
      if new:
        write_bot("Please describe your word!")
      guessed = True
    
      break
    
    elif b6:
      write_bot("I am sorry I couldn't help you this time. See you soon!")
      st.session_state.counters["word_count"] = 0
      st.session_state.counters["letter_count"] = 0
    st.session_state.actions.append('cue')
    
    if new:
      write_bot("Please describe your word!")
      st.session_state.counters["word_count"] = 0
      st.session_state.counters["letter_count"] = 0

    break

  # elif prompt == 'results':
  #   st.text("results")
  #   st.write("results")
  #   st.session_state.actions.append({'result': True})
  #   st.write(st.session_state.actions)
  #     with st.chat_message('user'):
  #       custom_response = "Results"
  #       st.markdown(custom_response)
  #     st.session_state.messages.append({'role': 'user', 'content': custom_response})

  #     with st.chat_message('assistant'):
  #       message_placeholder = st.empty()
  #       response = f"Here are my guesses about your word: {result_print}"
  #       message_placeholder.markdown(response + "|")
  #     st.session_state.messages.append({'role': 'assistant', 'content': response})
  #   elif st.button('Cue'):
  #     response = "Cue"
  #     with st.chat_message('user'):
  #       st.markdown(response)
  #     st.session_state.messages.append({'role': 'user', 'content': response})
  #     text = f'The first letter is {result[0][0]}.'
  #     bot.write(text)
  #     st.session_state.messages.append({'role': 'assistant', 'content': text})
  #     letter_count = 1
  #     word_count = 0
  # elif prompt == 'Results':
  #   with st.chat_message('assistant'):
  #       message_placeholder = st.empty()
  #       response = f"Here are my guesses about your word: {result_print}"
  #       message_placeholder.markdown(response + "|")
  #   st.session_state.messages.append({'role': 'assistant', 'content': response})

  # #if you don't wanna practice word naming
  # else:  
  #     with st.chat_message('assistant'):
  #       message_placeholder = st.empty()
  #       response = "See you next time!"
  #       message_placeholder.markdown(response + "|")
  #     st.session_state.messages.append({'role': 'assistant', 'content': response})
      

  
      # if st.button('Results'):
      #   bot.write("Here are my guesses about your word:")
      #   bot.write(result_print)
      # elif st.button('Cue'):
      #   bot.write(f'The first letter is {result[0][0]}.')
      #   letter_count = 1
      #   word_count = 0
      #   answer = st.chat_input('Does it help you remember the word? Type yes or no')
      #   if answer == "no":
      #     bot.write("What do you want to see?")
      #     if st.button('Next letter'):
      #       letter_count += 1
      #       bot.write(f'The word starts with {result[word_count][:letter_count]}')
      #     elif st.button('Next word'):
      #       letter_count = 1
      #       bot.write(f'The next word starts with {result[word_count][:letter_count]}')
      #       word_count += 1
      #     elif st.button('All words'):
      #       bot.write("Here are all my guesses about your word:")
      #       bot.write(result_print)
      #   bot.write("Does this help you remember your word?")
      #   answer = st.chat_input('Type yes/no/exit')
      #   if answer == 'Exit':
      #     st.write("I am sorry I couldn't help you. See you next time!")

      
  #write down assistant's responses
  #response = f'Echo: {prompt}' #echoes prompt
  # with st.chat_message('assistant'):
  #   message_placeholder = st.empty()
  #   full_response = "yeee"
  #   #here insert the loop with the model answers (for response in...)
  #   #this to imitate a cursor
  #   message_placeholder.markdown(full_response + "|")
    
  # #add to history
  # st.session_state.messages.append({'role': 'assistant', 'content': full_response})
  
  
  
##TODO: a button to delete history
# if prompt == 'Yes':
#   bot.write("Great! Please describe the word you have in mind.")
#   sent = st.chat_input('Description of your word')


# # adding the text that will show in the text box as default
# default_value = "Type the description of the word you have in mind!"

# sent = st.text_area("Text", default_value, height = 50)
# result = return_top_k(sent)
# result = ['animal', 'monster', 'creature', 'bird', 'cat', 'human', 'dog', 'spider', 'alien', 'meow']
# result = return_top_k(sent)
# result_print = dict(zip(range(1, 11), result))

# if st.button('Results'):
#   st.write("Here are my guesses about your word:")
#   st.write(result_print)
# elif st.button('Cue'):
#   st.write(f'The first letter is {result[0][0]}.')
#   letter_count = 1
#   word_count = 0
#   answer = st.text_area("Text", 'Does it help you remember the word? Type yes or no', height = 50)
#   if answer == 'No':
#     while answer == 'No':
#       option = st.selectbox(
#       'What do you want to see?',
#       ('Next letter', 'Next word', 'All words'))
#       if option == 'Next letter':
#         letter_count += 1
#         st.write(f'The word starts with {result[word_count][:letter_count]}')
#       elif option == 'Next word':
#         letter_count = 1
#         st.write(f'The next word starts with {result[word_count][:letter_count]}')
#         word_count += 1
#       else:
#         st.write("Here are all my guesses about your word:")
#         st.write(result_print)
#       answer = st.selectbox(
#         'Does it help you remember the word??',
#         ('Yes', 'No', 'Exit'))
#       if answer == 'Exit':
#         st.write("I am sorry I couldn't help you. See you next time!")
#         break
#   else:
#      st.write("I am happy I could be of help!")
# else:
#   st.write('Do you want to see my guesses or do you want a cue?')


#2

# option = st.selectbox(
#     'Do you want to see my guesses or do you want a cue?',
#     ('Results', 'Cue'))

# st.write('You selected:', option)

# if option == 'Results':
#   st.write("Here are my guesses about your word:")
#   st.write(result_print)
# elif option == 'Cue':
#   st.write(f'The first letter is {result[0][0]}.')
#   letter_count = 1
#   word_count = 0
#   answer = st.selectbox(
#     'Does it help you remember the word??',
#     ('Yes', 'No'))
#   if answer == 'No':
#     while answer == 'No':
#       option = st.selectbox(
#       'What do you want to see?',
#       ('Next letter', 'Next word', 'All words'))
#       if option == 'Next letter':
#         letter_count += 1
#         st.write(f'The word starts with {result[word_count][:letter_count]}')
#       elif option == 'Next word':
#         letter_count = 1
#         st.write(f'The next word starts with {result[word_count][:letter_count]}')
#         word_count += 1
#       else:
#         st.write("Here are all my guesses about your word:")
#         st.write(result_print)
#       answer = st.selectbox(
#         'Does it help you remember the word??',
#         ('Yes', 'No', 'Exit'))
#       if answer == 'Exit':
#         st.write("I am sorry I couldn't help you. See you next time!")
#         break
#   else:
#     st.write("I am happy I could be of help!")