Upload query_smoke_test.py

query_smoke_test.py

@@ -0,0 +1,81 @@
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+from datasets import load_dataset
+
+
+def main():
+    device = ("cuda" if torch.cuda.is_available() else "cpu")
+    tokenizer = AutoTokenizer.from_pretrained('vblagoje/retribert-base-uncased')
+    model = AutoModel.from_pretrained('vblagoje/retribert-base-uncased').to(device)
+    _ = model.eval()
+
+    index_file_name = "./data/kilt_wikipedia.faiss"
+    kilt_wikipedia = load_dataset("kilt_wikipedia", split="full")
+    columns = ['kilt_id', 'wikipedia_id', 'wikipedia_title', 'text', 'anchors', 'categories',
+                              'wikidata_info', 'history']
+
+    min_snippet_length = 20
+    topk = 21
+
+    def articles_to_paragraphs(examples):
+        ids, titles, sections, texts, start_ps, end_ps, start_cs, end_cs = [], [], [], [], [], [], [], []
+        for bidx, example in enumerate(examples["text"]):
+            last_section = ""
+            for idx, p in enumerate(example["paragraph"]):
+                if "Section::::" in p:
+                    last_section = p
+                ids.append(examples["wikipedia_id"][bidx])
+                titles.append(examples["wikipedia_title"][bidx])
+                sections.append(last_section)
+                texts.append(p)
+                start_ps.append(idx)
+                end_ps.append(idx)
+                start_cs.append(0)
+                end_cs.append(len(p))
+
+        return {"wikipedia_id": ids, "title": titles,
+                "section": sections, "text": texts,
+                "start_paragraph_id": start_ps, "end_paragraph_id": end_ps,
+                "start_character": start_cs,
+                "end_character": end_cs
+                }
+
+    kilt_wikipedia_paragraphs = kilt_wikipedia.map(articles_to_paragraphs, batched=True,
+                                                   remove_columns=columns,
+                                                   batch_size=256, cache_file_name=f"./wiki_kilt_paragraphs_full.arrow",
+                                                   desc="Expanding wiki articles into paragraphs")
+
+    # use paragraphs that are not simple fragments or very short sentences
+    kilt_wikipedia_paragraphs = kilt_wikipedia_paragraphs.filter(lambda x: x["end_character"] > 250)
+    kilt_wikipedia_paragraphs.load_faiss_index("embeddings", index_file_name, device=0)
+
+    def embed_questions_for_retrieval(questions):
+        query = tokenizer(questions, max_length=128, padding=True, truncation=True, return_tensors="pt")
+        with torch.no_grad():
+            q_reps = model.embed_questions(query["input_ids"].to(device),
+                                           query["attention_mask"].to(device)).cpu().type(torch.float)
+        return q_reps.numpy()
+
+    def query_index(question):
+        question_embedding = embed_questions_for_retrieval([question])
+        scores, wiki_passages = kilt_wikipedia_paragraphs.get_nearest_examples("embeddings", question_embedding, k=topk)
+        columns = ['wikipedia_id', 'title', 'text', 'section', 'start_paragraph_id', 'end_paragraph_id', 'start_character','end_character']
+        retrieved_examples = []
+        r = list(zip(wiki_passages[k] for k in columns))
+        for i in range(topk):
+            retrieved_examples.append({k: v for k, v in zip(columns, [r[j][0][i] for j in range(len(columns))])})
+        return retrieved_examples
+
+    questions = ["What causes the contrails (cirrus aviaticus) behind jets at high altitude? ",
+                 "Why does water heated to a room temeperature feel colder than the air around it?"]
+    res_list = query_index(questions[0])
+    res_list = [res for res in res_list if len(res["text"].split()) > min_snippet_length][:int(topk / 3)]
+    for res in res_list:
+        print("\n")
+        print(res)
+
+
+main()
+
+