Update app.py

app.py

@@ -1,19 +1,56 @@
 import streamlit as st
 import os
 import requests
+import re
 from langdetect import detect
 from PyPDF2 import PdfReader
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
 from sklearn.neighbors import NearestNeighbors
 import numpy as np
+from sentence_transformers import SentenceTransformer
+import faiss
+import hashlib
 
 # Load the Hugging Face token from environment variables
-huggingface_token = os.environ.get("Key2")  # Replace with your Hugging Face token
+huggingface_token = os.environ.get("Key2")
 
-# Function to query the Hugging Face API
+# Initialize Sentence Transformer model for better embeddings
+sentence_model = SentenceTransformer('all-MiniLM-L6-v2')
+
+# Cache PDF extraction
+@st.cache_data
+def extract_text_from_pdf(pdf_file):
+    pdf_reader = PdfReader(pdf_file)
+    text_data = []
+    for page_num, page in enumerate(pdf_reader.pages):
+        text = page.extract_text()
+        text = re.sub(r'\s+', ' ', text)  # Clean extra whitespace
+        text_data.append({
+            "page": page_num + 1,
+            "content": text
+        })
+    return text_data
+
+# Enhanced text chunking with overlap
+def split_text_into_chunks(text, chunk_size=500, overlap=100):
+    words = text.split()
+    chunks = []
+    for i in range(0, len(words), chunk_size - overlap):
+        chunks.append(" ".join(words[i:i + chunk_size]))
+    return chunks
+
+# Enhanced semantic search using sentence transformers
+def semantic_search(query, chunks, threshold=0.3):
+    query_embedding = sentence_model.encode([query])
+    chunk_embeddings = sentence_model.encode(chunks)
+    similarities = cosine_similarity(query_embedding, chunk_embeddings)[0]
+    results = [(chunks[i], similarities[i]) for i in np.argsort(similarities)[::-1]]
+    return [res for res in results if res[1] > threshold]
+
+# Improved query translation with error handling
 def query_huggingface_model(prompt, max_new_tokens=1000, temperature=0.7, top_k=50):
-    model_name = "HuggingFaceH4/zephyr-7b-alpha"  # Replace with your preferred model
+    model_name = "HuggingFaceH4/zephyr-7b-alpha"
     api_url = f"https://api-inference.huggingface.co/models/{model_name}"
     headers = {"Authorization": f"Bearer {huggingface_token}"}
     payload = {
@@ -24,206 +61,96 @@ def query_huggingface_model(prompt, max_new_tokens=1000, temperature=0.7, top_k=
             "top_k": top_k,
         },
     }
-    response = requests.post(api_url, headers=headers, json=payload)
-    if response.status_code == 200:
-        return response.json()[0]["generated_text"]
-    else:
-        st.error(f"Error: {response.status_code} - {response.text}")
-        return None
-
-# Function to detect language
-def detect_language(text):
     try:
-        return detect(text)
-    except:
-        return "en"  # Default to English if detection fails
-
-# Function to extract text from PDF with line and page numbers
-def extract_text_from_pdf(pdf_file):
-    pdf_reader = PdfReader(pdf_file)
-    text_data = []
-    for page_num, page in enumerate(pdf_reader.pages):
-        lines = page.extract_text().split('\n')
-        for line_num, line in enumerate(lines):
-            text_data.append({
-                "page": page_num + 1,
-                "line": line_num + 1,
-                "content": line
-            })
-    return text_data
-
-# Function to search for query in PDF content
-def search_pdf_content(pdf_text_data, query):
-    results = []
-    for entry in pdf_text_data:
-        if query.lower() in entry["content"].lower():
-            results.append(entry)
-    return results
-
-# Function to split text into chunks
-def split_text_into_chunks(text, chunk_size=500):
-    words = text.split()
-    chunks = [" ".join(words[i:i + chunk_size]) for i in range(0, len(words), chunk_size)]
-    return chunks
-
-# Function to compute cosine similarity between query and document chunks
-def compute_cosine_similarity(query, chunks):
-    vectorizer = TfidfVectorizer()
-    tfidf_matrix = vectorizer.fit_transform([query] + chunks)
-    cosine_similarities = cosine_similarity(tfidf_matrix[0:1], tfidf_matrix[1:]).flatten()
-    return cosine_similarities
-
-# Function to find KNN-based similar documents
-def find_knn_similar_documents(query, chunks, k=5):
-    vectorizer = TfidfVectorizer()
-    tfidf_matrix = vectorizer.fit_transform([query] + chunks)
-    knn = NearestNeighbors(n_neighbors=k, metric="cosine")
-    knn.fit(tfidf_matrix[1:])
-    distances, indices = knn.kneighbors(tfidf_matrix[0:1])
-    return indices.flatten(), distances.flatten()
+        response = requests.post(api_url, headers=headers, json=payload, timeout=30)
+        if response.status_code == 200:
+            return response.json()[0]["generated_text"]
+        else:
+            st.error(f"API Error: {response.status_code}")
+            return None
+    except Exception as e:
+        st.error(f"Connection Error: {str(e)}")
+        return None
 
-# Default system prompts for each query translation method
-DEFAULT_SYSTEM_PROMPTS = {
-    "Multi-Query": """You are an AI language model assistant. Your task is to generate five 
-different versions of the given user question to retrieve relevant documents from a vector 
-database. By generating multiple perspectives on the user question, your goal is to help
-the user overcome some of the limitations of the distance-based similarity search. 
-Provide these alternative questions separated by newlines. Original question: {question}""",
-    "RAG Fusion": """You are an AI language model assistant. Your task is to combine multiple 
-queries into a single, refined query to improve retrieval accuracy. Original question: {question}""",
-    "Decomposition": """You are an AI language model assistant. Your task is to break down 
-the given user question into simpler sub-questions. Provide these sub-questions separated 
-by newlines. Original question: {question}""",
-    "Step Back": """You are an AI language model assistant. Your task is to refine the given 
-user question by taking a step back and asking a more general question. Original question: {question}""",
-    "HyDE": """You are an AI language model assistant. Your task is to generate a hypothetical 
-document that would be relevant to the given user question. Original question: {question}""",
-}
+# Enhanced indexing strategies
+def create_index(text_chunks, method="Multi-Representation"):
+    if method == "Multi-Representation":
+        return TfidfVectorizer().fit_transform(text_chunks)
+    elif method == "Raptors":
+        embeddings = sentence_model.encode(text_chunks)
+        index = faiss.IndexFlatL2(embeddings.shape[1])
+        index.add(embeddings)
+        return index
+    elif method == "ColBERT":
+        return sentence_model.encode(text_chunks)
+
+# Improved similarity search with multiple methods
+def similarity_search(query, chunks, method="Cosine", index=None, k=5):
+    if method == "Cosine":
+        return semantic_search(query, chunks)
+    elif method == "KNN":
+        if isinstance(index, faiss.IndexFlatL2):
+            query_embedding = sentence_model.encode([query])
+            distances, indices = index.search(query_embedding, k)
+            return [(chunks[i], 1 - distances[0][j]) for j, i in enumerate(indices[0])]
+    return []
 
 # Streamlit App
 def main():
-    st.title("RAG Model with Advanced Query Translation and Indexing")
-    st.write("Enter a prompt and get a response from the model.")
-
-    # Sidebar for options
-    st.sidebar.title("Options")
-
-    # PDF Upload
-    st.sidebar.header("Upload PDF")
-    pdf_file = st.sidebar.file_uploader("Upload a PDF file", type="pdf")
-
-    # Query Translation Options
-    st.sidebar.header("Query Translation")
-    query_translation = st.sidebar.selectbox(
-        "Select Query Translation Method",
-        ["Multi-Query", "RAG Fusion", "Decomposition", "Step Back", "HyDE"]
-    )
-
-    # Indexing Options
-    st.sidebar.header("Indexing")
-    indexing_method = st.sidebar.selectbox(
-        "Select Indexing Method",
-        ["Multi-Representation", "Raptors", "ColBERT"]
-    )
-
-    # Similarity Search Options
-    st.sidebar.header("Similarity Search")
-    similarity_method = st.sidebar.selectbox(
-        "Select Similarity Search Method",
-        ["Cosine Similarity", "KNN"]
-    )
-    if similarity_method == "KNN":
-        k_value = st.sidebar.slider("Select K Value", 1, 10, 5)
-
-    # LLM Parameters
-    st.sidebar.header("LLM Parameters")
-    max_new_tokens = st.sidebar.slider("Max New Tokens", 10, 1000, 1000)
-    temperature = st.sidebar.slider("Temperature", 0.1, 1.0, 0.7)
-    top_k = st.sidebar.slider("Top K", 1, 100, 50)
-
-    # System Prompt
-    st.sidebar.header("System Prompt")
-    default_system_prompt = DEFAULT_SYSTEM_PROMPTS[query_translation]
-    system_prompt = st.sidebar.text_area("System Prompt", default_system_prompt)
-
-    # Main Content
-    st.header("Input Prompt")
-    prompt = st.text_input("Enter your prompt:")
+    st.title("Enhanced RAG Model with Advanced Features")
+    
+    # Sidebar configurations
+    st.sidebar.title("Configuration")
+    pdf_file = st.sidebar.file_uploader("Upload PDF", type="pdf")
+    query_translation = st.sidebar.selectbox("Query Translation", list(DEFAULT_SYSTEM_PROMPTS.keys()))
+    indexing_method = st.sidebar.selectbox("Indexing Method", ["Multi-Representation", "Raptors", "ColBERT"])
+    similarity_method = st.sidebar.selectbox("Similarity Search", ["Cosine", "KNN"])
+    similarity_threshold = st.sidebar.slider("Similarity Threshold", 0.0, 1.0, 0.3)
+    
+    # Main interface
+    prompt = st.text_input("Enter your query:")
+    
     if prompt:
-        st.write("**Prompt:**", prompt)
-
-        # Detect Language
-        language = detect_language(prompt)
-        st.write(f"**Detected Language:** {language}")
-
-        # Query Translation
-        if st.button("Apply Query Translation"): 
-            # Format the system prompt with the user's question
-            formatted_prompt = system_prompt.format(question=prompt)
-            st.write("**Formatted System Prompt:**", formatted_prompt)
-
-            # Query the Hugging Face model for query translation
-            translated_queries = query_huggingface_model(formatted_prompt, max_new_tokens, temperature, top_k)
-            if translated_queries:
-                st.write("**Translated Queries:**")
-                st.write(translated_queries.split("\n")[-1])  # Print only the updated question part
-
-        # Indexing
-        if st.button("Apply Indexing"):
-            st.write(f"**Applied Indexing Method:** {indexing_method}")
-            if pdf_file is not None:
-                # Extract and search PDF content
-                pdf_text_data = extract_text_from_pdf(pdf_file)
-                search_results = search_pdf_content(pdf_text_data, prompt)
-
-                if search_results:
-                    st.write("**Relevant Content from PDF:**")
-                    for result in search_results:
-                        st.write(f"**Page {result['page']}, Line {result['line']}:** {result['content']}")
-
-                    # Split text into chunks
-                    chunks = split_text_into_chunks("\n".join([result["content"] for result in search_results]))
-                    st.write("**Chunks Obtained from PDF:**")
-                    for i, chunk in enumerate(chunks):
-                        st.write(f"**Chunk {i + 1}:** {chunk}")
-
-                    # Perform similarity search
-                    if similarity_method == "Cosine Similarity":
-                        st.write("**Cosine Similarity Results:**")
-                        cosine_similarities = compute_cosine_similarity(prompt, chunks)
-                        for i, similarity in enumerate(cosine_similarities):
-                            st.write(f"**Chunk {i + 1} Similarity:** {similarity:.4f}")
-                    elif similarity_method == "KNN":
-                        st.write(f"**KNN Results (k={k_value}):**")
-                        indices, distances = find_knn_similar_documents(prompt, chunks, k_value)
-                        for i, (index, distance) in enumerate(zip(indices, distances)):
-                            st.write(f"**Chunk {index + 1} Distance:** {distance:.4f}")
+        with st.spinner("Processing..."):
+            # Query Translation
+            translated_prompt = query_huggingface_model(
+                DEFAULT_SYSTEM_PROMPTS[query_translation].format(question=prompt)
+            )
+            
+            if pdf_file:
+                # Process PDF
+                text_data = extract_text_from_pdf(pdf_file)
+                full_text = " ".join([p["content"] for p in text_data])
+                chunks = split_text_into_chunks(full_text)
+                
+                # Create index
+                index = create_index(chunks, indexing_method)
+                
+                # Perform search
+                if query_translation == "HyDE":
+                    hypothetical_answer = translated_prompt
+                    results = semantic_search(hypothetical_answer, chunks, similarity_threshold)
                 else:
-                    st.write("**No relevant content found in the PDF.**")
-            else:
-                st.write("**No PDF uploaded.**")
-
-        # Generate Response
-        if st.button("Generate Response"):
-            if pdf_file is not None:
-                # Extract and search PDF content
-                pdf_text_data = extract_text_from_pdf(pdf_file)
-                search_results = search_pdf_content(pdf_text_data, prompt)
-
-                if search_results:
-                    st.write("**Relevant Content from PDF:**")
-                    for result in search_results:
-                        st.write(f"**Page {result['page']}, Line {result['line']}:** \"{result['content']}\"")
-
-                    # Generate response based on PDF content
-                    pdf_context = "\n".join([result["content"] for result in search_results])
-                    response = query_huggingface_model(f"Based on the following context:\n{pdf_context}\n\nAnswer this question: {prompt}", max_new_tokens, temperature, top_k)
-                    if response:
-                        st.write("**Response:**", response)
+                    results = similarity_search(prompt, chunks, similarity_method, index)
+                
+                # Display results
+                if results:
+                    st.subheader("Top Results:")
+                    for i, (chunk, score) in enumerate(results[:3]):
+                        st.markdown(f"**Result {i+1}** (Score: {score:.2f}):")
+                        st.write(chunk)
+                    
+                    # Generate response
+                    context = "\n".join([chunk for chunk, _ in results[:3]])
+                    response = query_huggingface_model(
+                        f"Context: {context}\n\nQuestion: {prompt}\n\nAnswer:"
+                    )
+                    st.subheader("Generated Response:")
+                    st.write(response)
                 else:
-                    st.write("**No relevant content found in the PDF. No response generated.**")
+                    st.warning("No relevant documents found matching the query.")
             else:
-                st.write("**No PDF uploaded. No response generated.**")
+                st.error("Please upload a PDF document first.")
 
 if __name__ == "__main__":
     main()