gpu

modular_graph_and_candidates.py

@@ -40,6 +40,7 @@ from typing import Dict, List, Set, Tuple
 from sentence_transformers import SentenceTransformer, util
 from tqdm import tqdm
 import numpy as np
+import spaces
 
 # ────────────────────────────────────────────────────────────────────────────────
 # CONFIG
@@ -93,6 +94,7 @@ def similarity_clusters(bags: Dict[str, List[Set[str]]], thr: float) -> Dict[Tup
             out[(m1, m2)] = s
     return out
 
+@spaces.GPU
 def embedding_similarity_clusters(models_root: Path, missing: List[str], thr: float) -> Dict[Tuple[str, str], float]:
     model = SentenceTransformer("nomic-ai/nomic-embed-code")
     model.max_seq_length = 4096  # truncate overly long modeling files