citation improved

app.py

@@ -1,7 +1,7 @@
 # app.py
 
 import os
-# import re
+import re
 import torch
 # import pdfplumber
 from chromadb.utils import embedding_functions
@@ -18,6 +18,50 @@ groq_api_key = os.environ.get('GROQ_API_KEY')
 chat_client = Groq(api_key=groq_api_key)
 model = "llama-3.2-90b-text-preview"
 
+def edit_text(text):
+    # Find all citations and their positions
+    citation_matches = list(re.finditer(r'\[(\d+)\]', text))
+
+    # List to store indices of citations to remove
+    indices_to_remove = []
+
+    prev_num = None
+    prev_index = None
+
+    # Identify consecutive duplicate citations
+    for i in range(len(citation_matches)):
+        current_citation = citation_matches[i]
+        current_num = current_citation.group(1)
+
+        if prev_num == current_num:
+            # Mark the previous citation for removal
+            indices_to_remove.append(prev_index)
+        prev_num = current_num
+        prev_index = i
+
+    # Reconstruct the text with modifications
+    output_parts = []
+    last_end = 0
+
+    for i in range(len(citation_matches)):
+        m = citation_matches[i]
+        start, end = m.span()
+        if i in indices_to_remove:
+            # Remove citation
+            output_parts.append(text[last_end:start])
+        else:
+            # Keep and modify citation
+            output_parts.append(text[last_end:start])
+            page_num = m.group(1)
+            new_citation = '[Page ' + page_num + ']'
+            output_parts.append(new_citation)
+        last_end = end
+
+    # Append any remaining text after the last citation
+    output_parts.append(text[last_end:])
+
+    modified_text = ''.join(output_parts)
+    return modified_text
 
 # def parse_pdf(pdf_path):
 
@@ -153,9 +197,11 @@ def process_query(query):
         return "No relevant information found in the manual."
 
     # Extract chunks and metadata
-    chunks = results['documents'][0]
+    pre_chunks = results['documents'][0]
     metadatas = results['metadatas'][0]
 
+    chunks = [f'Page {y["page_number"]}:: {x}' for x,y in zip(pre_chunks,metadatas)]
+
     reranked_chunks = colbert_rerank(query, chunks)
     final_context = " ".join(reranked_chunks[:10])
 
@@ -168,14 +214,16 @@ def process_query(query):
     if last_complete != -1:
         answer = answer[:last_complete + 1].strip()
 
+    answer = edit_text(answer)
     # Prepare citations
-    citations = [
-        f"Page {meta.get('page_number', 'N/A')}" for meta in metadatas[:5]
-    ]
+    # citations = [
+    #     f"Page {meta.get('page_number', 'N/A')}" for meta in metadatas[:5]
+    # ]
 
-    citations_text = "Pages cited from:\n" + "\n".join(citations)
+    # citations_text = "Pages cited from:\n" + "\n".join(citations)
 
-    return f"{answer}\n\n{citations_text}"
+    # return f"{answer}\n\n{citations_text}"
+    return answer
 
 # Initialize global variables
 def initialize():