Update app.py

app.py

@@ -4,6 +4,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 from PyPDF2 import PdfReader
 import gradio as gr
 from datasets import Dataset, load_from_disk
+from sentence_transformers import SentenceTransformer
 
 # Extract text from PDF
 def extract_text_from_pdf(pdf_path):
@@ -19,13 +20,19 @@ model_name = "scb10x/llama-3-typhoon-v1.5x-8b-instruct"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
+# Load a sentence transformer model for embedding generation
+embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+
 # Extract text from the provided PDF
 pdf_path = "/home/user/app/TOPF 2564.pdf"  # Ensure this path is correct
 pdf_text = extract_text_from_pdf(pdf_path)
 passages = [{"title": "", "text": line} for line in pdf_text.split('\n') if line.strip()]
 
-# Create a Dataset
-dataset = Dataset.from_dict({"title": [p["title"] for p in passages], "text": [p["text"] for p in passages]})
+# Convert text to embeddings
+embeddings = embedding_model.encode([passage["text"] for passage in passages])
+
+# Create a Dataset with embeddings
+dataset = Dataset.from_dict({"title": [p["title"] for p in passages], "text": [p["text"] for p in passages], "embeddings": embeddings.tolist()})
 
 # Save the dataset and create an index in the current working directory
 dataset_path = "/home/user/app/rag_document_dataset"
@@ -37,15 +44,15 @@ os.makedirs(index_path, exist_ok=True)
 
 # Save the dataset to disk and create an index
 dataset.save_to_disk(dataset_path)
-dataset.load_from_disk(dataset_path).add_faiss_index(column="text").save(index_path)
+dataset = load_from_disk(dataset_path)
+dataset.add_faiss_index(column="embeddings").save(index_path)
 
 # Custom retriever
 def retrieve(query):
     # Use FAISS index to retrieve relevant passages
-    query_embedding = tokenizer(query, return_tensors="pt")["input_ids"]
-    # Perform retrieval (this is a placeholder, actual retrieval code will be more complex)
-    # retrieved_passages = faiss_search(query_embedding)
-    retrieved_passages = " ".join([passage['text'] for passage in passages])  # Simplified for demo
+    query_embedding = embedding_model.encode([query])
+    scores, samples = dataset.get_nearest_examples("embeddings", query_embedding, k=5)
+    retrieved_passages = " ".join([sample["text"] for sample in samples])
     return retrieved_passages
 
 # Define the chat function