Use zerogpu

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

app.py

@@ -1,3 +1,4 @@
+import spaces
 import os
 import json
 import gradio as gr
@@ -23,6 +24,7 @@ def get_model(model_name: str = None):
 
     if _MODEL.get(model_name) is None:
         _MODEL[model_name] = GLiNER.from_pretrained(model_name, cache_dir=_CACHE_DIR)
+        _MODEL[model_name].to("cuda")
 
     return _MODEL[model_name]
 
@@ -34,15 +36,20 @@ def get_country(country_name: str):
         return None
 
 
-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]]:
+@spaces.GPU
+def predict_entities(model_name: str, query: str, labels: Union[str, list], threshold: float = 0.3, nested_ner: bool = False, model_name: str = None):
     model = get_model(model_name)
 
     if isinstance(labels, str):
         labels = [i.strip() for i in labels.split(",")]
 
-    _entities = model.predict_entities(query, labels, threshold=threshold)
+    return model.predict_entities(query, labels, threshold=threshold, flat_ner=not nested_ner)
+
+
+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)
 
     for entity in _entities:
         if entity["label"] == "country":

requirements.txt

@@ -1,4 +1,5 @@
 gliner
 pycountry
 scipy==1.12
-gradio
+gradio
+spaces