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 | |
break | |