Minor fixes

backend/query_llm.py

@@ -1,17 +1,16 @@
 import os
-
-import requests
 from os import getenv
 
+import gradio as gr
+import requests
 
 API_URL = getenv('API_URL')
 BEARER = getenv('BEARER')
 
-
 headers = {
-	"Authorization": f"Bearer {BEARER}",
-	"Content-Type": "application/json"
-}
+    "Authorization": f"Bearer {BEARER}",
+    "Content-Type": "application/json"
+    }
 
 
 def call_jais(payload):
@@ -26,7 +25,7 @@ def call_jais(payload):
 
 
 def generate(prompt: str):
-    payload = {'inputs': '', 'prompt':prompt}
+    payload = {'inputs': '', 'prompt': prompt}
     response = call_jais(payload)
     return response
 

backend/semantic_search.py

@@ -1,9 +1,10 @@
 import logging
-from pathlib import Path
 import time
+from pathlib import Path
 
 import lancedb
 from sentence_transformers import SentenceTransformer
+
 import spaces
 
 
@@ -17,7 +18,7 @@ start_time = time.perf_counter()
 proj_dir = Path(__file__).parents[1]
 
 # Log the time taken to load the QdrantDocumentStore
-db = lancedb.connect(proj_dir/"lancedb")
+db = lancedb.connect(proj_dir / "lancedb")
 tbl = db.open_table('arabic-wiki')
 lancedb_loading_time = time.perf_counter() - start_time
 logger.info(f"Time taken to load LanceDB: {lancedb_loading_time:.6f} seconds")
@@ -25,23 +26,35 @@ logger.info(f"Time taken to load LanceDB: {lancedb_loading_time:.6f} seconds")
 # Start the timer for loading the EmbeddingRetriever
 start_time = time.perf_counter()
 
-name="sentence-transformers/paraphrase-multilingual-minilm-l12-v2"
-st_model = SentenceTransformer(name, device='cuda')
+name = "sentence-transformers/paraphrase-multilingual-minilm-l12-v2"
+st_model_gpu = SentenceTransformer(name, device='mps')
+st_model_cpu = SentenceTransformer(name, device='cpu')
+
 
 # used for both training and querying
+def call_embed_func(query):
+    try:
+        return embed_func(query)
+    except:
+        logger.warning(f'Using CPU')
+        return st_model_cpu.encode(query)
+
+
 @spaces.GPU
 def embed_func(query):
-    return st_model.encode(query)
+    return st_model_gpu.encode(query)
+
 
 def vector_search(query_vector, top_k):
     return tbl.search(query_vector).limit(top_k).to_list()
 
+
 def retriever(query, top_k=3):
-    query_vector = embed_func(query)
+    query_vector = call_embed_func(query)
     documents = vector_search(query_vector, top_k)
     return documents
 
 
 # Log the time taken to load the EmbeddingRetriever
 retriever_loading_time = time.perf_counter() - start_time
-logger.info(f"Time taken to load EmbeddingRetriever: {retriever_loading_time:.6f} seconds")
+logger.info(f"Time taken to load EmbeddingRetriever: {retriever_loading_time:.6f} seconds")