Add time logs and duration

Signed-off-by: Aivin V. Solatorio <[email protected]>

app.py

@@ -36,20 +36,24 @@ def get_country(country_name: str):
         return None
 
 
-@spaces.GPU
-def predict_entities(model_name: str, query: str, labels: Union[str, list], threshold: float = 0.3, nested_ner: bool = False):
-    model = get_model(model_name)
+@spaces.GPU(duration=3)
+def predict_entities(model, query: str, labels: Union[str, list], threshold: float = 0.3, nested_ner: bool = False):
+    start = datetime.now()
 
     if isinstance(labels, str):
         labels = [i.strip() for i in labels.split(",")]
 
-    return model.predict_entities(query, labels, threshold=threshold, flat_ner=not nested_ner)
+    entities = model.predict_entities(query, labels, threshold=threshold, flat_ner=not nested_ner)
+
+    print(f"{datetime.now()} :: predict_entities :: {datetime.now() - start}")
+
+    return entities
 
 
 def parse_query(query: str, labels: Union[str, list], threshold: float = 0.3, nested_ner: bool = False, model_name: str = None) -> Dict[str, Union[str, list]]:
 
     entities = []
-    _entities = predict_entities(model_name=model_name, query=query, labels=labels, threshold=threshold, nested_ner=nested_ner)
+    _entities = predict_entities(model=get_model(model_name), query=query, labels=labels, threshold=threshold, nested_ner=nested_ner)
 
     for entity in _entities:
         if entity["label"] == "country":
@@ -61,7 +65,7 @@ def parse_query(query: str, labels: Union[str, list], threshold: float = 0.3, ne
             entities.append(entity)
 
     payload = {"query": query, "entities": entities}
-    print(f"{datetime.now()} :: {json.dumps(payload)}\n")
+    print(f"{datetime.now()} :: parse_query :: {json.dumps(payload)}\n")
 
     return payload