From ef935ec0c19349b6a7ae2f5b4e407ef35cc0064b Mon Sep 17 00:00:00 2001
From: Fraser <tomatoesneverlie@gmail.com>
Date: Fri, 31 Mar 2023 19:40:05 -0400
Subject: [PATCH] context sensitive dataset load

---
 .gitignore  |  1 +
 api/main.py | 45 ++++++++++++++++++++++++++++-----------------
 2 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/.gitignore b/.gitignore
index b8708f4..4f9ddb1 100644
--- a/.gitignore
+++ b/.gitignore
@@ -136,3 +136,4 @@ dmypy.json
 src/tmp.py
 
 .vercel/
+api/dataset.pkl
diff --git a/api/main.py b/api/main.py
index 00799a1..332da65 100644
--- a/api/main.py
+++ b/api/main.py
@@ -8,35 +8,46 @@ import openai
 import pickle
 import requests
 
+# ---------------------------------- env setup ---------------------------------
 if os.path.exists('.env'):
     from dotenv import load_dotenv
     load_dotenv()
 
-app = Flask(__name__)
-cors = CORS(app)
-app.config['CORS_HEADERS'] = 'Content-Type'
-
-
-# -------------------------------- general setup -------------------------------
-
-
-
 OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
 openai.api_key = OPENAI_API_KEY
 
-print('Downloading dataset...')
 
-url = os.environ.get('DATASET_URL')
-if url is None:
-    print('No dataset url provided.')
-    exit()
+# -------------------------------- load dataset --------------------------------
+
+
+if os.path.exists('dataset.pkl'):
+    print('Found dataset.pkl')
+    print('Loading dataset...')
+    with open('dataset.pkl', 'rb') as f:
+        dataset_dict = pickle.load(f)
+
+else:
+    print('No dataset.pkl found on disk.')
+    print('Downloading dataset...')
+
+    url = os.environ.get('DATASET_URL')
+    if url is None:
+        print('No dataset url provided.')
+        exit()
+
+    dataset_dict_bytes = requests.get(url).content
+    print('Unpacking dataset...')
+    dataset_dict = pickle.load(dataset_dict_bytes)
 
-dataset_dict_bytes = requests.get(url).content
-print('Unpacking dataset...')
-dataset_dict = pickle.loads(dataset_dict_bytes)
 print('Done!')
 
 
+# ---------------------------------- web setup ---------------------------------
+
+app = Flask(__name__)
+cors = CORS(app)
+app.config['CORS_HEADERS'] = 'Content-Type'
+
 
 # ------------------------------- semantic search ------------------------------