RAG-Chatbot / app.py
adarshzolekar's picture
Create app.py
43b9dec verified
raw
history blame contribute delete
865 Bytes
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()