adding citation feature

app.py

@@ -18,26 +18,34 @@ for path in file_names:
 model = SentenceTransformer('all-MiniLM-L6-v2')
 
 def preprocessText(text):
-    cleanedText = text.strip()
-    chunks = cleanedText.split("\n")
-    # removes section headings from the chunks
-    cleanedChunks = [chunk.strip() for chunk in chunks if chunk.strip() and not chunk.strip().startswith("===")]
-    return cleanedChunks
+    cleanedText = text.strip().split("\n")
+    cleanedChunks = []
+    sectionLabels = []
+    for chunk in cleanedText: 
+        chunk = chunk.strip()
+        if not chunk: 
+            continue
+        if chunk.startsith("===") and "SECTION:" in chunk:
+            currentLabel = chunk.replace("===", "",).replace("SECTION:", "").strip()
+            continue
+        cleanedChunks.append(chunk)
+        sectionLabels.append(currentLabel)
+    return cleanedChunks, sectionLabels
 
 def createEmbeddings(textChunks):
     chunkEmbeddings = model.encode(cleanedChunks, convert_to_tensor = True)
-    return chunkEmbeddings
 
-def getTopChunks(query, chunkEmbeddings, textChunks):
+def getTopChunks(query, chunkEmbeddings, textChunks, sectionLabels):
     queryEmbedding = model.encode(query, convert_to_tensor = True)
     queryEmbeddingNormalized = queryEmbedding / queryEmbedding.norm()
     chunkEmbeddingsNormalized = chunkEmbeddings / chunkEmbeddings.norm(dim = 1, keepdim = True)
     similarities = torch.matmul(chunkEmbeddingsNormalized, queryEmbeddingNormalized)
     topIndices = torch.topk(similarities, k=3).indices
     topChunks = [textChunks[i] for i in topIndices]
-    return topChunks
+    topSections = [sectionlabels[i] for i in topIndices]
+    return topChunks, topSections
 
-cleanedChunks = preprocessText(financialText)
+cleanedChunks, sectionLabels = preprocessText(financialText)
 chunkEmbeddings = createEmbeddings(cleanedChunks)
 
 client = Groq(api_key = os.environ.get("SF_TOKEN"))
@@ -54,8 +62,7 @@ def respond(message, history):
         messages.extend([{"role": h["role"], "content": h["content"]} for h in history])
         # helping grok to retain message history
 
-    topResults = getTopChunks(message, chunkEmbeddings, cleanedChunks)
-    print(f'Pulled chunks: {topResults}')
+    topResults, topSections = getTopChunks(message, chunkEmbeddings, cleanedChunks, sectionLabels)
     context = "\n\n".join(topResults)
 
     messages.append({"role": "system",
@@ -76,6 +83,10 @@ def respond(message, history):
             response += token
             yield response
 
+    usedSections = list(dict.fromskeys(topSections))
+    citation = "\n\n* Sources: " + ", ".join(usedSections) + "*"
+    yield response + citation
+
 chatbot = gr.ChatInterface(respond, title = "Student Formula Bot 🔬",
                           description = 'Welcome to the core component of \"The Student Formula\": the RAG chatbot! With the ability to act as a finance tutor, accountability buddy, and goal-setting partner all in one, it\'s designed to best suit your needs on the way to productivity and success. To get started, ask about the basic principles of creating a budget!')
 
@@ -102,9 +113,7 @@ with gr.Blocks(css=custom_css) as demo:
             "What is educational investment?"
         ]
     )
-
-
-
+    
     gr.HTML(
         """
     <iframe