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