import gradio as gr
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
from transformers import pipeline

docs = [
    "RAG stands for Retrieval Augmented Generation.",
    "This chatbot runs on Hugging Face Spaces.",
    "FAISS is a vector database."
]

embedder = SentenceTransformer("all-MiniLM-L6-v2")
doc_embeddings = embedder.encode(docs)

index = faiss.IndexFlatL2(doc_embeddings.shape[1])
index.add(np.array(doc_embeddings))

llm = pipeline("text2text-generation", model="google/flan-t5-base")

def chat(q):
    q_emb = embedder.encode([q])
    _, I = index.search(np.array(q_emb), k=2)
    context = " ".join([docs[i] for i in I[0]])
    prompt = f"Context: {context}\nQuestion: {q}\nAnswer:"
    return llm(prompt, max_length=120)[0]["generated_text"]

gr.Interface(chat, "text", "text", title="Mobile RAG Bot").launch()