rag_test_task / test_rag.py
alexandraroze's picture
added test_rag
5847e55
raw
history blame contribute delete
843 Bytes
from src.rag import FaissDB
import argparse
from dotenv import load_dotenv
import os
load_dotenv()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--path_to_index", type=str, required=True)
args = parser.parse_args()
path_to_index = args.path_to_index
try:
faiss_db = FaissDB(emb_model=os.getenv("OPENAI_EMBEDDINGS_MODEL"))
faiss_db.load_index(path_to_index)
except Exception as e:
print(f"Error loading index: {e}")
exit(1)
while True:
query = input("Enter query: ")
if query == "exit":
break
try:
documents = faiss_db.similarity_search(query)
print("\n\n".join(documents))
except Exception as e:
print(f"Error searching for documents: {e}")
continue