app.py

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


@st.cache_resource
def get_models():
  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")

  st.write("The assistant is loaded and ready to use!")
  return model, tokenizer

model, tokenizer = get_models()

def return_top_k(sentence, k=10):
  
  if sentence[-1] != ".":
    sentence = sentence + "."

  inputs = [f"Descripton : {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]
    
    for pred in predictions:
      if (len(pred) < 2) | (pred in sentence.split()):
        predictions.pop(predictions.index(pred))  
     
  return predictions[:10]

  
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": 1, "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!")

with st.chat_message('user', avatar='simon.jpg'):
  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'], avatar='simon.jpg'):
      st.markdown(message['content'])
  else:
    with st.chat_message(message['role']):
      st.markdown(message['content'])

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!")
  elif n:
    st.session_state.actions.append('cue')
  elif new:
    write_bot("Please describe your word!")
    # st.session_state.is_helpful['ask'] = False

if st.session_state.actions[-1] == "result":
  a1 = st.button('Results', key=10)
  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)
    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'] == True:
    ask_if_helped()


if st.session_state.actions[-1] == 'cue':
  guessed = False
  write_bot('What do you want to see?', remember=False, blink=False)
  b1 = st.button("Next letter", key="1")
  b2 = st.button("Next word", key="2")
  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')
  while guessed == False:
    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 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
  
#display user message in chat message container
prompt = get_text() 
if prompt:
  with st.chat_message('user', avatar='simon.jpg'):
    st.markdown(prompt)
  #add to history
  st.session_state.messages.append({'role': 'user', 'content': prompt})
  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")

    

  # 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!")