Create app.py

app.py

@@ -0,0 +1,239 @@
+import streamlit as st
+import nltk
+import fitz
+import spacy
+import json
+import subprocess
+import re
+import numpy as np
+from summa import keywords
+from nltk.tokenize import sent_tokenize, word_tokenize
+from sentence_transformers import SentenceTransformer, util
+import time
+
+# Download required NLTK data
+nltk.download('punkt_tab', quiet=True)
+nltk.download('stopwords', quiet=True)
+
+# Load models
+nlp = spacy.load("en_core_web_sm")
+sbert_model = SentenceTransformer("paraphrase-MiniLM-L6-v2")
+
+CHARS_TO_REMOVE = "(){},;-'\":‘’“”"
+
+# Text processing functions (unchanged from your code)
+def clean_text(text):
+    text = "".join(char if char not in CHARS_TO_REMOVE else " " for char in text)
+    text = re.sub(r'\s+', ' ', text).strip()
+    return text
+
+def extract_text_from_pdf(pdf_file):
+    try:
+        doc = fitz.open(stream=pdf_file.read(), filetype="pdf")
+        text = "\n".join(page.get_text("text") for page in doc)
+        return text.strip()
+    except Exception as e:
+        st.error(f"Error reading PDF: {e}")
+        return ""
+
+def custom_sent_tokenize(text):
+    sentence_endings = re.compile(r'(?<!\b[A-Z])(?<!\b[A-Z]\.)(?<!\b[A-Z]\.[A-Z])(?<=\.)\s+')
+    sentences = sentence_endings.split(text)
+    return [s.strip() for s in sentences if s.strip()]
+
+def chunk_text(text, chunk_size=15, max_words=150, overlap=10):
+    sentences = custom_sent_tokenize(text)
+    chunks = []
+    i = 0
+    while i < len(sentences):
+        chunk = sentences[i:i + chunk_size]
+        chunk_text = " ".join(chunk)
+        words = word_tokenize(chunk_text)
+        if len(words) > max_words:
+            chunk.pop()
+            chunk_text = " ".join(chunk)
+        if chunks:
+            prev_words = word_tokenize(chunks[-1])[-overlap:]
+            chunk_text = " ".join(prev_words) + " " + chunk_text
+        chunks.append(chunk_text)
+        i += chunk_size
+    return chunks
+
+# Keyword extraction functions (unchanged)
+def normalize_text(text):
+    return text.lower().strip()
+
+def lemmatize_keywords(keywords_list):
+    doc = nlp(" ".join(keywords_list))
+    return {token.lemma_ for token in doc if token.is_alpha}
+
+def extract_keywords(chunk):
+    doc = nlp(chunk)
+    ner_keywords = {normalize_text(ent.text) for ent in doc.ents}
+    singlerank_keywords = {normalize_text(kw) for kw in keywords.keywords(chunk, scores=False).split("\n")}
+    all_tokens = {normalize_text(token.text) for token in doc if token.is_alpha}
+    all_keywords = ner_keywords | singlerank_keywords | all_tokens
+    return lemmatize_keywords(all_keywords)
+
+# Embedding generation (unchanged)
+def get_chunk_embeddings(chunks):
+    return [sbert_model.encode(chunk, convert_to_tensor=True) for chunk in chunks]
+
+# Levenshtein distance and keyword correction (unchanged)
+def levenshtein_distance(s1, s2):
+    if len(s1) < len(s2):
+        return levenshtein_distance(s2, s1)
+    if len(s2) == 0:
+        return len(s1)
+    previous_row = range(len(s2) + 1)
+    for i, c1 in enumerate(s1):
+        current_row = [i + 1]
+        for j, c2 in enumerate(s2):
+            insertions = previous_row[j + 1] + 1
+            deletions = current_row[j] + 1
+            substitutions = previous_row[j] + (c1 != c2)
+            current_row.append(min(insertions, deletions, substitutions))
+        previous_row = current_row
+    return previous_row[-1]
+
+def correct_keywords(query_keywords, stored_keywords, threshold=2):
+    corrected_keywords = set()
+    for qk in query_keywords:
+        if qk in stored_keywords:
+            corrected_keywords.add(qk)
+        else:
+            min_dist = float('inf')
+            best_match = qk
+            for sk in stored_keywords:
+                dist = levenshtein_distance(qk, sk)
+                if dist < min_dist:
+                    min_dist = dist
+                    best_match = sk
+            if min_dist <= threshold:
+                corrected_keywords.add(best_match)
+            else:
+                corrected_keywords.add(qk)
+    return corrected_keywords
+
+# Bit Vector-based search and retrieval (adapted for multiple PDFs)
+def process_pdf(pdf_file):
+    text = extract_text_from_pdf(pdf_file)
+    text = clean_text(text)
+    chunks = chunk_text(text)
+    n_chunks = len(chunks)
+
+    keyword_bitmaps = {}
+    chunk_keywords = []
+    for i, chunk in enumerate(chunks):
+        keywords = extract_keywords(chunk)
+        chunk_keywords.append(keywords)
+        for kw in keywords:
+            if kw not in keyword_bitmaps:
+                keyword_bitmaps[kw] = np.zeros(n_chunks, dtype=bool)
+            keyword_bitmaps[kw][i] = 1
+
+    chunk_embeddings = get_chunk_embeddings(chunks)
+    all_keywords = set().union(*chunk_keywords)
+    return chunks, chunk_embeddings, keyword_bitmaps, chunk_keywords, all_keywords
+
+def search_relevant_chunks(query, chunks, chunk_embeddings, keyword_bitmaps, chunk_keywords, all_keywords, top_k=5):
+    query_keywords = extract_keywords(query)
+    corrected_query_keywords = correct_keywords(query_keywords, all_keywords)
+
+    matched_bitmap = np.zeros(len(chunks), dtype=bool)
+    for keyword in corrected_query_keywords:
+        if keyword in keyword_bitmaps:
+            matched_bitmap |= keyword_bitmaps[keyword]
+
+    matched_chunk_indices = set(np.where(matched_bitmap)[0])
+    chunk_scores = {idx: len(corrected_query_keywords & chunk_keywords[idx]) for idx in matched_chunk_indices}
+    matched_chunks = [(chunks[idx], idx) for idx in sorted(chunk_scores, key=chunk_scores.get, reverse=True)]
+
+    if len(matched_chunks) >= top_k:
+        return [chunk for chunk, _ in matched_chunks[:top_k]]
+
+    remaining_slots = top_k - len(matched_chunks)
+    unmatched_indices = [i for i in range(len(chunks)) if i not in matched_chunk_indices]
+    query_embedding = sbert_model.encode(query, convert_to_tensor=True)
+    similarities = [util.pytorch_cos_sim(query_embedding, chunk_emb).item() for chunk_emb in chunk_embeddings]
+    top_indices = sorted(unmatched_indices, key=lambda i: similarities[i], reverse=True)[:remaining_slots]
+    similar_chunks = [chunks[i] for i in top_indices]
+    return [chunk for chunk, _ in matched_chunks] + similar_chunks[:top_k]
+
+# Mistral API query (unchanged)
+def query_mistral(prompt, MISTRAL_API_KEY):
+    payload = {"model": "mistral-large-latest", "messages": [{"role": "user", "content": prompt}]}
+    curl_command = [
+        "curl", "--location", "https://api.mistral.ai/v1/chat/completions",
+        "--header", "Content-Type: application/json",
+        "--header", "Accept: application/json",
+        "--header", f"Authorization: Bearer {MISTRAL_API_KEY}",
+        "--data", json.dumps(payload)
+    ]
+    response = subprocess.run(curl_command, capture_output=True, text=True)
+    if response.returncode == 0:
+        try:
+            response_json = json.loads(response.stdout)
+            return response_json['choices'][0]['message']['content']
+        except (KeyError, json.JSONDecodeError):
+            return "Error parsing the LLM response."
+    return f"Error: {response.stderr}"
+
+# Streamlit app
+st.title("PDF Query System")
+st.write("Upload PDFs and ask questions about their content.")
+
+# File uploader for multiple PDFs
+uploaded_files = st.file_uploader("Upload PDF files", type="pdf", accept_multiple_files=True)
+
+# Store processed PDFs in session state
+if 'processed_pdfs' not in st.session_state:
+    st.session_state.processed_pdfs = {}
+
+# Process uploaded PDFs
+if uploaded_files:
+    for pdf_file in uploaded_files:
+        if pdf_file.name not in st.session_state.processed_pdfs:
+            with st.spinner(f"Processing {pdf_file.name}..."):
+                start_time = time.time()
+                chunks, chunk_embeddings, keyword_bitmaps, chunk_keywords, all_keywords = process_pdf(pdf_file)
+                st.session_state.processed_pdfs[pdf_file.name] = {
+                    "chunks": chunks,
+                    "chunk_embeddings": chunk_embeddings,
+                    "keyword_bitmaps": keyword_bitmaps,
+                    "chunk_keywords": chunk_keywords,
+                    "all_keywords": all_keywords
+                }
+                end_time = time.time()
+                st.success(f"Processed {pdf_file.name} in {end_time - start_time:.4f} seconds")
+
+# Query input
+query = st.text_input("Enter your query:")
+
+# Mistral API key (you may want to secure this differently in production)
+MISTRAL_API_KEY = "S3vzsvK7rP5in24joHgL55dVCjqYSi1F"
+
+if st.button("Search") and query and st.session_state.processed_pdfs:
+    with st.spinner("Searching..."):
+        all_relevant_chunks = []
+        for pdf_name, data in st.session_state.processed_pdfs.items():
+            start_search = time.time()
+            relevant_chunks = search_relevant_chunks(
+                query, data["chunks"], data["chunk_embeddings"], 
+                data["keyword_bitmaps"], data["chunk_keywords"], data["all_keywords"]
+            )
+            end_search = time.time()
+            all_relevant_chunks.extend(relevant_chunks)
+            st.write(f"Search time for {pdf_name}: {end_search - start_search:.4f} seconds")
+
+        context = "\n".join(all_relevant_chunks)
+        start_response_time = time.time()
+        llm_prompt = f"Only Based on the following context, answer the query:\n{context}\n\nQuery: {query}"
+        response = query_mistral(llm_prompt, MISTRAL_API_KEY)
+        end_response_time = time.time()
+
+        st.subheader("Response:")
+        st.write(response)
+        st.write(f"Response time: {end_response_time - start_response_time:.4f} seconds")
+elif st.button("Search") and not st.session_state.processed_pdfs:
+    st.warning("Please upload at least one PDF before searching.")