Update main.py

main.py

@@ -2,6 +2,7 @@ import json  # to work with JSON
 import threading  # to allow streaming response
 import time  # to pave the deliver of the message
 
+import datasets  # for loading RAG database
 import faiss  # to create a search index
 import gradio  # for the interface
 import numpy  # to work with vectors
@@ -23,7 +24,7 @@ EMBEDDING_MODEL_NAME = "allenai-specter"
 LLM_MODEL_NAME = "Qwen/Qwen2-7B-Instruct"
 
 # Load the dataset and convert to pandas
-data = pandas.read_parquet("hf://datasets/ccm/rag-idetc/data/train-00000-of-00001.parquet")
+data = datasets.load_dataset("ccm/rag-idetc")["train"].to_pandas()
 
 # Load the model for later use in embeddings
 model = sentence_transformers.SentenceTransformer(EMBEDDING_MODEL_NAME)