Update app.py

added rag part

app.py

@@ -1,152 +1,294 @@
-import time
-import gradio as gr
+# import time
+# import gradio as gr
+
+# import torch
+# from transformers import GPT2LMHeadModel, GPT2Tokenizer
+# from peft import PeftModel, LoraConfig
+
+
+# tokenizer = GPT2Tokenizer.from_pretrained("gpt2-medium")
+# tokenizer.pad_token = tokenizer.eos_token
+# base_model = GPT2LMHeadModel.from_pretrained("gpt2-medium")
+
+# lora_config = LoraConfig(
+#     r=8,
+#     lora_alpha=16,
+#     target_modules=["c_fc", "c_proj", "c_attn"],
+#     lora_dropout=0.1,
+#     task_type="CAUSAL_LM"
+# )
+    
+# finetuned_model = PeftModel.from_pretrained(base_model, "./lora_ft_weights", config=lora_config)
+# finetuned_model.eval()
+
+
+
+# # -------------------------------
+# # Simulated QA Models
+# # -------------------------------
+# def qa_system(method, question):
+#     start_time = time.time()
+
+#     if not question.strip():
+#         return "**Error:** Please enter a question.", 0.0, "0 seconds", ""
+
+#     # Simulated response based on method
+#     if method == "Retrieval-Augmented Generation (RAG)":
+#         answer = "Using RAG: Based on retrieved financial documents, the answer is $95,000,000."
+#         model_name = "RAG-based Model"
+#         confidence = 0.92
+#     else:
+#         model_name = "GPT 2-finetuned"
+#         # Input guradrails
+#         financial_keywords = [
+#             'revenue', 'profit', 'earnings', 'financial', 'income', 'balance', 
+#             'cash', 'debt', 'equity', 'assets', 'market', 'investment', 'sales',
+#             'cost', 'margin', 'growth', 'compliance', 'risk', 'customer'
+#         ]
+#         for text in question:            
+#         # Check for financial content
+#             text_lower = text.lower()
+#             if any(pattern in text_lower for pattern in financial_keywords):
+#                 return "This question does not seem to be related to Finance"
+
+#         prompt = f"You are a financial assistant.\nUse the context below to answer the question.\n\nQuestion: {question}\nAnswer:"
+        
+#         inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
+    
+#         with torch.no_grad():
+#             outputs = finetuned_model.generate(
+#                 **inputs,
+#                 max_length=inputs['input_ids'].shape[1] + 100,
+#                 temperature=0.7,
+#                 do_sample=True,
+#                 pad_token_id=tokenizer.eos_token_id,
+#             )
+        
+#         generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+#         answer = generated_text.split("Answer:")[-1].strip()
+        
+#     end_time = time.time()
+#     response_time = round(end_time - start_time, 2)
 
+#     return (
+#         f"**Method:** {model_name}",
+#         0.95,
+#         f"{response_time} seconds",
+#         answer
+#     )
+
+# # -------------------------------
+# # Gradio UI
+# # -------------------------------
+# with gr.Blocks(css="""
+#     .radio-vertical .wrap {
+#         flex-direction: column !important;
+#     }
+#     .radio-vertical .wrap > label {
+#         margin-bottom: 8px !important;
+#         margin-right: 0 !important;
+#     }
+#     .small-btn {
+#         max-width: fit-content !important;
+#         width: auto !important;
+#     }
+#     .small-btn button {
+#         width: auto !important; 
+#         min-width: unset !important; 
+#         padding: 8px 16px !important; 
+#         font-size: 16px !important;
+#         white-space: nowrap !important;
+#         max-width: fit-content !important;
+#     }
+# """) as demo:
+#     gr.Markdown(
+#         """
+#         # 📊 Comparative Financial QA System  
+#         An implementation comparing **Retrieval-Augmented Generation (RAG)** and a **Fine-Tuned on LoRA and Replay-Based Learning** GPT 2 model for answering questions on financial reports.
+#         """
+#     )
+
+#     # Radio buttons displayed vertically
+#     method = gr.Radio(
+#         choices=["Retrieval-Augmented Generation (RAG)", "Fine-Tuned Model"],
+#         label="Choose QA Method:",
+#         value="Fine-Tuned Model",
+#         interactive=True,
+#         elem_classes="radio-vertical"
+#     )
+
+#     # Question input
+#     question = gr.Textbox(
+#         label="Ask a question about Nice's 2023-2024 financials:",
+#         placeholder="e.g., What was the total revenue in 2023?"
+#     )
+
+#     # Get Answer button — auto-sized
+#     submit_btn = gr.Button("Get Answer", elem_classes="small-btn")
+
+#     # Output section - initially hidden
+#     with gr.Group(visible=False) as output_section:
+#         method_output = gr.Markdown()
+#         confidence_output = gr.Number(label="Model Confidence")
+#         response_time_output = gr.Textbox(label="Response Time")
+#         answer_output = gr.Markdown(label="Answer")
+
+#     # Button click handler
+#     def handle_submit(method_val, question_val):
+#         # Show output section and get results
+#         results = qa_system(method_val, question_val)
+#         return [gr.Group(visible=True)] + list(results)
+    
+#     submit_btn.click(
+#         handle_submit,
+#         inputs=[method, question],
+#         outputs=[output_section, method_output, confidence_output, response_time_output, answer_output]
+#     )
+
+# # -------------------------------
+# # Launch for Hugging Face Spaces
+# # -------------------------------
+# if __name__ == "__main__":
+#     demo.launch()
+
+
+
+import time, re, numpy as np, pickle, faiss
+import gradio as gr
 import torch
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
 from peft import PeftModel, LoraConfig
+from sentence_transformers import SentenceTransformer, CrossEncoder
+from rank_bm25 import BM25Okapi
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-
+# -------------------------------
+# Fine-Tuned GPT2 Setup
+# -------------------------------
 tokenizer = GPT2Tokenizer.from_pretrained("gpt2-medium")
 tokenizer.pad_token = tokenizer.eos_token
 base_model = GPT2LMHeadModel.from_pretrained("gpt2-medium")
 
 lora_config = LoraConfig(
-    r=8,
-    lora_alpha=16,
-    target_modules=["c_fc", "c_proj", "c_attn"],
-    lora_dropout=0.1,
-    task_type="CAUSAL_LM"
+    r=8, lora_alpha=16,
+    target_modules=["c_fc","c_proj","c_attn"],
+    lora_dropout=0.1, task_type="CAUSAL_LM"
 )
-    
 finetuned_model = PeftModel.from_pretrained(base_model, "./lora_ft_weights", config=lora_config)
 finetuned_model.eval()
 
+# -------------------------------
+# RAG Setup
+# -------------------------------
+# Load chunks + FAISS
+with open("financial_chunks.pkl","rb") as f:
+    chunks = pickle.load(f)
+dense_index = faiss.read_index("financial_index.faiss")
+
+# Sparse BM25
+tokenized_corpus = [c["text"].split(" ") for c in chunks]
+bm25_index = BM25Okapi(tokenized_corpus)
+
+# Embedding + reranker
+embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
+
+# Generator
+gen_name = "google/flan-t5-large"
+rag_tok = AutoTokenizer.from_pretrained(gen_name)
+rag_model = AutoModelForSeq2SeqLM.from_pretrained(gen_name)
+
+def detect_numeric_query(query: str) -> bool:
+    return any(k in query.lower() for k in ["revenue","income","profit","eps","earnings","assets","liabilities","cash","dividend","margin","cost","expenses","sales","percentage"])
+
+def hybrid_retrieve(query, top_k=20, alpha=0.6):
+    query_embedding = embedder.encode([query], convert_to_numpy=True)
+    D_dense, I_dense = dense_index.search(query_embedding.astype(np.float32), top_k)
+    dense_results = [{'chunk': chunks[i], 'score': float(D_dense[0][j])} for j, i in enumerate(I_dense[0])]
 
+    tokenized_query = query.lower().split(" ")
+    sparse_scores_list = bm25_index.get_scores(tokenized_query)
+    sparse_indices = np.argsort(sparse_scores_list)[-top_k:][::-1]
+    sparse_results = [{'chunk': chunks[i], 'score': float(sparse_scores_list[i])} for i in sparse_indices]
+
+    fused = {c['chunk']['id']: 0 for c in dense_results + sparse_results}
+    for res in dense_results:
+        fused[res['chunk']['id']] += alpha * (1 - res['score']/(1+res['score']))
+    for res in sparse_results:
+        fused[res['chunk']['id']] += (1-alpha) * res['score']
+
+    candidates = [c for c in chunks if c["id"] in fused]
+    pairs = [(query, c["text"]) for c in candidates]
+    scores = reranker.predict(pairs)
+    ranked = sorted(zip(candidates, scores), key=lambda x: x[1], reverse=True)
+    return [c for c, s in ranked[:6]]
+
+def rag_pipeline(query: str):
+    if not query.strip():
+        return "Empty query", 0.0, 0.0
+
+    t0 = time.time()
+    ctxs = hybrid_retrieve(query)
+
+    if not detect_numeric_query(query):
+        pairs = [(query, c["text"]) for c in ctxs]
+        scores = reranker.predict(pairs)
+        ctxs = [c for c, _ in sorted(zip(ctxs, scores), key=lambda x: x[1], reverse=True)[:1]]
+
+    context = " ".join([c["text"] for c in ctxs])
+
+    if detect_numeric_query(query):
+        prompt = f"Context:\n{context}\n\nQuestion: {query}\nAnswer with the exact numeric value."
+    else:
+        prompt = f"Context:\n{context}\n\nQuestion: {query}\nAnswer in 2-3 sentences."
+
+    inputs = rag_tok(prompt, return_tensors="pt", max_length=1024, truncation=True)
+    outputs = rag_model.generate(**inputs, max_new_tokens=120)
+    ans = rag_tok.decode(outputs[0], skip_special_tokens=True)
+
+    return ans.strip(), 0.9, round(time.time()-t0,2)
 
 # -------------------------------
-# Simulated QA Models
+# Unified QA System
 # -------------------------------
 def qa_system(method, question):
-    start_time = time.time()
-
-    if not question.strip():
-        return "**Error:** Please enter a question.", 0.0, "0 seconds", ""
-
-    # Simulated response based on method
     if method == "Retrieval-Augmented Generation (RAG)":
-        answer = "Using RAG: Based on retrieved financial documents, the answer is $95,000,000."
-        model_name = "RAG-based Model"
-        confidence = 0.92
+        answer, confidence, elapsed = rag_pipeline(question)
+        return f"**Method:** RAG-based Model", confidence, f"{elapsed} seconds", answer
     else:
-        model_name = "GPT 2-finetuned"
-        # Input guradrails
-        financial_keywords = [
-            'revenue', 'profit', 'earnings', 'financial', 'income', 'balance', 
-            'cash', 'debt', 'equity', 'assets', 'market', 'investment', 'sales',
-            'cost', 'margin', 'growth', 'compliance', 'risk', 'customer'
-        ]
-        for text in question:            
-        # Check for financial content
-            text_lower = text.lower()
-            if any(pattern in text_lower for pattern in financial_keywords):
-                return "This question does not seem to be related to Finance"
-
-        prompt = f"You are a financial assistant.\nUse the context below to answer the question.\n\nQuestion: {question}\nAnswer:"
-        
+        # Fine-tuned GPT2
+        prompt = f"You are a financial assistant.\nQuestion: {question}\nAnswer:"
         inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
-    
         with torch.no_grad():
-            outputs = finetuned_model.generate(
-                **inputs,
-                max_length=inputs['input_ids'].shape[1] + 100,
-                temperature=0.7,
-                do_sample=True,
-                pad_token_id=tokenizer.eos_token_id,
-            )
-        
-        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        answer = generated_text.split("Answer:")[-1].strip()
-        
-    end_time = time.time()
-    response_time = round(end_time - start_time, 2)
-
-    return (
-        f"**Method:** {model_name}",
-        0.95,
-        f"{response_time} seconds",
-        answer
-    )
+            outputs = finetuned_model.generate(**inputs, max_length=inputs['input_ids'].shape[1]+100,
+                                               temperature=0.7, do_sample=True,
+                                               pad_token_id=tokenizer.eos_token_id)
+        answer = tokenizer.decode(outputs[0], skip_special_tokens=True).split("Answer:")[-1].strip()
+        return f"**Method:** Fine-Tuned Model", 0.95, "N/A", answer
 
 # -------------------------------
-# Gradio UI
+# Gradio UI (same as before)
 # -------------------------------
-with gr.Blocks(css="""
-    .radio-vertical .wrap {
-        flex-direction: column !important;
-    }
-    .radio-vertical .wrap > label {
-        margin-bottom: 8px !important;
-        margin-right: 0 !important;
-    }
-    .small-btn {
-        max-width: fit-content !important;
-        width: auto !important;
-    }
-    .small-btn button {
-        width: auto !important; 
-        min-width: unset !important; 
-        padding: 8px 16px !important; 
-        font-size: 16px !important;
-        white-space: nowrap !important;
-        max-width: fit-content !important;
-    }
-""") as demo:
-    gr.Markdown(
-        """
-        # 📊 Comparative Financial QA System  
-        An implementation comparing **Retrieval-Augmented Generation (RAG)** and a **Fine-Tuned on LoRA and Replay-Based Learning** GPT 2 model for answering questions on financial reports.
-        """
-    )
-
-    # Radio buttons displayed vertically
-    method = gr.Radio(
-        choices=["Retrieval-Augmented Generation (RAG)", "Fine-Tuned Model"],
-        label="Choose QA Method:",
-        value="Fine-Tuned Model",
-        interactive=True,
-        elem_classes="radio-vertical"
-    )
-
-    # Question input
-    question = gr.Textbox(
-        label="Ask a question about Nice's 2023-2024 financials:",
-        placeholder="e.g., What was the total revenue in 2023?"
-    )
-
-    # Get Answer button — auto-sized
-    submit_btn = gr.Button("Get Answer", elem_classes="small-btn")
-
-    # Output section - initially hidden
+with gr.Blocks() as demo:
+    gr.Markdown("# 📊 Comparative Financial QA System")
+
+    method = gr.Radio(["Retrieval-Augmented Generation (RAG)", "Fine-Tuned Model"],
+                      label="Choose QA Method:", value="Fine-Tuned Model")
+    question = gr.Textbox(label="Ask a question:")
+    submit_btn = gr.Button("Get Answer")
+
     with gr.Group(visible=False) as output_section:
         method_output = gr.Markdown()
         confidence_output = gr.Number(label="Model Confidence")
         response_time_output = gr.Textbox(label="Response Time")
         answer_output = gr.Markdown(label="Answer")
 
-    # Button click handler
     def handle_submit(method_val, question_val):
-        # Show output section and get results
         results = qa_system(method_val, question_val)
         return [gr.Group(visible=True)] + list(results)
-    
-    submit_btn.click(
-        handle_submit,
-        inputs=[method, question],
-        outputs=[output_section, method_output, confidence_output, response_time_output, answer_output]
-    )
 
-# -------------------------------
-# Launch for Hugging Face Spaces
-# -------------------------------
+    submit_btn.click(handle_submit, [method, question],
+                     [output_section, method_output, confidence_output, response_time_output, answer_output])
+
 if __name__ == "__main__":
     demo.launch()