Update app.py

app.py

@@ -1,6 +1,6 @@
-import os 
+import os
 import PyPDF2
-import pandas as pd 
+import pandas as pd
 import warnings
 import re
 from transformers import DPRContextEncoder, DPRContextEncoderTokenizer
@@ -8,115 +8,123 @@ from transformers import DPRQuestionEncoder, DPRQuestionEncoderTokenizer
 import torch
 import gradio as gr
 from typing import Union
-from datasets import Dataset
+import numpy as np
+from cassandra.cluster import Cluster
+from cassandra.auth import PlainTextAuthProvider
+from dotenv import load_dotenv, find_dotenv
+
 warnings.filterwarnings("ignore")
 
+# Load environment variables
+load_dotenv(find_dotenv())
+ASTRADB_TOKEN = os.getenv("ASTRADB_TOKEN")
+ASTRADB_API_ENDPOINT = os.getenv("ASTRADB_API_ENDPOINT")
 
+# AstraDB connection setup using token and endpoint
+auth_provider = PlainTextAuthProvider(username="token", password=ASTRADB_TOKEN)
+cluster = Cluster([ASTRADB_API_ENDPOINT], auth_provider=auth_provider)
+session = cluster.connect("your_keyspace_name")
 
+# Load DPR models and tokenizers
 ctx_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
 ctx_tokenizer = DPRContextEncoderTokenizer.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
 
-
 q_encoder = DPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
 q_tokenizer = DPRQuestionEncoderTokenizer.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
 
-
-
-def process_pdfs(parent_dir: Union[str,list]):
-    """ processess the PDF files and returns a dataframe with the text of each page in a
-    different line"""  # XD
-    # creating a pdf file object
-    df = pd.DataFrame(columns = ["title","text"])
-    if type(parent_dir) == str :
+def process_pdfs(parent_dir: Union[str, list]):
+    """Processes the PDF files and returns a dataframe with the text of each page in a different line."""
+    df = pd.DataFrame(columns=["title", "text"])
+    if type(parent_dir) == str:
         parent_dir = [parent_dir]
     for file_path in parent_dir:
-        if ".pdf" not in file_path : # skip non pdf files
+        if ".pdf" not in file_path:  # Skip non-pdf files
             raise Exception("only pdf files are supported")
-        # creating a pdf file object
         pdfFileObj = open(file_path, 'rb')
-        
-        # creating a pdf reader object
         pdfReader = PyPDF2.PdfReader(pdfFileObj)
-        # printing number of pages in pdf file
         num_pages = len(pdfReader.pages)
-        for i in range(num_pages) : 
+        for i in range(num_pages):
             pageObj = pdfReader.pages[i]
-            # extracting text from page
-            txt = pageObj.extract_text()
-            txt = txt.replace("\n","") # strip return to line
-            txt = txt.replace("\t","") # strip tabs
-            txt = re.sub(r" +"," ",txt) # strip extra space
-            # 512 is related to the positional encoding "facebook/dpr-ctx_encoder-single-nq-base" model
+            txt = pageObj.extract_text().replace("\n", "").replace("\t", "")
+            txt = re.sub(r" +", " ", txt)  # Strip extra space
             file_name = file_path.split("/")[-1]
-            if len(txt) < 512 :
-                new_data = pd.DataFrame([[f"{file_name}-page-{i}",txt]],columns=["title","text"])
-                df = pd.concat([df,new_data],ignore_index=True)
-            else :
-                while len(txt) > 512 :
-                    new_data = pd.DataFrame([[f"{file_name}-page-{i}",txt[:512]]],columns=["title","text"])
-                    df = pd.concat([df,new_data],ignore_index=True)
+            if len(txt) < 512:
+                new_data = pd.DataFrame([[f"{file_name}-page-{i}", txt]], columns=["title", "text"])
+                df = pd.concat([df, new_data], ignore_index=True)
+            else:
+                while len(txt) > 512:
+                    new_data = pd.DataFrame([[f"{file_name}-page-{i}", txt[:512]]], columns=["title", "text"])
+                    df = pd.concat([df, new_data], ignore_index=True)
                     txt = txt[512:]
-        
-        # closing the pdf file object
         pdfFileObj.close()
     return df
 
-def process(example):
-  """process the bathces of the dataset and returns the embeddings"""
-  try : 
-    tokens = ctx_tokenizer(example["text"], return_tensors="pt")
-    embed = ctx_encoder(**tokens)[0][0].detach().numpy()
-    return {'embeddings': embed}
-  except Exception as e:
-    raise Exception(f"error in process: {e}")
-
 def process_dataset(df):
-    """processess the dataframe and returns a dataset variable"""
-    if len(df) == 0 : 
+    """Processes the dataframe and stores embeddings in AstraDB."""
+    if len(df) == 0:
         raise Exception("empty pdf files, or can't read text from them")
-    ds = Dataset.from_pandas(df)
-    ds = ds.map(process) 
-    ds.add_faiss_index(column='embeddings') # add faiss index
-    return ds
-
-def search(query, ds, k=3):
-    """searches the query in the dataset and returns the k most similar"""
-    try :
+
+    for _, row in df.iterrows():
+        title = row['title']
+        text = row['text']
+        tokens = ctx_tokenizer(text, return_tensors="pt")
+        embed = ctx_encoder(**tokens)[0][0].detach().numpy().tolist()
+
+        query = "INSERT INTO your_table_name (title, text, embeddings) VALUES (%s, %s, %s)"
+        session.execute(query, (title, text, embed))
+
+    return df
+
+def search(query, k=3):
+    """Searches the query in the database and returns the k most similar."""
+    try:
         tokens = q_tokenizer(query, return_tensors="pt")
-        query_embed = q_encoder(**tokens)[0][0].detach().numpy()
-        scores, retrieved_examples = ds.get_nearest_examples("embeddings", query_embed, k=k)
-        out = f"""**title** : {retrieved_examples["title"][0]},\ncontent: {retrieved_examples["text"][0]}\n\n\n**similar resources:** {retrieved_examples["title"]}
+        query_embed = q_encoder(**tokens)[0][0].detach().numpy().tolist()
+
+        # Perform vector search in AstraDB
+        query = """
+        SELECT title, text, embeddings
+        FROM your_table_name
+        ORDER BY embeddings ANN OF %s LIMIT %s
+        """
+        rows = session.execute(query, (query_embed, k))
+
+        retrieved_examples = []
+        for row in rows:
+            retrieved_examples.append({
+                "title": row.title,
+                "text": row.text,
+                "embeddings": np.array(row.embeddings)
+            })
+
+        out = f"""**title** : {retrieved_examples[0]["title"]},\ncontent: {retrieved_examples[0]["text"]}\n\n\n**similar resources:** {[example["title"] for example in retrieved_examples]}
         """
     except Exception as e:
         out = f"error in search: {e}"
     return out
-    
-def predict(query,file_paths, k=3):
-    """predicts the most similar files to the query"""
-    try :
+
+def predict(query, file_paths, k=3):
+    """Predicts the most similar files to the query."""
+    try:
         df = process_pdfs(file_paths)
-        ds = process_dataset(df)
-        out =  search(query,ds,k=k)
+        process_dataset(df)
+        out = search(query, k=k)
     except Exception as e:
         out = f"error in predict: {e}"
     return out
 
-with gr.Blocks() as demo : 
+# Gradio interface
+with gr.Blocks() as demo:
     gr.Markdown("<h1 style='text-align: center'> PDF Search Engine </h1>")
     with gr.Row():
         with gr.Column(): 
-            files = gr.Files(label="Upload PDFs",type="filepath",file_count="multiple")
+            files = gr.Files(label="Upload PDFs", type="filepath", file_count="multiple")
             query = gr.Text(label="query")
-            with gr.Accordion("number of references",open=False):
-                k = gr.Number(value=3,show_label=False,precision=0,minimum=1,container=False)
+            with gr.Accordion("number of references", open=False):
+                k = gr.Number(value=3, show_label=False, precision=0, minimum=1, container=False)
             button = gr.Button("search")
         with gr.Column():
             output = gr.Markdown(label="output")
-    button.click(predict, [query,files,k],outputs=output)
+    button.click(predict, [query, files, k], outputs=output)
 
 demo.launch()
-
-
-
-
-