Initialize models

Signed-off-by: Aivin V. Solatorio <avsolatorio@gmail.com>

app.py

@@ -11,8 +11,10 @@ from gliner import GLiNER
 
 _MODEL = {}
 _CACHE_DIR = os.environ.get("CACHE_DIR", None)
+THRESHOLD = 0.3
 LABELS = ["country", "year", "statistical indicator", "geographic region"]
 QUERY = "gdp, co2 emissions, and mortality rate of the philippines vs. south asia in 2024"
+MODELS = ["urchade/gliner_base", "urchade/gliner_medium-v2.1"]
 
 print(f"Cache directory: {_CACHE_DIR}")
 
@@ -36,6 +38,13 @@ def get_model(model_name: str = None):
     return _MODEL[model_name]
 
 
+# Initialize model here.
+print("Initializing models...")
+for model_name in MODELS:
+    model = get_model(model_name=model_name)
+    model.predict_entities(QUERY, LABELS, threshold=THRESHOLD)
+
+
 def get_country(country_name: str):
     try:
         return pycountry.countries.search_fuzzy(country_name)
@@ -43,7 +52,7 @@ def get_country(country_name: str):
         return None
 
 
-@spaces.GPU(enable_queue=True, duration=3)
+@spaces.GPU(enable_queue=True, duration=5)
 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)
@@ -99,7 +108,7 @@ with gr.Blocks(title="GLiNER-query-parser") as demo:
     )
     with gr.Row() as row:
         model_name = gr.Radio(
-            choices=["urchade/gliner_base", "urchade/gliner_medium-v2.1"],
+            choices=MODELS,
             value="urchade/gliner_base",
             label="Model",
         )
@@ -112,7 +121,7 @@ with gr.Blocks(title="GLiNER-query-parser") as demo:
         threshold = gr.Slider(
             0,
             1,
-            value=0.3,
+            value=THRESHOLD,
             step=0.01,
             label="Threshold",
             info="Lower threshold may extract more false-positive entities from the query.",