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()