Replace DatasetDict by Dataset

functions.py

@@ -5,7 +5,7 @@ import torch
 from bs4 import BeautifulSoup
 from peft import PeftConfig, PeftModel
 from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig, AutoModel
-from datasets import DatasetDict
+from datasets import DatasetDict, Dataset
 
 # os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 
@@ -57,7 +57,8 @@ def build_faiss_index(text):
     emb_list = []
     for item in text_list:
         emb_list.append({"embeddings": get_embeddings(item)})
-    dataset = DatasetDict({'train': emb_list})
+    # dataset = DatasetDict({'train': emb_list})
+    dataset = Dataset.from_dict(emb_list)
     dataset.add_faiss_index(column="embeddings")
     shared['embeddings_dataset'] = dataset
     print(['build_faiss_index', 'end'])