Update app.py

app.py

@@ -3,7 +3,8 @@ from fastapi.staticfiles import StaticFiles
 from fastapi.responses import FileResponse
 from pydantic import BaseModel
 from fastapi.middleware.cors import CORSMiddleware
-from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModel
+import torch
 import lancedb
 import os
 
@@ -25,7 +26,14 @@ class BookRequest(BaseModel):
     query: str
     limit: int = 5
 
-model = SentenceTransformer("all-MiniLM-L6-v2")
+tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
+model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
+
+def encode_query(text):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    return outputs.last_hidden_state.mean(dim=1).squeeze().numpy()
 
 db = lancedb.connect(
     uri=os.getenv("LANCEDB_URI", "your-lancedb-uri"),
@@ -53,7 +61,8 @@ def search_books(book_request: BookRequest):
         raise HTTPException(status_code=400, detail="Query string is required")
 
     # Get query embeddings
-    query_embedding = model.encode(book_request.query)
+    query_embedding = encode_query(book_request.query) 
+    #query_embedding = model.encode(book_request.query)
 
     table = db.open_table(os.getenv("LANCEDB_TABLE", "book_db"))
     results = (table.search(query_embedding)