import gradio as gr
import math
import spacy
from datasets import load_dataset
from sentence_transformers import SentenceTransformer
from sentence_transformers import InputExample
from sentence_transformers import losses
from sentence_transformers import util
from transformers import pipeline
from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification
from transformers import TrainingArguments, Trainer
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
import numpy as np
import evaluate
import nltk
from nltk.corpus import stopwords
import subprocess
import sys
import random
# !pip install https://huggingface.co/spacy/en_core_web_sm/resolve/main/en_core_web_sm-any-py3-none-any.whl
subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'https://huggingface.co/spacy/en_core_web_sm/resolve/main/en_core_web_sm-any-py3-none-any.whl'])
# tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
model_base = "bert-analogies"
nltk.download('stopwords')
nlp = spacy.load("en_core_web_sm")
stops = stopwords.words("english")
ROMAN_CONSTANTS = (
( "", "I", "II", "III", "IV", "V", "VI", "VII", "VIII", "IX" ),
( "", "X", "XX", "XXX", "XL", "L", "LX", "LXX", "LXXX", "XC" ),
( "", "C", "CC", "CCC", "CD", "D", "DC", "DCC", "DCCC", "CM" ),
( "", "M", "MM", "MMM", "", "", "-", "", "", "" ),
( "", "i", "ii", "iii", "iv", "v", "vi", "vii", "viii", "ix" ),
( "", "x", "xx", "xxx", "xl", "l", "lx", "lxx", "lxxx", "xc" ),
( "", "c", "cc", "ccc", "cd", "d", "dc", "dcc", "dccc", "cm" ),
( "", "m", "mm", "mmm", "", "", "-", "", "", "" ),
)
# answer = "Pizza"
guesses = []
return_guesses = []
answer = "Moon"
word1 = "Black"
word2 = "White"
word3 = "Sun"
base_prompts = ["Sun is to Moon as ", "Black is to White as ", "Atom is to Element as",
"Athens is to Greece as ", "Cat is to Dog as ", "Robin is to Bird as",
"Hunger is to Ambition as "]
#Mean Pooling - Take attention mask into account for correct averaging
def mean_pooling(model_output, attention_mask):
token_embeddings = model_output['token_embeddings'] #First element of model_output contains all token embeddings
input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
def normalize(comment, lowercase, remove_stopwords):
if lowercase:
comment = comment.lower()
comment = nlp(comment)
lemmatized = list()
for word in comment:
lemma = word.lemma_.strip()
if lemma:
if not remove_stopwords or (remove_stopwords and lemma not in stops):
lemmatized.append(lemma)
return " ".join(lemmatized)
# def tokenize_function(examples):
# return tokenizer(examples["text"])
def compute_metrics(eval_pred):
logits, labels = eval_pred
predictions = np.argmax(logits, axis=-1)
metric = evaluate.load("accuracy")
return metric.compute(predictions=predictions, references=labels)
def get_model():
global model_base
model = SentenceTransformer(model_base)
gpu_available = torch.cuda.is_available()
device = torch.device("cuda" if gpu_available else "cpu")
model = model.to(device)
return model
def cosine_scores(model, sentence):
global word1
global word2
global word3
# sentence1 = f"{word1} is to {word2} as"
embeddings1 = model.encode(sentence, convert_to_tensor=True)
def embeddings(model, sentences):
gpu_available = torch.cuda.is_available()
device = torch.device("cuda" if gpu_available else "cpu")
# device = torch.device('cuda:0')
embeddings = model.encode(sentences)
global word1
global word2
global word3
global model_base
# Load model from HuggingFace Hub
tokenizer = AutoTokenizer.from_pretrained(model_base)
encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
# token_ids = tokenizer.encode(sentences, return_tensors='pt')
# blank_id = tokenizer.mask_token_id
# blank_id_idx = torch.where(encoded_input["input_ids"] == blank_id)[1]
encoded_input["input_ids"] = encoded_input["input_ids"].to(device)
encoded_input["attention_mask"] = encoded_input["attention_mask"].to(device)
encoded_input['token_type_ids'] = encoded_input['token_type_ids'].to(device)
encoded_input['input'] = {'input_ids':encoded_input['input_ids'], 'attention_mask':encoded_input['attention_mask']}
del encoded_input['input_ids']
del encoded_input['token_type_ids']
del encoded_input['attention_mask']
with torch.no_grad():
# output = model(encoded_input)
print(encoded_input)
model_output = model(**encoded_input)
# output = model(encoded_input_topk)
unmasker = pipeline('fill-mask', model=model_base)
guesses = unmasker(sentences)
print(guesses)
# Perform pooling
sentence_embeddings = mean_pooling(model_output, encoded_input['input']["attention_mask"])
# Normalize embeddings
sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)
potential_words = []
for guess in guesses:
temp_word = guess['token_str']
if temp_word[0].isalpha() and temp_word not in stops and temp_word not in ROMAN_CONSTANTS:
potential_words.append(guess['token_str'])
rand_index = random.randint(0, len(potential_words) - 1)
print("THE LENGTH OF POTENTIAL WORDS FOR", sentences, "IS", len(potential_words), "AND THE RANDOM INDEX CHOSEN IS", rand_index)
chosen_word = potential_words[rand_index]
return chosen_word
def random_word():
global model_base
with open(model_base + '/vocab.txt', 'r') as file:
line = ""
content = file.readlines()
length = len(content)
while line == "":
rand_line = random.randrange(0, length)
if content[rand_line][0].isalpha() and content[rand_line][:-1] not in stops and content[rand_line][:-1] not in ROMAN_CONSTANTS:
line = content[rand_line]
else:
print(f"{content[rand_line]} is not alpha or is a stop word")
# for num, aline in enumerate(file, 1997):
# if random.randrange(num) and aline.isalpha():
# continue
# # elif not aline.isalpha():
# line = aline
print(line)
return line[:-1]
def generate_prompt(model):
global word1
global word2
global word3
global answer
global base_prompts
word1 = random_word()
# word2 = random_word()
random_line = random.randint(0, len(base_prompts) - 1)
word2 = embeddings(model, f"{base_prompts[random_line]}{word1} is to [MASK].")
word3 = random_word()
sentence = f"{word1} is to {word2} as {word3} is to [MASK]."
print(sentence)
answer = embeddings(model, sentence)
print("ANSWER IS", answer)
return f"# {word1} is to {word2} as {word3} is to ___."
# cosine_scores(model, sentence)
def new_prompt(name):
model = get_model()
return generate_prompt(model)
def check_answer(guess:str):
global guesses
global answer
global return_guesses
global word1
global word2
global word3
model = get_model()
output = ""
protected_guess = guess
sentence = f"{word1} is to {word2} as [MASK] is to {guess}."
other_word = embeddings(model, sentence)
guesses.append(guess)
for guess in return_guesses:
output += (guess)
# output = output[:-1]
prompt = f"{word1} is to {word2} as {word3} is to ___."
# print("IS", protected_guess, "EQUAL TO", answer, ":", protected_guess.lower() == answer.lower())
if protected_guess.lower() == answer.lower():
return_guesses.append(f"- {protected_guess}: {word1} is to {word2} as {word3} is to {protected_guess}.
")
output += f"- {return_guesses[-1]}
"
new_prompt = generate_prompt(model)
return new_prompt, "Correct!", output
else:
return_guess = f"- {protected_guess}: {word1} is to {word2} as {other_word} is to {protected_guess}.
"
return_guesses.append(return_guess)
output += (return_guess)
return prompt, "Try again!", output
def main():
global word1
global word2
global word3
global answer
# answer = "Moon"
global guesses
# num_rows, data_type, value, example, embeddings = training()
# sent_embeddings = embeddings()
model = get_model()
generate_prompt(model)
prompt = f"# {word1} is to {word2} as {word3} is to ____"
print(prompt)
print("TESTING EMBEDDINGS")
with gr.Blocks() as iface:
mark_question = gr.Markdown(prompt)
with gr.Tab("Guess"):
text_input = gr.Textbox()
text_output = gr.Textbox()
text_button = gr.Button("Submit")
prompt_button = gr.Button("New Prompt")
with gr.Accordion("Open for previous guesses"):
text_guesses = gr.Markdown()
# with gr.Tab("Testing"):
# gr.Markdown(f"""The Embeddings are {sent_embeddings}.""")
text_button.click(check_answer, inputs=[text_input], outputs=[mark_question, text_output, text_guesses])
prompt_button.click(new_prompt, inputs=[], outputs=[mark_question])
# iface = gr.Interface(fn=greet, inputs="text", outputs="text")
iface.launch()
if __name__ == "__main__":
main()