Adding pytorch model to cloud and other minor changes

Author: mojavad

.gitignore

@@ -14,6 +14,9 @@ Icon
 # Keys
 keyfile.json
 
+# Pytorch Model
+models/*.bin
+
 # Thumbnails
 ._*
 

app.py

@@ -3,7 +3,7 @@
 from search import *
 from build import *
 from utils import docs2text, id2details
-from cloud_storage import test_file_exists, download_blob, upload_blob, pull_indices
+from cloud_storage import test_file_exists, download_blob, upload_blob, pull_indices, download_pytorch_model
 
 # from deeppavlov import build_model
 # dp_model = build_model('models/squad_torch_bert.json', download=True)
@@ -16,7 +16,7 @@
 # GCP test connections
 test_file_exists()
 download_blob("symptomizer_indices_bucket-1", "hello.txt", "test.txt")
-
+download_pytorch_model()
 pull_indices()
 
 bert_model = load_bert_model()

build.py

@@ -103,6 +103,12 @@
 #         else:
 #             yield collection.find(query,projection)[chunks[i-1]:chunks.stop]
 
+def upload_indices_and_vectors():
+    upload_blob("symptomizer_indices_bucket-1", "models/bert.index", "bert.index")
+    upload_blob("symptomizer_indices_bucket-1", "models/tfidf.index", "tfidf.index")
+    upload_blob("symptomizer_indices_bucket-1", "models/ids.joblib", "ids.joblib")
+    upload_blob("symptomizer_indices_bucket-1", "models/tfidf_model.joblib", "tfidf_model.joblib")
+    print("Completed Uploading indices to bucket")
 
 def build_tfidf_model(num_docs=2000, max_features=1500):
     print("Building TF-IDF model...")
@@ -125,16 +131,11 @@ def load_bert_model():
     print("Completed BERT model.")
     return model
 
-def upload_faiss():
-    upload_blob("symptomizer_indices_bucket-1", "models/bert.index", "bert.index")
-    upload_blob("symptomizer_indices_bucket-1", "models/tfidf.index", "tfidf.index")
-    upload_blob("symptomizer_indices_bucket-1", "models/ids.joblib", "ids.joblib")
-    print("Completed Uploading indices to bucket")
 
 def build_faiss(tfidf_model, bert_model):
     tr = tracker.SummaryTracker()
     print(f"Building indices ...")
-    c = collection.find().count()
+    c = 1500
     batch_size = 500
     encoder = None
     bert_index =  None
@@ -181,7 +182,7 @@ def build_faiss(tfidf_model, bert_model):
     faiss.write_index(tfidf_index,f"models/tfidf.index")
     dump(ids,'models/ids.joblib')
     print(f"Completed indices.")
-    upload_faiss()  
+    upload_indices_and_vectors()
     return [tfidf_index, bert_index]
 
 def load_faiss(tfidf_model, bert_model):

cloud_storage.py

@@ -2,6 +2,16 @@
 import os.path
 import os
 import json
+import wget
+
+def download_pytorch_model():
+    if(not os.path.exists("models/pytorch_model.bin")):
+        print("Pytorch QA model not detected. Downloading...")
+        url = 'https://storage.googleapis.com/symptomizer_model_bucket/pytorch_model.bin'
+        wget.download(url, 'models/pytorch_model.bin')
+        print("Download Complete.")
+    else:
+        print("Pytorch model exists. Skipping download.")
 
 def fix_json():
     with open('keyfile.json') as json_file:
@@ -11,11 +21,7 @@ def fix_json():
 
 
 def test_file_exists():
-    print("Key File Exists")
-    print(os.path.isfile("keyfile.json"))
-    
     f = open("keyfile.json", "r")
-    print(f.read())
     fix_json()
 
 def check_if_exists(bucket_name, file_name):
@@ -42,7 +48,7 @@ def upload_blob(bucket_name, source_file_name, destination_blob_name):
         )
     )
 
-# upload_blob("symptomizer_indices_bucket-1", "hello.txt", "hello.txt")
+# upload_blob("symptomizer_model_bucket", "models/pytorch_model.bin", "pytorch_model.bin")
 # print(check_if_exists("symptomizer_indices_bucket-1", "hello.txt"))
 # download_blob("symptomizer_indices_bucket-1", "hello.txt", "downloaded.txt")
 
@@ -76,5 +82,7 @@ def pull_indices():
         download_blob("symptomizer_indices_bucket-1", "tfidf.index", "models/tfidf.index")
         download_blob("symptomizer_indices_bucket-1", "bert.index", "models/bert.index")
         download_blob("symptomizer_indices_bucket-1", "ids.joblib", "models/ids.joblib")
+        download_blob("symptomizer_indices_bucket-1", "tfidf_model.joblib", "models/tfidf_model.joblib")
     else:
         print("No PULL_INDS env found. Rebuilding indices")
+

models/pytorch_model.bin

@@ -13,4 +13,5 @@ sentence-transformers
 pytorch_transformers
 dask[complete]
 Pympler
-google-cloud-storage
+google-cloud-storage
+wget