RAGtest / app.py
willco-afk's picture
Create app.py
2b509dc verified
import gradio as gr
import chromadb
import numpy as np
from sentence_transformers import SentenceTransformer
from transformers import pipeline
import pickle
# Load pre-trained model and embeddings
model = SentenceTransformer("all-MiniLM-L6-v2") # You can upload this model from HF Hub if available
generator = pipeline("text-generation", model="gpt2")
# Initialize ChromaDB client (using the Chroma database uploaded as a file)
client = chromadb.Client()
collection = client.create_collection("documents")
# Manually load your embeddings and document data from the HF Space files
with open("embeddings.pkl", "rb") as f:
embeddings = pickle.load(f)
# Example of adding embeddings to FAISS (if using FAISS as the indexer)
faiss_index = faiss.IndexFlatL2(512) # Adjust dimension if needed
faiss_index.add(np.array(embeddings))
# Example documents loaded manually or fetched via API
documents = ["What is RAG?", "How does FAISS work?", "Introduction to Chroma."]
def generate_answer(query):
query_embedding = model.encode([query])
D, I = faiss_index.search(np.array(query_embedding), k=1) # Retrieve the closest document
retrieved_doc = documents[I[0][0]]
prompt = f"Context: {retrieved_doc}\nQuestion: {query}\nAnswer:"
response = generator(prompt, max_length=50)
return response[0]['generated_text']
# Gradio interface for manual file uploads and query input
iface = gr.Interface(fn=generate_answer, inputs="text", outputs="text")
iface.launch()