Update app.py

app.py

@@ -1,10 +1,30 @@
 import streamlit as st
 import chromadb
 import torch
-from transformers import pipeline
+from transformers import pipeline, AutoModel, AutoTokenizer
+import numpy as np
 from PyPDF2 import PdfReader
 import os
 
+# Load sentence transformer model for embeddings
+def load_embedding_model():
+    model = AutoModel.from_pretrained("cross-encoder/qnli-electra-base")
+    tokenizer = AutoTokenizer.from_pretrained("cross-encoder/qnli-electra-base")
+    return model, tokenizer
+
+# Generate embeddings for text
+def generate_embedding(model, tokenizer, text):
+    # Tokenize the text
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
+    
+    # Generate embeddings
+    with torch.no_grad():
+        outputs = model(**inputs)
+        
+    # Use the last hidden state as embedding
+    embeddings = outputs.last_hidden_state.mean(dim=1).squeeze().numpy()
+    return embeddings
+
 # Initialize Hugging Face pipeline for question answering
 def load_qa_pipeline():
     return pipeline("question-answering", model="deepset/roberta-base-squad2")
@@ -24,8 +44,8 @@ def split_text_into_chunks(text, chunk_size=500, overlap=100):
         chunks.append(text[i:i+chunk_size])
     return chunks
 
-# Create ChromaDB collection
-def create_chroma_collection(chunks):
+# Create ChromaDB collection with embeddings
+def create_chroma_collection(chunks, model, tokenizer):
     # Use persistent client to avoid memory issues
     client = chromadb.PersistentClient(path="./chroma_db")
     
@@ -35,19 +55,27 @@ def create_chroma_collection(chunks):
     # Create collection
     collection = client.create_collection(name=collection_name)
     
-    # Add chunks to collection
+    # Add chunks to collection with embeddings
     for i, chunk in enumerate(chunks):
+        # Generate embedding for the chunk
+        embedding = generate_embedding(model, tokenizer, chunk)
+        
         collection.add(
             ids=[f"chunk_{i}"],
-            documents=[chunk]
+            documents=[chunk],
+            embeddings=[embedding.tolist()]
         )
     
     return client, collection, collection_name
 
 # Retrieve most relevant context
-def retrieve_context(collection, question, top_k=3):
+def retrieve_context(collection, question, model, tokenizer, top_k=3):
+    # Generate embedding for the question
+    question_embedding = generate_embedding(model, tokenizer, question)
+    
+    # Query the collection
     results = collection.query(
-        query_texts=[question],
+        query_embeddings=[question_embedding.tolist()],
         n_results=top_k
     )
     return results['documents'][0]
@@ -56,6 +84,9 @@ def retrieve_context(collection, question, top_k=3):
 def main():
     st.title("PDF Question Answering App")
     
+    # Load embedding model
+    embedding_model, tokenizer = load_embedding_model()
+    
     # File uploader
     uploaded_file = st.file_uploader("Upload PDF", type=['pdf'])
     
@@ -75,11 +106,15 @@ def main():
                 # Split text into chunks
                 text_chunks = split_text_into_chunks(pdf_text)
                 
-                # Create ChromaDB collection
-                client, collection, collection_name = create_chroma_collection(text_chunks)
+                # Create ChromaDB collection with embeddings
+                client, collection, collection_name = create_chroma_collection(
+                    text_chunks, embedding_model, tokenizer
+                )
                 
                 # Retrieve context
-                contexts = retrieve_context(collection, question)
+                contexts = retrieve_context(
+                    collection, question, embedding_model, tokenizer
+                )
                 
                 # Prepare answers
                 answers = []