app.py

# 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"
)
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)

# -------------------------------
# Unified QA System
# -------------------------------
def qa_system(method, question):
    start_time = time.time()
    if method == "Retrieval-Augmented Generation (RAG)":
        answer, confidence, elapsed = rag_pipeline(question)
        return f"**Method:** RAG-based Model", confidence, f"{elapsed} seconds", answer
    else:
        # 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)
        answer = tokenizer.decode(outputs[0], skip_special_tokens=True).split("Answer:")[-1].strip()
        if "Question:" in answer:
            answer = answer.split("Question:")[0].strip()
        
        end_time = time.time()
        response_time = round(end_time - start_time, 2)
        return f"**Method:** Fine-Tuned Model", 0.95, response_time, answer.strip()

# -------------------------------
# Gradio UI (same as before)
# -------------------------------
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")

    def handle_submit(method_val, question_val):
        results = qa_system(method_val, question_val)
        return [gr.Group(visible=True)] + list(results)

    submit_btn.click(handle_submit, [method, question],
                     [output_section, method_output, confidence_output, response_time_output, answer_output])

if __name__ == "__main__":
    demo.launch()