Run model on GPU and add fp16 and trust_remote_code options

app.py

@@ -12,6 +12,8 @@ from pathlib import Path
 from dataclasses import dataclass
 from itertools import batched, chain
 import os
+import torch
+from sys import stderr
 
 
 class IndexParameters(TypedDict):
@@ -105,8 +107,8 @@ def get_env_var[T, U](
     return var
 
 
-def get_model(model_name: str, device: str) -> SentenceTransformer:
-    return SentenceTransformer(model_name, device=device)
+def get_model(model_name: str, trust_remote_code: bool) -> SentenceTransformer:
+    return SentenceTransformer(model_name, trust_remote_code=trust_remote_code)
 
 
 def get_index(dir: Path, search_time_s: float) -> Dataset:
@@ -207,13 +209,20 @@ def format_response(neighbors: list[Work], distances: list[float]) -> str:
 def main():
     # TODO: figure out some better defaults?
     model_name = get_env_var("MODEL_NAME", default="all-MiniLM-L6-v2")
+    trust_remote_code = get_env_var("TRUST_REMOTE_CODE", bool, default=False)
+    fp16 = get_env_var("FP16", bool, default=False)
     dir = get_env_var("DIR", Path, default=Path("index"))
     search_time_s = get_env_var("SEARCH_TIME_S", float, default=1)
     k = get_env_var("K", int, default=20)  # TODO: can't go higher than 20 yet
     mailto = get_env_var("MAILTO", str, None)
 
-    model = get_model(model_name, "cpu")
+    model = get_model(model_name, trust_remote_code)
     index = get_index(dir, search_time_s)
+    if torch.cuda.is_available():
+        model = model.half().cuda() if fp16 else model.bfloat16().cuda()
+        # TODO: if huggingface datasets exposes an fp16 gpu option, use it here
+    elif fp16:
+        print('warning: used "FP16" on CPU-only system, ignoring...', file=stderr)
 
     # function signature: (expanded tuple of input batches) -> tuple of output batches
     def search(query: list[str]) -> tuple[list[str]]:

requirements.txt

@@ -1,3 +1,3 @@
 sentence-transformers
-faiss-cpu
+faiss-gpu
 datasets