From 81501403e907f0304e2b1c5e0fadf7f6361f9662 Mon Sep 17 00:00:00 2001
From: Prashanth R <shanth@google.com>
Date: Fri, 29 Mar 2024 10:34:05 -0700
Subject: [PATCH] NL Server: use small spacy model instead of large (#4079)

Per SO (https://stackoverflow.com/a/57337084) and spacy website
(https://spacy.io/models/en), the accuracy difference for
Parts-Of-Speech task (for verb detection) between large vs. small is
very small. The verb test sanity passes.

This should cut down the memory usage per NL server a bunch, since the
small model is apparently just 10MB (vs. 780MB).

I see that heap-usage drops from 970MB to 230MB.

Also, drop a super edge case handling for test environment which loads
SentenceTransformer always.
---
 nl_server/gcs.py                | 18 +-----------------
 nl_server/loader.py             |  2 +-
 nl_server/nl_attribute_model.py |  4 ++--
 nl_server/requirements.txt      |  4 ++--
 nl_server/tests/verb_test.py    |  8 ++------
 5 files changed, 8 insertions(+), 28 deletions(-)

diff --git a/nl_server/gcs.py b/nl_server/gcs.py
index 8f936bcd84..5c0bfa69c4 100644
--- a/nl_server/gcs.py
+++ b/nl_server/gcs.py
@@ -82,23 +82,7 @@ def download_model_folder(model_folder: str) -> str:
   # Only download if needed.
   model_path = os.path.join(directory, model_folder)
   if os.path.exists(model_path):
-    if os.environ.get('FLASK_ENV') not in [
-        'local', 'test', 'integration_test', 'webdriver'
-    ]:
-      # If a production or production-like enrivonment,
-      # just return the model_path.
-      return model_path
-
-    # Check if this path can still be loaded as a Sentence Transformer
-    # model. If not, delete it and download anew.
-    try:
-      _ = SentenceTransformer(model_path)
-      return model_path
-    except:
-      print(f"Could not load the model from ({model_path}).")
-      print("Deleting this path and re-downloading.")
-      shutil.rmtree(model_path)
-      assert (not os.path.exists(model_path))
+    return model_path
 
   print(
       f"Model ({model_folder}) was either not previously downloaded or cannot successfully be loaded. Downloading to: {model_path}"
diff --git a/nl_server/loader.py b/nl_server/loader.py
index 7ec1d7b3bf..701aa9bac6 100644
--- a/nl_server/loader.py
+++ b/nl_server/loader.py
@@ -14,7 +14,7 @@
 
 import logging
 import os
-from typing import Any, Dict
+from typing import Dict
 
 from diskcache import Cache
 from flask import Flask
diff --git a/nl_server/nl_attribute_model.py b/nl_server/nl_attribute_model.py
index 00a6209dfd..a9596328fb 100644
--- a/nl_server/nl_attribute_model.py
+++ b/nl_server/nl_attribute_model.py
@@ -15,13 +15,13 @@
 
 from typing import List
 
-import en_core_web_lg
+import en_core_web_sm
 
 
 class NLAttributeModel:
 
   def __init__(self) -> None:
-    self.spacy_model_ = en_core_web_lg.load()
+    self.spacy_model_ = en_core_web_sm.load()
 
   def detect_verbs(self, query: str) -> List[str]:
     try:
diff --git a/nl_server/requirements.txt b/nl_server/requirements.txt
index 1a1c10f37f..b97feccb84 100644
--- a/nl_server/requirements.txt
+++ b/nl_server/requirements.txt
@@ -15,5 +15,5 @@ spacy==3.5.0
 Werkzeug==3.0.1
 # Downloading the named-entity recognition (NER) library spacy and the large EN model
 # using the guidelines here: https://spacy.io/usage/models#production
--f https://github.com/explosion/spacy-models/releases/download/en_core_web_lg-3.5.0/en_core_web_lg-3.5.0-py3-none-any.whl
-en_core_web_lg==3.5.0
\ No newline at end of file
+-f https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.5.0/en_core_web_sm-3.5.0-py3-none-any.whl
+en_core_web_sm==3.5.0
diff --git a/nl_server/tests/verb_test.py b/nl_server/tests/verb_test.py
index 4bee51a569..c582d525bf 100644
--- a/nl_server/tests/verb_test.py
+++ b/nl_server/tests/verb_test.py
@@ -51,13 +51,9 @@ def setUpClass(cls) -> None:
           ['give']
       ],
       [('Elaborate  how  your  roles  in  management  accounting  covering  – '
-        'planning,  organizing  and directing,  and  controlling  can  assist'
+        'planning,  organizing  and directing,  and  controlling  can  assist '
         'the  above  organization  in  achieving  their  goals  and objectives.'
-       ),
-       [
-           'Elaborate', 'covering', 'organizing', 'directing', 'controlling',
-           'achieving'
-       ]],
+       ), ['covering', 'organizing', 'controlling', 'assist', 'achieving']],
       ['How to write scholarship essay', ['write']]
   ])
   def test_verb_detection(self, query_str, expected):