Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import torch | |
| from tqdm import tqdm | |
| from transformers import pipeline | |
| 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 | |
| def get_models_context(llama=False): | |
| st.write('Loading the model...') | |
| model = pipeline("fill-mask") | |
| st.write("_The assistant is loaded and ready to use! :tada:_") | |
| return model | |
| model_context = get_models_context() | |
| 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_context(sentence, 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 : "] | |
| output = model_context(sentence) | |
| output = [output[i]['token_str'].strip() for i in range(len(output))] | |
| return output | |
| # 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_context' 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_context.append({'role': 'assistant', 'content': full_response}) | |
| def ask_if_helped_context(): | |
| 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!") | |
| again = st.button('Play again') | |
| if again: | |
| write_bot("Please give a sentence using a <mask> instead of the word you have in mind!") | |
| st.session_state.is_helpful_context['ask'] = False | |
| elif n: | |
| st.session_state.actions_context.append('cue') | |
| st.session_state.is_helpful_context['ask'] = False | |
| #cue_generation() | |
| elif new: | |
| write_bot("Please give a sentence using a <mask> instead of the word you have in mind!") | |
| st.session_state.is_helpful_context['ask'] = False | |
| ## removed: if st.session_state.actions_context[-1] == "result": | |
| # JS | |
| # def get_related_words_llama(relation, target, device, num=5): | |
| # prompt_context = 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_context], 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_context[-1] == 'cue': | |
| if 'messages_context' not in st.session_state: | |
| st.session_state.messages_context = [] | |
| if 'results_context' not in st.session_state: | |
| st.session_state.results_context = {'results_context': False, 'results_context_print': False} | |
| if 'actions_context' not in st.session_state: | |
| st.session_state.actions_context = [""] | |
| if 'counter_context' not in st.session_state: | |
| st.session_state.counter_context = {"letter_count": 0, "word_count": 0} | |
| if 'is_helpful_context' not in st.session_state: | |
| st.session_state.is_helpful_context = {'ask':False} | |
| if 'descriptions_context' not in st.session_state: | |
| st.session_state.descriptions_context = [] | |
| 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_context | |
| for message in st.session_state.messages_context: | |
| 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_context = get_text() | |
| if prompt_context: | |
| #JS: would replace Simon by some neutral character | |
| with st.chat_message('user'): | |
| st.markdown(prompt_context) | |
| #add to history | |
| st.session_state.messages_context.append({'role': 'user', 'content': prompt_context}) | |
| #TODO: replace it with zero-shot classifier | |
| yes = ['yes', 'again', 'Yes', 'sure', 'new word', 'yes!', 'yep', 'yeah'] | |
| if prompt_context in yes: | |
| write_bot("Please give a sentence using a <mask> instead of the word you have in mind!") | |
| elif prompt_context == '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_context | |
| elif (st.session_state.messages_context[-2]['content'] == "Please give a sentence using a <mask> instead of the word you have in mind!") & (st.session_state.messages_context[-1]['content'] != "no"): | |
| write_bot("Great! Let me think what it could be...") | |
| st.session_state.descriptions_context.append(prompt_context) | |
| st.session_state.results_context['results_context'] = return_top_k_context(st.session_state.descriptions_context[-1]) | |
| st.session_state.results_context['results_context_print'] = dict(zip(range(1, len(st.session_state.results_context['results_context'])+1), st.session_state.results_context['results_context'])) | |
| 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_context.append("result") | |
| if st.session_state.actions_context[-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_context['results_context_print']) | |
| time.sleep(1) | |
| write_bot('Does it help you remember the word?', remember=False) | |
| st.session_state.is_helpful_context['ask'] = True | |
| elif a2: | |
| write_bot(f'The first letter is {st.session_state.results_context["results_context"][0][0]}.') | |
| time.sleep(1) | |
| # st.session_state.actions_context.append('cue') | |
| #cue_generation() | |
| write_bot('Does it help you remember the word?', remember=False) | |
| st.session_state.is_helpful_context['ask'] = True | |
| if st.session_state.is_helpful_context['ask'] == True: | |
| ask_if_helped_context() | |
| if st.session_state.actions_context[-1] == 'cue': | |
| guessed = False | |
| write_bot('What do you want to see?', remember=False, blink=False) | |
| 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') | |
| while guessed == False: | |
| # JS | |
| word_count = st.session_state.counter_context["word_count"] | |
| target = st.session_state.results_context["results_context"][word_count] | |
| if b1: | |
| st.session_state.counter_context["letter_count"] += 1 | |
| #word_count = st.session_state.counter_context["word_count"] | |
| letter_count = st.session_state.counter_context["letter_count"] | |
| if letter_count < len(target): | |
| write_bot(f'The word starts with {st.session_state.results_context["results_context"][word_count][:letter_count]}. \n Does this help you remember the word?', remember=False) | |
| #ask_if_helped_context() | |
| st.session_state.is_helpful_context['ask'] = True | |
| else: | |
| write_bot(f'This is my predicted word: "{target}". Does this match your query?') | |
| #ask_if_helped_context() | |
| st.session_state.is_helpful_context['ask'] = True | |
| elif b2: | |
| rels = return_top_k_context(st.session_state.descriptions_context[-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_context() | |
| st.session_state.is_helpful_context['ask'] = True | |
| elif b3: | |
| st.session_state.counter_context["letter_count"] = 1 | |
| letter_count = st.session_state.counter_context["letter_count"] | |
| st.session_state.counter_context["word_count"] += 1 | |
| word_count = st.session_state.counter_context["word_count"] | |
| #write_bot(f'The next word starts with {st.session_state.results_context["results_context"][word_count][:letter_count]}', remember=False) | |
| if letter_count < len(target): | |
| write_bot(f'The next word starts with {st.session_state.results_context["results_context"][word_count][:letter_count]}. \n Does this help you remember the word?', remember=False) | |
| #ask_if_helped_context() | |
| st.session_state.is_helpful_context['ask'] = True | |
| else: | |
| write_bot(f'This is my predicted word: "{target}". Does this match your query?') | |
| #ask_if_helped_context() | |
| st.session_state.is_helpful_context['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_context['results_context_print']}") | |
| st.session_state.is_helpful_context['ask'] = True | |
| elif b5: | |
| write_bot("Yay! I am happy I could be of help!") | |
| st.session_state.counter_context["word_count"] = 0 | |
| st.session_state.counter_context["letter_count"] = 0 | |
| new = st.button('Play again', key=63) | |
| if new: | |
| write_bot("Please give a sentence using a <mask> instead of the word you have in mind!") | |
| guessed = True | |
| break | |
| elif b6: | |
| write_bot("I am sorry I couldn't help you this time. See you soon!") | |
| st.session_state.counter_context["word_count"] = 0 | |
| st.session_state.counter_context["letter_count"] = 0 | |
| st.session_state.actions_context.append('cue') | |
| if new: | |
| write_bot("Please give a sentence using a <mask> instead of the word you have in mind!") | |
| st.session_state.counter_context["word_count"] = 0 | |
| st.session_state.counter_context["letter_count"] = 0 | |