# rag_module.py

import json
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer, CrossEncoder
from sklearn.feature_extraction.text import TfidfVectorizer
from transformers import AutoTokenizer, T5ForConditionalGeneration
import gradio as gr

from rag_pipeline import RAGPipeline


# Gradio function
# ==== Instantiate RAG ====
rag = RAGPipeline(
    embedder_model="infly/inf-retriever-v1-1.5b",
    reranker_model="cross-encoder/ms-marco-MiniLM-L-6-v2",
    generator_model="google/flan-t5-base"
)

# ==== Gradio App ====
def answer_question(query, top_k):
    return rag.generate_answer(query, top_k)

gr.Interface(
    fn=answer_question,
    inputs=[
        gr.Textbox(label="Enter your question"),
        gr.Slider(1, 5, step=1, value=3, label="Top-K Retrieved Chunks")
    ],
    outputs=gr.Textbox(label="📘 Generated Answer"),
    title="RAG-based Question Answering",
    description="A lightweight RAG system using FAISS + TF-IDF + FLAN-T5. Ask a question, get an answer."
).launch()