Update src/models/predict_model.py

src/models/predict_model.py

@@ -5,15 +5,16 @@ from src.models.bm25_utils import BM25Gensim
 from src.models.qa_model import *
 from tqdm.auto import tqdm
 tqdm.pandas()
+from datasets import load_dataset
 
-df_wiki_windows = pd.read_csv("/home/user/app/src/data/processed/wikipedia_20220620_cleaned_v2.csv")
-df_wiki = pd.read_csv("/home/user/app/src/data/wikipedia_20220620_short.csv")
+df_wiki_windows = load_dataset("foxxy-hm/e2eqa-wiki",  data_files="processed/wikipedia_20220620_cleaned_v2.csv")["train"].to_pandas()
+df_wiki = load_dataset("foxxy-hm/e2eqa-wiki",  data_files="processed/wikipedia_20220620_short.csv")["train"].to_pandas()
 df_wiki.title = df_wiki.title.apply(str)
 
-entity_dict = json.load(open("/home/user/app/src/data/processed/entities.json"))
+entity_dict = load_dataset("foxxy-hm/e2eqa-wiki",  data_files="processed/entities.json")["train"].to_dict()
 new_dict = dict()
 for key, val in entity_dict.items():
-    val = val.replace("wiki/", "").replace("_", " ")
+    val = val[0].replace("wiki/", "").replace("_", " ")
     entity_dict[key] = val
     key = preprocess(key)
     new_dict[key.lower()] = val
@@ -22,15 +23,15 @@ title2idx = dict([(x.strip(), y) for x, y in zip(df_wiki.title, df_wiki.index.va
 
 qa_model = QAEnsembleModel("nguyenvulebinh/vi-mrc-large", ["src/models/qa_model_robust.bin"], entity_dict)
 pairwise_model_stage1 = PairwiseModel("nguyenvulebinh/vi-mrc-base").half()
-pairwise_model_stage1.load_state_dict(torch.load("/home/user/app/src/models/pairwise_v2.bin"))
+pairwise_model_stage1.load_state_dict(torch.load("/src/models/pairwise_v2.bin"))
 pairwise_model_stage1.eval()
 
 pairwise_model_stage2 = PairwiseModel("nguyenvulebinh/vi-mrc-base").half()
-pairwise_model_stage2.load_state_dict(torch.load("/home/user/app/src/models/pairwise_stage2_seed0.bin"))
+pairwise_model_stage2.load_state_dict(torch.load("/src/models/pairwise_stage2_seed0.bin"))
 
-bm25_model_stage1 = BM25Gensim("/home/user/app/src/models/bm25_stage1/", entity_dict, title2idx)
-bm25_model_stage2_full = BM25Gensim("/home/user/app/src/models/bm25_stage2/full_text/", entity_dict, title2idx)
-bm25_model_stage2_title = BM25Gensim("/home/user/app/src/models/bm25_stage2/title/", entity_dict, title2idx)
+bm25_model_stage1 = BM25Gensim("/src/models/bm25_stage1/", entity_dict, title2idx)
+bm25_model_stage2_full = BM25Gensim("/src/models/bm25_stage2/full_text/", entity_dict, title2idx)
+bm25_model_stage2_title = BM25Gensim("/src/models/bm25_stage2/title/", entity_dict, title2idx)
 
 def get_answer_e2e(question):
     #Bm25 retrieval for top200 candidates