Add zeroGPU

app.py

@@ -14,11 +14,18 @@ from datasets import Dataset
 from datasets.search import FaissIndex
 import faiss
 from huggingface_hub import snapshot_download
+import numpy as np
+import numpy.typing as npt
 import gradio as gr
 import requests
 from sentence_transformers import SentenceTransformer
 import torch
 
+try:
+    import spaces
+except ImportError:
+    spaces = None
+
 
 class IndexParameters(TypedDict):
     recall: float  # in this case 10-recall@10
@@ -267,11 +274,14 @@ def main():
         print('warning: used "FP16" on CPU-only system, ignoring...', file=stderr)
     model.compile(mode="reduce-overhead")
 
+    def encode(query: list[str]) -> npt.NDArray[np.float16 | np.float32]:
+        return model.encode(query, prompt_name, normalize_embeddings=normalize)
+    if spaces:
+        encode = spaces.GPU(encode)
+
     # function signature: (expanded tuple of input batches) -> tuple of output batches
     def search(query: list[str]) -> tuple[list[str]]:
-        query_embedding = model.encode(
-            query, prompt_name, normalize_embeddings=normalize
-        )
+        query_embedding = encode(query)
         distances, faiss_ids = index.search_batch("embedding", query_embedding, k)
 
         faiss_ids_flat = list(chain(*faiss_ids))