preloading models

app.py

@@ -2,15 +2,22 @@ import gradio as gr
 from transformers import AutoModel, AutoTokenizer
 from sklearn.neighbors import NearestNeighbors
 
+available_models = ['cardiffnlp/twitter-roberta-base-2019-90m',
+                    'cardiffnlp/twitter-roberta-base-jun2020']
+
+models = {}
+tokenizers = {}
+
+for MODEL in available_models:
+    models[MODEL] = AutoModel.from_pretrained(MODEL)
+    tokenizers[MODEL] = AutoTokenizer.from_pretrained(MODEL)
 
-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
+    # model = AutoModel.from_pretrained(MODEL)
+    # tokenizer = AutoTokenizer.from_pretrained(MODEL)
+    embedding_matrix = models[MODEL].embeddings.word_embeddings.weight
     embedding_matrix = embedding_matrix.detach().numpy()
 
     knn_model = NearestNeighbors(n_neighbors=500,
@@ -22,7 +29,7 @@ def topk_model(MODEL):
 
     distances, indices = nbrs.kneighbors(embedding_matrix)
 
-    return distances,indices,tokenizer
+    return distances,indices,tokenizers[MODEL]
 
 
 title = "How does a word's meaning change with time?"
@@ -52,7 +59,7 @@ def topk(word,model):
 # demo.launch()
 
 interface = gr.Interface(fn=topk, 
-                        inputs=[gr.Textbox(label="Word"), gr.Dropdown(models)],
+                        inputs=[gr.Textbox(label="Word"), gr.Dropdown(available_models)],
                         outputs=gr.outputs.Textbox()
                         )
 interface.launch()