adding two models to check speed

app.py

@@ -3,28 +3,34 @@ from transformers import AutoModel, AutoTokenizer
 from sklearn.neighbors import NearestNeighbors
 
 
+models = ['cardiffnlp/twitter-roberta-base-jun2022',
+ 'cardiffnlp/twitter-roberta-base-2019-90m']
 
-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()
+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)
+    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)
+    distances, indices = nbrs.kneighbors(embedding_matrix)
 
+    return distances,indices,tokenizer
 
-title = "How does a word's meaning change with time?"
 
+title = "How does a word's meaning change with time?"
 
-def topk(word):
+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))
@@ -42,6 +48,6 @@ with gr.Blocks() as demo:
     with gr.Row():
         greet_btn = gr.Button("Compute")
     with gr.Row():
-        greet_btn.click(fn=topk, inputs=[word], outputs=gr.outputs.Textbox())
+        greet_btn.click(fn=topk, inputs=[word,gr.Dropdown(models)], outputs=gr.outputs.Textbox())
 
 demo.launch()