preloading models brute force

app.py

@@ -2,48 +2,51 @@ 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']
+available_models = ['2019',
+                    '2020']
 
-models = {}
-tokenizers = {}
+model_2019 = AutoModel.from_pretrained('cardiffnlp/twitter-roberta-base-2019-90m')
+tokenizers_2019 = AutoTokenizer.from_pretrained('cardiffnlp/twitter-roberta-base-2019-90m')
+embedding_matrix_2019 = model_2019.embeddings.word_embeddings.weight
+embedding_matrix_2019 = embedding_matrix_2019.detach().numpy()
+knn_model_2019 = NearestNeighbors(n_neighbors=500,
+                        metric='cosine',
+                        algorithm='auto',
+                        n_jobs=3)                
+nbrs_2019 = knn_model_2019.fit(embedding_matrix_2019)
+distances_2019, indices_2019 = nbrs_2019.kneighbors(embedding_matrix_2019)
 
-for MODEL in available_models:
-    models[MODEL] = AutoModel.from_pretrained(MODEL)
-    tokenizers[MODEL] = AutoTokenizer.from_pretrained(MODEL)
 
-
-def topk_model(MODEL):
-    # MODEL = "cardiffnlp/twitter-roberta-base-jun2022"
-    # 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,
-                            metric='cosine',
-                            algorithm='auto',
-                            n_jobs=3)
-                    
-    nbrs = knn_model.fit(embedding_matrix)
-
-    distances, indices = nbrs.kneighbors(embedding_matrix)
-
-    return distances,indices,tokenizers[MODEL]
+model_2020 = AutoModel.from_pretrained('cardiffnlp/twitter-roberta-base-jun2020')
+tokenizers_2020 = AutoTokenizer.from_pretrained('cardiffnlp/twitter-roberta-base-jun2020')
+embedding_matrix_2020 = model_2020.embeddings.word_embeddings.weight
+embedding_matrix_2020 = embedding_matrix_2020.detach().numpy()
+knn_model_2020 = NearestNeighbors(n_neighbors=500,
+                        metric='cosine',
+                        algorithm='auto',
+                        n_jobs=3)              
+nbrs_2020 = knn_model_2020.fit(embedding_matrix_2020)
+distances_2020, indices_2020 = nbrs_2020.kneighbors(embedding_matrix_2020)
 
 
 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
+    if model == '2019':
+        index = tokenizers_2019.encode(f'{word}')
+        for i in indices_2019[index[1]]:
+            outs.append(tokenizers_2019.decode(i))
+            print(tokenizers_2019.decode(i))
+        return outs
+    
+    if model == '2020':
+        index = tokenizers_2020.encode(f'{word}')
+        for i in indices_2020[index[1]]:
+            outs.append(tokenizers_2020.decode(i))
+            print(tokenizers_2020.decode(i))
+        return outs
 
 # with gr.Blocks() as demo:
 #     gr.Markdown(f" # {title}")