adding lance

app.py

@@ -5,28 +5,29 @@ from datasets import load_dataset
 from sentence_transformers import SentenceTransformer
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import os
+import lancedb
 
 os.environ["HF_TOKEN"] = os.getenv("auth")
-dataset = load_dataset("ariG23498/pis-blogs-chunked")
+
+db = lancedb.connect("embedding_dataset")
+tbl = db.open_table("my_table")
+
 embedding_model = SentenceTransformer(model_name_or_path="all-mpnet-base-v2", device="cuda")
 tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
 model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it", torch_dtype=torch.bfloat16, device_map="auto")
 
 @spaces.GPU(duration=300)
 def process_query(query):
-    text_embeddings = embedding_model.encode(dataset["train"]["text"])
     query_embedding = embedding_model.encode(query)
+    search_hits = tbl.search(query_embedding).metric("cosine").limit(5).to_list()
 
-    similarity_scores = embedding_model.similarity(query_embedding, text_embeddings)
-    top_indices = (-similarity_scores).argsort()[0][:5]
-    
-    context = dataset["train"]["text"][top_indices[0]]
-    url = dataset["train"]["url"][top_indices[0]]
+    context = search_hits[0]["text"]
+    url = search_hits[0]["url"]
     
     input_text = (
-        f"Based on the context provided, '{context}', how would"
-        f"you address the user's query regarding '{query}'? Please"
-        " provide a detailed and contextually relevant response."
+        f"You are being provided a query: {query}"
+        f"YOu are being provided context to the query: {context}"
+        "Please provide a detailed and contextually relevant response."
     )
     
     input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")

requirements.txt

@@ -2,4 +2,5 @@ accelerate
 transformers
 gradio
 sentence-transformers
-datasets
+datasets
+lancedb