Update main.py

main.py

@@ -24,13 +24,11 @@ data.reset_index(inplace=True)
 # Create a FAISS index for fast similarity search
 metric = faiss.METRIC_INNER_PRODUCT
 vectors = numpy.stack(data["embedding"].tolist(), axis=0)
-gpu_index = faiss.IndexFlatL2(len(data["embedding"][0]))
-# res = faiss.StandardGpuResources()  # use a single GPU
-# gpu_index = faiss.index_cpu_to_gpu(res, 0, index)
-gpu_index.metric_type = metric
+index = faiss.IndexFlatL2(len(data["embedding"][0]))
+index.metric_type = metric
 faiss.normalize_L2(vectors)
-gpu_index.train(vectors)
-gpu_index.add(vectors)
+index.train(vectors)
+index.add(vectors)
 
 # Load the model for later use in embeddings
 model = sentence_transformers.SentenceTransformer("allenai-specter")
@@ -39,7 +37,7 @@ model = sentence_transformers.SentenceTransformer("allenai-specter")
 def search(query: str, k: int) -> tuple[str]:
     query = numpy.expand_dims(model.encode(query), axis=0)
     faiss.normalize_L2(query)
-    D, I = gpu_index.search(query, k)
+    D, I = index.search(query, k)
     top_five = data.loc[I[0]]
 
     search_results = "You are an AI assistant who delights in helping people" \
@@ -78,6 +76,7 @@ def postprocess(response: str, bypass_from_preprocessing: str) -> str:
     """Applies a postprocessing step to the LLM's response before the user receives it"""
     return response + bypass_from_preprocessing
 
+@spaces.GPU
 def predict(message: str, history: list[str]) -> str:
     """This function is responsible for crafting a response"""