File size: 865 Bytes
43b9dec
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
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()