implementation of a RAG system

from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

# Load document embeddings and corpus
embedder = SentenceTransformer('all-MiniLM-L6-v2')
document_embeddings = np.load('document_embeddings.npy')
with open('document_corpus.txt', 'r') as f:
    documents = f.readlines()

def retrieve_relevant_docs(query, top_k=3):
    query_embedding = embedder.encode([query])
    similarities = cosine_similarity(query_embedding, document_embeddings)[0]
    top_indices = similarities.argsort()[-top_k:][::-1]
    return [documents[i] for i in top_indices]

def rag_response(query):
    relevant_docs = retrieve_relevant_docs(query)
    context = "\n".join(relevant_docs)

    full_prompt = f"""
    Based on the following information, please answer the user's question:

    Context:
    {context}

    User Question: {query}

    Answer:
    """

    response = llm.generate(full_prompt)
    return response

# Example usage
user_query = "What are the key provisions of the latest data privacy regulation?"
answer = rag_response(user_query)
print(answer)