Update app.py

app.py

@@ -22,16 +22,37 @@ def generate_answer_rag(prompt, max_tokens=200):
 # -----------------------------
 # Combined Answer Function
 # -----------------------------
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
 def combined_generate(prompt, max_tokens):
     with ThreadPoolExecutor() as executor:
-        # Submit both tasks in parallel
-        ft_future = executor.submit(fine_tuning.generate_answer, model, tokenizer, device, prompt, max_tokens)
-        rag_future = executor.submit(generate_answer_rag, prompt, max_tokens)
+        start_times = {"Fine-tuned": time.time(), "RAG": time.time()}
+
+        futures = {
+            executor.submit(generate_answer, model, tokenizer, device, prompt, max_tokens): "Fine-tuned",
+            executor.submit(generate_answer_rag, prompt, max_tokens): "RAG",
+        }
+
+        answers = {"Fine-tuned": "", "RAG": ""}
+        times = {"Fine-tuned": None, "RAG": None}
+
+        for future in as_completed(futures):
+            model_name = futures[future]
+            answers[model_name] = future.result()
+            times[model_name] = round(time.time() - start_times[model_name], 2)
+
+            # Format answers with time taken
+            ft_display = answers["Fine-tuned"]
+            if times["Fine-tuned"] is not None:
+                ft_display += f"\n\n⏱ Took {times['Fine-tuned']}s"
+
+            rag_display = answers["RAG"]
+            if times["RAG"] is not None:
+                rag_display += f"\n\n⏱ Took {times['RAG']}s"
 
-        fine_tuned_answer = ft_future.result()
-        rag_answer = rag_future.result()
+            yield ft_display, rag_display
 
-    return fine_tuned_answer, rag_answer
 
 
 # -----------------------------