Update app.py

app.py

@@ -25,7 +25,7 @@ model = AutoModelForCausalLM.from_pretrained(
     model_name,
     quantization_config=quantization_config,
     device_map="auto",
-    max_memory={0: "22GiB", "cpu": "6GiB"}  # Prevent VRAM overflow 
+    max_memory={0: "22GiB", "cpu": "6GiB"}  # Prevent VRAM overflow
 )
 text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
@@ -62,15 +62,15 @@ def start_new_session():
 def get_embedding(text):
     return embedding_model.encode(text, normalize_embeddings=True)
 
-def store_chat_in_session(user_input, response):
+def store_chat_in_session(user_input, response, reference):
     if current_session_id is None:
         start_new_session()
-    chat_sessions[current_session_id].append((user_input, response))
+    chat_sessions[current_session_id].append((user_input, response, reference))
     chat_index.add(np.array([get_embedding(response)]))
 
 def get_recent_chat_history():
     if current_session_id in chat_sessions:
-        return "\n".join([f"User: {q}\nAI: {r}" for q, r in chat_sessions[current_session_id]])
+        return "\n".join([f"User: {q}\nAI: {r}\nReference: {ref}" for q, r, ref in chat_sessions[current_session_id]])
     return ""
 
 # Document Processing
@@ -93,8 +93,8 @@ def retrieve_relevant_passage(query, top_k=3):
     D, I = doc_index.search(np.array([query_embedding]), top_k)
     valid_indices = [i for i in I[0] if 0 <= i < len(doc_texts)]
     if valid_indices:
-        return "\n".join([f"- {doc_texts[i]}" for i in valid_indices])
-    return "No relevant document found."
+        return "\n".join([f"- {doc_texts[i]}" for i in valid_indices]), "\n".join([doc_texts[i] for i in valid_indices])
+    return "No relevant document found.", ""
 
 # Retrieve Chat Context
 def retrieve_chat_context(user_input, top_k=3):
@@ -111,8 +111,8 @@ def retrieve_chat_context(user_input, top_k=3):
 def chat_with_pdf(user_input, chat_history=[]):
     if not authenticated:
         return "Access Denied!", chat_history
-    relevant_passage = retrieve_relevant_passage(user_input)
-    past_chat_context = retrieve_chat_context(user_input)
+    relevant_passage, reference = retrieve_relevant_passage(user_input)
+    past_chat_context = get_recent_chat_history()
     prompt = (
         "You are an HR assistant. Provide responses based on company policies. If unsure, say 'Please contact HR'.\n\n"
         f"Recent Chat:\n{past_chat_context}\nHR Policy Context:\n{relevant_passage}\nUser Inquiry: {user_input}\nAI Response:"
@@ -123,8 +123,10 @@ def chat_with_pdf(user_input, chat_history=[]):
             prompt, max_new_tokens=1024, do_sample=True, temperature=0.3, top_p=0.85, repetition_penalty=1.2,
             return_full_text=False
         )
-        for token in response[0]['generated_text'].split():
-            yield token + " "
+        answer = response[0]['generated_text'].split("AI Response:")[-1].strip()
+        store_chat_in_session(user_input, answer, reference)
+        formatted_response = f"{answer}\n\n*Reference:* _{reference}_"
+        yield formatted_response
     
     return response_generator(), chat_history