zerogpu decorator

app.py

@@ -1,13 +1,18 @@
 import gradio as gr
+import spaces
 import pandas as pd
 from typing import List, Dict
 from flow_judge import Hf, FlowJudge, EvalInput
 from flow_judge.metrics import CustomMetric, RubricItem
 
-try:
-    model = Hf(flash_attn=False)
-except Exception as e:
-    raise RuntimeError(f"Failed to initialize Hf Model: {e}")
+
+@spaces.GPU
+def load_model():
+    try:
+        model = Hf(flash_attn=False)
+        return model
+    except Exception as e:  
+        raise RuntimeError(f"Failed to initialize Hf Model: {e}")
 
 EXAMPLES = [
     {
@@ -33,8 +38,9 @@ def populate_fields(example_index: int):
         example["evaluation_criteria"],
         [[str(i), description] for i, description in enumerate(example["rubric"])]
     )
-
-def evaluate(task_inputs: pd.DataFrame, task_output: pd.DataFrame, evaluation_criteria: str, rubric: pd.DataFrame) -> tuple:
+    
+@spaces.GPU
+def evaluate(model, task_inputs: pd.DataFrame, task_output: pd.DataFrame, evaluation_criteria: str, rubric: pd.DataFrame) -> tuple:
     # Convert inputs to the expected format
     eval_input = EvalInput(
         inputs=[{row['Name']: row['Value']} for _, row in task_inputs.iterrows()],
@@ -94,6 +100,7 @@ def reset_evaluation_criteria():
     )
 
 with gr.Blocks() as demo:
+    model = load_model()
     with gr.Row():
         example_buttons = [gr.Button(f"{example['emoji']} Example {i+1}") for i, example in enumerate(EXAMPLES)]
 
@@ -177,7 +184,7 @@ with gr.Blocks() as demo:
 
     evaluate_btn.click(
         evaluate,
-        inputs=[task_inputs, task_output, evaluation_criteria, rubric],
+        inputs=[model,task_inputs, task_output, evaluation_criteria, rubric],
         outputs=[feedback, score]
     )