rachith's picture
adding two models to check speed
76e5451
raw
history blame
1.6 kB
import gradio as gr
from transformers import AutoModel, AutoTokenizer
from sklearn.neighbors import NearestNeighbors
models = ['cardiffnlp/twitter-roberta-base-jun2022',
'cardiffnlp/twitter-roberta-base-2019-90m']
def topk_model(MODEL):
# MODEL = "cardiffnlp/twitter-roberta-base-jun2022"
model = AutoModel.from_pretrained(MODEL)
tokenizer = AutoTokenizer.from_pretrained(MODEL)
embedding_matrix = model.embeddings.word_embeddings.weight
embedding_matrix = embedding_matrix.detach().numpy()
knn_model = NearestNeighbors(n_neighbors=500,
metric='cosine',
algorithm='auto',
n_jobs=3)
nbrs = knn_model.fit(embedding_matrix)
distances, indices = nbrs.kneighbors(embedding_matrix)
return distances,indices,tokenizer
title = "How does a word's meaning change with time?"
def topk(word,model):
outs = []
distances, indices, tokenizer = topk_model(model)
index = tokenizer.encode(f'{word}')
for i in indices[index[1]]:
outs.append(tokenizer.decode(i))
print(tokenizer.decode(i))
return outs
with gr.Blocks() as demo:
gr.Markdown(f" # {title}")
# gr.Markdown(f" ## {description1}")
# gr.Markdown(f"{description2}")
# gr.Markdown(f"{description3}")
with gr.Row():
word = gr.Textbox(label="Word")
with gr.Row():
greet_btn = gr.Button("Compute")
with gr.Row():
greet_btn.click(fn=topk, inputs=[word,gr.Dropdown(models)], outputs=gr.outputs.Textbox())
demo.launch()