zeroGPU

interfaces/cap_minor_media.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import spaces
 
 import os
 import torch
@@ -56,6 +57,7 @@ def check_huggingface_path(checkpoint_path: str):
 def build_huggingface_path(language: str, domain: str):
     return ("poltextlab/xlm-roberta-large-pooled-cap-media", "poltextlab/xlm-roberta-large-pooled-cap-minor-v3")
 
+@spaces.GPU
 def predict(text, major_model_id, minor_model_id, tokenizer_id, HF_TOKEN=None):
     device = torch.device("cpu")
 

requirements.txt

@@ -8,3 +8,4 @@ huspacy
 numpy==1.26.4
 safetensors
 huggingface_hub
+spaces