Update app.py

app.py

@@ -13,24 +13,24 @@ import gradio as gr
 
 from rag_pipeline import RAGPipeline
 from adversarial_framework import *
+# Load all models and retrievers ONCE
+rag = RAGRetriever(
+    embedding_model="paraphrase-MiniLM-L3-v2",
+    cross_encoder_model="cross-encoder/ms-marco-MiniLM-L-6-v2",
+    generator_model="google/flan-t5-small"
+)
 
-adv_pipeline = AdversarialAttackPipeline(answer_generator=RAGPipeline().generate_answer)
+adv_pipeline = AdversarialAttackPipeline(answer_generator=rag.generate_answer)
 
+# Define the Gradio wrapper
 def gradio_wrapper(query, method, k):
     stats_text, auc, fig, pert_q, pert_r, adv_r = adv_pipeline.evaluate_adversarial_robustness(
         query=query,
         method=method,
         k=k
     )
+    return stats_text, f"{auc}", fig, pert_q, pert_r, adv_r
 
-    return (
-        stats_text,
-        f"{auc}",
-        fig,
-        f"🟠 Perturbed Query:\n\n{pert_q}",
-        f"🟢 Perturbed Response:\n\n{pert_r}",
-        f"🔴 Directly Perturbed Response of Normal Output:\n\n{adv_r}"
-    )
 
 gr.Interface(
     fn=gradio_wrapper,