Update app.py

app.py

@@ -3,59 +3,94 @@ import io
 import re
 import numpy as np
 import faiss
-import time
-import gradio as gr
+import torch
 from pypdf import PdfReader
 from rank_bm25 import BM25Okapi
 from sentence_transformers import SentenceTransformer
 from accelerate import Accelerator
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from bert_score import score
+import gradio as gr
 
-def chunk_text(text, chunk_size=700, overlap_size=150):
-    """Chunks text without breaking words in the middle (corrected overlap)."""
-    chunks = []
-    start = 0
-    text_length = len(text)
-
-    while start < text_length:
-        end = min(start + chunk_size, text_length)
-
-        # Ensure we do not split words
-        if end < text_length and text[end].isalnum():
-            last_space = text.rfind(" ", start, end)  # Find last space within the chunk
-            if last_space != -1:  # If a space is found, adjust the end
-                end = last_space
-
-        chunk = text[start:end].strip()
-        if chunk:  # Avoid empty chunks
-            chunks.append(chunk)
-
-        if end == text_length:
-            break
-
-        # Corrected overlap calculation
-        overlap_start = max(0, end - overlap_size)
-        if overlap_start < end: # Prevent infinite loop if overlap_start is equal to end.
-            last_overlap_space = text.rfind(" ", 0, overlap_start)
-            if last_overlap_space != -1 and last_overlap_space > start:
-                start = last_overlap_space + 1
+# --- Preload Data ---
+DEFAULT_PDF_URLS = [
+    "https://www.latentview.com/wp-content/uploads/2023/07/LatentView-Annual-Report-2022-23.pdf",
+    "https://www.latentview.com/wp-content/uploads/2024/08/LatentView-Annual-Report-2023-24.pdf"
+] 
+
+def preload_data(pdf_urls):
+    embedding_model = SentenceTransformer("BAAI/bge-large-en")
+
+    def download_pdf(url):
+        response = requests.get(url, stream=True)
+        response.raise_for_status()
+        return response.content
+
+    def extract_text_from_pdf(pdf_bytes):
+        pdf_file = io.BytesIO(pdf_bytes)
+        reader = PdfReader(pdf_file)
+        text = ""
+        for page in reader.pages:
+            text += page.extract_text() or ""
+        return text
+
+    def preprocess_text(text):
+        financial_symbols = r"\$\€\₹\£\¥\₩\₽\₮\₦\₲"
+        text = re.sub(fr"[^\w\s{financial_symbols}.,%/₹$€¥£-]", "", text)
+        text = re.sub(r'\s+', ' ', text).strip()
+        return text
+
+    def chunk_text(text, chunk_size=700, overlap_size=150):
+        chunks = []
+        start = 0
+        text_length = len(text)
+        while start < text_length:
+            end = min(start + chunk_size, text_length)
+            if end < text_length and text[end].isalnum():
+                last_space = text.rfind(" ", start, end)
+                if last_space != -1:
+                    end = last_space
+            chunk = text[start:end].strip()
+            if chunk:
+                chunks.append(chunk)
+            if end == text_length:
+                break
+            overlap_start = max(0, end - overlap_size)
+            if overlap_start < end:
+                last_overlap_space = text.rfind(" ", 0, overlap_start)
+                if last_overlap_space != -1 and last_overlap_space > start:
+                    start = last_overlap_space + 1
+                else:
+                    start = end
             else:
-                start = end # If no space found, start at the last end.
-        else:
-            start = end
+                start = end
+        return chunks
 
-    return chunks
+    all_data = []
+    for url in pdf_urls:
+        pdf_bytes = download_pdf(url)
+        text = extract_text_from_pdf(pdf_bytes)
+        preprocessed_text = preprocess_text(text)
+        all_data.append(preprocessed_text)
 
-chunks = []
-for data in all_data:
-  chunks.extend(chunk_text(data))
+    chunks = []
+    for data in all_data:
+        chunks.extend(chunk_text(data))
 
-embedding_model = SentenceTransformer("BAAI/bge-large-en")
-# embedding_model = SentenceTransformer('multi-qa-mpnet-base-dot-v1')
-embeddings = embedding_model.encode(chunks)
+    embeddings = embedding_model.encode(chunks)
+    index = faiss.IndexFlatL2(embeddings.shape[1])
+    index.add(embeddings)
+
+    return index, chunks
 
-index = faiss.IndexFlatL2(embeddings.shape[1])
-index.add(embeddings)
+index, chunks = preload_data(DEFAULT_PDF_URLS)
+embedding_model = SentenceTransformer("BAAI/bge-large-en")
+accelerator = Accelerator()
+MODEL_NAME = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto", trust_remote_code=True, cache_dir="./my_models")
+model = accelerator.prepare(model)
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 def bm25_retrieval(query, documents, top_k=3):
     tokenized_docs = [doc.split() for doc in documents]
@@ -75,47 +110,26 @@ def rerank(query, results):
     similarities = np.dot(result_embeddings, query_embedding.T).flatten()
     return [results[i] for i in np.argsort(similarities)[::-1]], similarities
 
-#Chunk merging.
 def merge_chunks(retrieved_chunks, overlap_size=100):
-    """Merges overlapping chunks properly by detecting the actual overlap."""
     merged_chunks = []
     buffer = retrieved_chunks[0] if retrieved_chunks else ""
-
     for i in range(1, len(retrieved_chunks)):
         chunk = retrieved_chunks[i]
-
-        # Find actual overlap
-        overlap_start = buffer[-overlap_size:]  # Get the last `overlap_size` chars of the previous chunk
-        overlap_index = chunk.find(overlap_start)  # Find where this part appears in the new chunk
-
+        overlap_start = buffer[-overlap_size:]
+        overlap_index = chunk.find(overlap_start)
         if overlap_index != -1:
-            # Merge only the non-overlapping part
             buffer += chunk[overlap_index + overlap_size:]
         else:
-            # Store completed merged chunk and start a new one
             merged_chunks.append(buffer)
             buffer = chunk
-
     if buffer:
         merged_chunks.append(buffer)
-
     return merged_chunks
 
-# def calculate_confidence(query, context, similarities):
-#     return np.mean(similarities)  # Averaged similarity scores
 def calculate_confidence(query, answer):
     P, R, F1 = score([answer], [query], lang="en", verbose=False)
     return F1.item()
 
-# Load SLM
-accelerator = Accelerator()
-accelerator.free_memory()
-MODEL_NAME = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto", cache_dir="./my_models")
-model = accelerator.prepare(model)
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
-
 def generate_response(query, context):
     prompt = f"""Your task is to analyze the given Context and answer the Question concisely in plain English. 
     **Guidelines:**
@@ -130,12 +144,7 @@ def generate_response(query, context):
     answer = response.split("Answer:")[1].strip()
     return answer
 
-def process_query(pdf_urls_text, query):
-    pdf_urls = [url.strip() for url in pdf_urls_text.split("\n") if url.strip()]
-    if not pdf_urls:
-        return "Please enter at least one PDF URL."
-
-    index, chunks = load_and_index_data(pdf_urls)
+def process_query(query):
     retrieved_chunks = adaptive_retrieval(query, index, chunks)
     merged_chunks = merge_chunks(retrieved_chunks, 150)
     reranked_chunks, similarities = rerank(query, merged_chunks)
@@ -147,11 +156,11 @@ def process_query(pdf_urls_text, query):
 
 iface = gr.Interface(
     fn=process_query,
-    inputs=[gr.Textbox(lines=3, placeholder="Enter PDF URLs (one per line)"), gr.Textbox(placeholder="Enter your financial question")],
+    inputs=gr.Textbox(placeholder="Enter your financial question"),
     outputs="text",
     title="Financial Document Q&A Chatbot",
-    description="Enter PDF URLs and your question to get answers from the documents."
+    description="Ask questions about the preloaded financial documents."
 )
 
 iface.launch()
-accelerator.free_memory()
+accelerator.free_memory()