Add option for local testing.

Author: alfredo-f

.gitignore

@@ -48,4 +48,8 @@ var/
 # pytest
 *pytest_cache
 
+# Credentials
 key_openai.txt
+
+# Models saved locally
+models/

README.md

@@ -36,6 +36,50 @@ To run REMO, you will need the following:
 2. Interact with the API using a REST client or web browser: `http://localhost:8000`
 
 
+## Models
+
+### Embedding Model
+
+REMO currently uses the
+Universal Sentence Encoder v5 for generating embeddings.
+
+#### Loading from TensorFlow Hub
+
+This is the default option.
+
+When 
+
+```python
+ARE_YOU_TESTING__LOAD_MODEL_LOCAL = False
+```
+
+in file `utils.py`, the model is loaded from TensorFlow Hub.
+
+#### Loading from a local file
+
+Downloading the model from TensorFlow Hub every time you need to spin up
+the microservice would be expensive and time-consuming.
+
+1. Download the `.tar.gz` file from
+   TensorFlow Hub: https://tfhub.dev/google/universal-sentence-encoder-large/5
+
+   ![img.png](docs/images/embedding_local_1.png)
+
+2. Extract the file to the folder
+   ```
+   models/universal-sentence-encoder-large_5/
+   ```
+   with
+   ```shell
+   tar -xvzf universal-sentence-encoder-large_5.tar.gz
+   ```
+
+3. Set
+   ```python
+   ARE_YOU_TESTING__LOAD_MODEL_LOCAL = True
+   ```
+   in file `utils.py`.
+
 ## API Endpoints
 
 - **POST /add_message**: Add a new message to REMO. Speaker, timestamp, and content required.

docs/images/embedding_local_1.png

@@ -1,6 +1,7 @@
 from fastapi import FastAPI
 import utils
 import os
+import uvicorn
 
 app = FastAPI()
 root_folder = os.getcwd()
@@ -48,3 +49,7 @@ async def maintain_tree():
     utils.maintain_tree(root_folder)
 
     return {"detail": "Tree maintenance completed"}
+
+
+if __name__ == '__main__':
+    uvicorn.run(app, host='0.0.0.0', port=8000)

remo.py

@@ -1,4 +1,6 @@
 import os
+from pathlib import Path
+
 import yaml
 import shutil
 import openai
@@ -7,12 +9,24 @@
 from typing import Dict, Any, List
 from sklearn.metrics.pairwise import cosine_similarity
 from sklearn.cluster import KMeans
+import tensorflow as tf
 import tensorflow_hub as hub
 
 
-embedding_model = hub.load(
-    "https://tfhub.dev/google/universal-sentence-encoder-large/5"
-)
+ARE_YOU_TESTING__LOAD_MODEL_LOCAL = True
+
+ROOT_REPO_PATH = Path().parent.absolute()
+
+if ARE_YOU_TESTING__LOAD_MODEL_LOCAL:
+    embedding_model = tf.saved_model.load(
+        ROOT_REPO_PATH
+        / "models/universal-sentence-encoder-large_5"
+    )
+
+else:
+    embedding_model = hub.load(
+        "https://tfhub.dev/google/universal-sentence-encoder-large/5"
+    )
 
 
 def open_file(filepath):