Update app.py

app.py

@@ -1,5 +1,6 @@
 
 
+
 import streamlit as st
 from PyPDF2 import PdfReader
 import requests
@@ -10,14 +11,14 @@ from transformers import AutoTokenizer, AutoModel
 import torch
 from rank_bm25 import BM25Okapi
 import json
-from transformers import DPRContextEncoder, DPRContextEncoderTokenizerFast, DPRQuestionEncoder, DPRQuestionEncoderTokenizerFast
+import time
+from transformers import pipeline
 
 my_token = os.getenv('my_repo_token')
 # Use Mistral API for serverless architecture
 API_URL_MISTRAL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.3"
 headers = {"Authorization": f"Bearer {my_token}"}
 
-# Query Mistral API
 def query_mistral(payload):
     response = requests.post(API_URL_MISTRAL, headers=headers, json=payload)
     try:
@@ -30,124 +31,102 @@ def query_mistral(payload):
     except json.JSONDecodeError:
         return "Error: Failed to parse response"
 
-# Generate response from Mistral
 def generate_response(prompt):
-    payload = {"inputs": prompt, "parameters": {"max_length": 200, "temperature": 0.7}}
+    prompt += "\nIf the answer is not in the provided text, say 'I don't know'."
+    payload = {"inputs": prompt, "parameters": {"max_length": 200, "temperature": 0.3}}
     return query_mistral(payload)
 
-# Dense Retrieval - FAISS Embeddings
 def get_embeddings(texts, model_name='sentence-transformers/all-MiniLM-L6-v2'):
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModel.from_pretrained(model_name)
-    
-    texts = [str(text) for text in texts if text]  # Ensure valid text input
     inputs = tokenizer(texts, return_tensors='pt', padding=True, truncation=True)
-
     with torch.no_grad():
         outputs = model(**inputs)
-    
     embeddings = outputs.last_hidden_state.mean(dim=1).cpu().numpy()
     return embeddings
 
-# Sparse Retrieval - BM25
-def find_most_relevant_context_bm25(contexts, question):
-    tokenized_corpus = [doc.split() for doc in contexts]
-    bm25 = BM25Okapi(tokenized_corpus)
-    tokenized_question = question.split()
-    return bm25.get_top_n(tokenized_question, contexts, n=min(3, len(contexts)))
-
-# Dense Retrieval - FAISS
 def find_most_relevant_context_faiss(contexts, question, model_name='sentence-transformers/all-MiniLM-L6-v2'):
     all_texts = [question] + contexts
-    all_texts = [str(text) for text in all_texts if text]  # Ensure valid input
-    
-    if not all_texts:
-        return []
-
     embeddings = get_embeddings(all_texts, model_name=model_name)
     question_embedding = embeddings[0]
     context_embeddings = embeddings[1:]
-
     dimension = context_embeddings.shape[1]
     index = faiss.IndexFlatL2(dimension)
     index.add(context_embeddings)
-
     _, indices = index.search(question_embedding.reshape(1, -1), min(3, len(context_embeddings)))
     return [contexts[idx] for idx in indices[0] if idx < len(contexts)]
 
-# **ColBERT Retrieval - Contextualized Word-Level Matching**
-def find_most_relevant_context_colbert(contexts, question):
-    # Load DPR question and context encoders (ColBERT-like approach)
-    question_encoder = DPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
-    context_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
-    question_tokenizer = DPRQuestionEncoderTokenizerFast.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
-    context_tokenizer = DPRContextEncoderTokenizerFast.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
-
-    # Encode question
-    question_inputs = question_tokenizer(question, return_tensors="pt", truncation=True, padding=True)
-    with torch.no_grad():
-        question_embedding = question_encoder(**question_inputs).pooler_output
+def find_most_relevant_context_bm25(contexts, question):
+    tokenized_corpus = [doc.split() for doc in contexts]
+    bm25 = BM25Okapi(tokenized_corpus)
+    tokenized_question = question.split()
+    top_docs = bm25.get_top_n(tokenized_question, contexts, n=min(3, len(contexts)))
+    return top_docs
 
-    # Encode each context
-    context_embeddings = []
+def rerank_results(contexts, question):
+    scores = []
     for context in contexts:
-        context_inputs = context_tokenizer(context, return_tensors="pt", truncation=True, padding=True)
-        with torch.no_grad():
-            context_embeddings.append(context_encoder(**context_inputs).pooler_output)
-    
-    context_embeddings = torch.stack(context_embeddings)
-
-    # Compute similarity scores
-    scores = torch.matmul(question_embedding, context_embeddings.T).squeeze(0)
-    top_k_indices = scores.topk(min(3, len(contexts))).indices.tolist()
-
-    return [contexts[idx] for idx in top_k_indices]
-
-# Hybrid Retrieval - FAISS + BM25 + ColBERT
-def hybrid_search(contexts, question):
-    faiss_results = find_most_relevant_context_faiss(contexts, question)
-    bm25_results = find_most_relevant_context_bm25(contexts, question)
-    colbert_results = find_most_relevant_context_colbert(contexts, question)
-    
-    # Combine results, ensuring uniqueness
-    combined_results = list(set(faiss_results + bm25_results + colbert_results))
-    return combined_results if combined_results else ["No relevant context found."]
+        response = generate_response(f"Does this text answer the question: '{question}'? Text: {context}")
+        if "yes" in response.lower():
+            scores.append((context, 1))
+        else:
+            scores.append((context, 0))
+    ranked_contexts = [context for context, score in scores if score > 0]
+    return ranked_contexts if ranked_contexts else ["I don't know."]
+
+def is_response_confident(response):
+    low_confidence_phrases = [
+        "i'm not sure", "i think", "possibly", "maybe", "it seems", "likely",
+        "i guess", "i assume", "it's possible", "as far as i know"
+    ]
+    return not any(phrase in response.lower() for phrase in low_confidence_phrases)
+
+def validate_response(prompt):
+    response1 = generate_response(prompt)
+    response2 = generate_response(prompt)
+    if response1.strip().lower() != response2.strip().lower():
+        return "I am not confident about the answer. Please refer to the document."
+    return response1
 
-# Generate answer based on retrieved context
 def answer_question_from_pdf(pdf_text, question):
-    return generate_response(f"{pdf_text} The Question is: {question} Provide the answer with max length of about 100 words.")
+    if not pdf_text.strip():
+        return "I could not find relevant information in the document."
+    prompt = f"{pdf_text} The Question is: {question} Provide the answer with max length of about 100 words."
+    response = validate_response(prompt)
+    if not is_response_confident(response):
+        return "I'm not sure about this. Please refer to the original document."
+    return response
 
-# Extract text from PDF
 def extract_text_from_pdf(pdf_file):
     pdf_reader = PdfReader(pdf_file)
     pdf_arr = [pdf_reader.pages[page_num].extract_text() for page_num in range(len(pdf_reader.pages)) if pdf_reader.pages[page_num].extract_text()]
     return pdf_arr if pdf_arr else ["No text extracted from the PDF"]
 
+def hybrid_search(contexts, question):
+    faiss_results = find_most_relevant_context_faiss(contexts, question)
+    bm25_results = find_most_relevant_context_bm25(contexts, question)
+    combined_results = list(set(faiss_results + bm25_results))
+    if not combined_results:
+        return ["I don't know."]
+    return rerank_results(combined_results, question)
+
 # Streamlit chatbot UI
-st.title("PDF Chatbot (Serverless) - Hybrid RAG (FAISS + BM25 + ColBERT)")
+st.title("PDF Chatbot (Serverless)")
 
 uploaded_file = st.file_uploader("Upload a PDF", type="pdf")
 
 if uploaded_file is not None:
     pdf_arr = extract_text_from_pdf(uploaded_file)
-    st.write("📄 **PDF Uploaded Successfully. Start chatting!**")
-
+    st.write("PDF Uploaded Successfully. Start chatting!")
     if "chat_history" not in st.session_state:
         st.session_state.chat_history = []
-    
     question = st.text_input("Ask a question about the PDF")
     if st.button("Send") and question:
-        combined_results = hybrid_search(pdf_arr, question)  # Hybrid Retrieval
-        
-        # Generate response using the best retrieved contexts
-        response = answer_question_from_pdf(" ".join(combined_results), question)
-
-        # Store the conversation history
+        combined_results = hybrid_search(pdf_arr, question)
+        response = answer_question_from_pdf(" ".join(combined_results), question) if combined_results else "No relevant context found."
         st.session_state.chat_history.append((question, response))
-
-    # Display chat history
     for q, r in st.session_state.chat_history:
         st.write(f"**You:** {q}")
         st.write(f"**Bot:** {r}")
 else:
-    st.write("📥 **Please upload a PDF file to start the chat.**")
+    st.write("Please upload a PDF file to start the chat.")