Update app.py

app.py

@@ -38,6 +38,7 @@ cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
 
 dataset = load_dataset("sritang/hack_df", split='train')
 mypassages = list(dataset.to_pandas()['paragraph'])
+mysource = list(dataset.to_pandas()['Source'])
 
 dataset_embed = load_dataset("sritang/hack_policy_embed_sri", split='train')
 dataset_embed_pd = dataset_embed.to_pandas()
@@ -101,7 +102,7 @@ def get_text_fmt(qry, passages = mypassages, doc_embedding=mycorpus_embeddings):
   prediction_text = []
   for hit in predictions:
     page_content = passages[hit['corpus_id']]
-    metadata = {"source": hit['corpus_id']}
+    metadata = {"source": mysource[hit['corpus_id']]}
     result = Document(page_content=page_content, metadata=metadata)
     prediction_text.append(result)
   return prediction_text