Spaces:
Build error
Build error
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() |