Refactor

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

app.py

@@ -3,6 +3,7 @@ import os
 import json
 import gradio as gr
 import pycountry
+import torch
 from datetime import datetime
 from typing import Dict, Union
 from gliner import GLiNER
@@ -17,16 +18,20 @@ print(f"Cache directory: {_CACHE_DIR}")
 
 
 def get_model(model_name: str = None):
+    start = datetime.now()
+
     if model_name is None:
         model_name = "urchade/gliner_base"
 
     global _MODEL
 
     if _MODEL.get(model_name) is None:
-        start = datetime.now()
         _MODEL[model_name] = GLiNER.from_pretrained(model_name, cache_dir=_CACHE_DIR)
-        _MODEL[model_name].to("cuda")
-        print(f"{datetime.now()} :: get_model :: {datetime.now() - start}")
+
+    if torch.cuda.is_available() and _MODEL[model_name].device.type.startswith("cuda"):
+        _MODEL[model_name] = _MODEL[model_name].to("cuda")
+
+    print(f"{datetime.now()} :: get_model :: {datetime.now() - start}")
 
     return _MODEL[model_name]
 
@@ -38,9 +43,10 @@ def get_country(country_name: str):
         return None
 
 
-@spaces.GPU
-def predict_entities(model, query: str, labels: Union[str, list], threshold: float = 0.3, nested_ner: bool = False):
+@spaces.GPU(enable_queue=True)
+def predict_entities(model_name: str, query: str, labels: Union[str, list], threshold: float = 0.3, nested_ner: bool = False):
     start = datetime.now()
+    model = get_model(model_name)
 
     if isinstance(labels, str):
         labels = [i.strip() for i in labels.split(",")]
@@ -55,8 +61,7 @@ def predict_entities(model, query: str, labels: Union[str, list], threshold: flo
 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 = []
-    model = get_model(model_name)
-    _entities = predict_entities(model=model, query=query, labels=labels, threshold=threshold, nested_ner=nested_ner)
+    _entities = predict_entities(model_name=model_name, query=query, labels=labels, threshold=threshold, nested_ner=nested_ner)
 
     for entity in _entities:
         if entity["label"] == "country":