feat: further optimize for gpu speed

app.py

@@ -8,13 +8,18 @@ from sentence_transformers import SentenceTransformer
 import gradio as gr
 from threading import Thread
 
-index = faiss.read_index("vector_db/index.faiss")
+cpu_index = faiss.read_index("vector_db/index.faiss")
+res = faiss.StandardGpuResources()
+index = faiss.index_cpu_to_gpu(res, 0, cpu_index)
+
 with open("vector_db/chunks.pkl", "rb") as f:
     chunks = pickle.load(f)
 with open("vector_db/metadata.pkl", "rb") as f:
     metadata_dict = pickle.load(f)
 
-ST = SentenceTransformer("BAAI/bge-large-en-v1.5")
+ST = SentenceTransformer("BAAI/bge-large-en-v1.5", device="cuda")
+ST.cuda()
+ST.compile()
 
 model_id = "nvidia/Llama-3.1-Nemotron-Nano-8B-v1"
 bnb = BitsAndBytesConfig(
@@ -39,18 +44,20 @@ SYS = (
     "Answer this Question:"
 )
 
+@spaces.GPU(duration=20)
 def retrieve(q, k=3):
-    emb = ST.encode(q)
-    D, I = index.search(np.array([emb], dtype="float32"), k)
+    emb = ST.encode(q, convert_to_tensor=True)
+    emb = emb.unsqueeze(0).to('cuda')
+    D, I = index.search(emb, k)
+    ids = I.cpu().numpy().tolist()[0]
     docs, file_sources = [], []
-    for i in I[0]:
+    for i in ids:
         chunk = chunks[i]
         meta = metadata_dict[i]
         docs.append({"title": chunk, "pages": chunk})
         file_sources.append(meta["source"])
     return docs, file_sources
 
-
 def make_prompt(q, docs, reasoning_mode):
     context = "\n\n".join(f"Title: {d['title']}\nPages: {d['pages']}" for d in docs)
     prompt = f"detailed thinking {reasoning_mode}\n"
@@ -59,7 +66,7 @@ def make_prompt(q, docs, reasoning_mode):
     prompt += f"Instruct: {SYS} {q} based on the following documents:\n{context}\nOutput:"
     return prompt
 
-@spaces.GPU()
+@spaces.GPU(duration=20)
 def qa_fn(question, reasoning_mode, top_k, temperature, max_tokens):
     docs, file_sources = retrieve(question, top_k)
     prompt = make_prompt(question, docs, reasoning_mode)[:8000]