update app.py

app.py

@@ -7,6 +7,7 @@ import tqdm
 import torch
 import numpy as np
 from time import sleep
+from datetime import datetime
 import threading
 import gc
 import os
@@ -34,7 +35,6 @@ from averitec.data.sample_claims import CLAIMS_Type
 from utils import create_user_id
 user_id = create_user_id()
 
-from datetime import datetime
 from azure.storage.fileshare import ShareServiceClient
 try:
     from dotenv import load_dotenv
@@ -86,8 +86,21 @@ LABEL = [
     "Not Enough Evidence",
     "Conflicting Evidence/Cherrypicking",
 ]
-# Veracity
+
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
+# Question Generation
+qg_tokenizer = BloomTokenizerFast.from_pretrained("bigscience/bloom-1b1")
+qg_model = BloomForCausalLM.from_pretrained("bigscience/bloom-1b1", torch_dtype=torch.bfloat16).to(device)
+# qg_tokenizer = BloomTokenizerFast.from_pretrained("bigscience/bloom-7b1")
+# qg_model = BloomForCausalLM.from_pretrained("bigscience/bloom-7b1", torch_dtype=torch.bfloat16).to(device)
+
+# rerank
+rerank_tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+rereank_bert_model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2, problem_type="single_label_classification")  # Must specify single_label for some reason
+best_checkpoint = "averitec/pretrained_models/bert_dual_encoder.ckpt"
+rerank_trained_model = DualEncoderModule.load_from_checkpoint(best_checkpoint, tokenizer=rerank_tokenizer, model=rereank_bert_model).to(device)
+
+# Veracity
 veracity_tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
 bert_model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=4, problem_type="single_label_classification")
 veracity_model = SequenceClassificationModule.load_from_checkpoint("averitec/pretrained_models/bert_veracity.ckpt",
@@ -99,7 +112,6 @@ best_checkpoint = 'averitec/pretrained_models/bart_justifications_verdict-epoch=
 justification_model = JustificationGenerationModule.load_from_checkpoint(best_checkpoint, tokenizer=justification_tokenizer, model=bart_model).to(device)
 # ---------------------------------------------------------------------------
 
-
 # Set up Gradio Theme
 theme = gr.themes.Base(
     primary_hue="blue",
@@ -182,7 +194,7 @@ class SequenceClassificationDataLoader(pl.LightningDataModule):
         )
 
 
-def averitec_veracity_prediction(claim, qa_evidence):
+def google_veracity_prediction(claim, qa_evidence):
     bert_model_name = "bert-base-uncased"
     tokenizer = BertTokenizer.from_pretrained(bert_model_name)
     bert_model = BertForSequenceClassification.from_pretrained(bert_model_name, num_labels=4,
@@ -340,7 +352,7 @@ def extract_claim_str(claim, qa_evidence, verdict_label):
     return claim_str
 
 
-def averitec_justification_generation(claim, qa_evidence, verdict_label):
+def google_justification_generation(claim, qa_evidence, verdict_label):
     #
     claim_str = extract_claim_str(claim, qa_evidence, verdict_label)
     claim_str.strip()
@@ -383,21 +395,19 @@ def QAprediction(claim, evidence, sources):
     for i, evi in enumerate(evidence, 1):
         part = f"""<span>Doc {i}</span>"""
         subpart = f"""<a href="#doc{i}" class="a-doc-ref" target="_self"><span class='doc-ref'><sup>{i}</sup></span></a>"""
-        # subpart = f"""<span class='doc-ref'>{i}</sup></span>"""
         subparts = "".join([part, subpart])
         parts.append(subparts)
 
     evidence_part = ", ".join(parts)
 
     prediction_title = f"""<h5>Prediction:</h5>"""
-    # if 'Google' in sources or 'AVeriTeC' in sources:
-    #     verdict_label = averitec_veracity_prediction(claim, evidence)
-    #     justification_label = averitec_justification_generation(claim, evidence, verdict_label)
-    #     # justification_label = "See retrieved docs."
+    # if 'Google' in sources:
+    #     verdict_label = google_veracity_prediction(claim, evidence)
+    #     justification_label = google_justification_generation(claim, evidence, verdict_label)
     #     justification_part = f"""<span>Justification: {justification_label}</span>"""
     # if 'WikiPedia' in sources:
-    #     # verdict_label = fever_veracity_prediction(claim, evidence)
-    #     justification_label = averitec_justification_generation(claim, evidence, verdict_label)
+    #     verdict_label = wikipedia_veracity_prediction(claim, evidence)
+    #     justification_label = wikipedia_justification_generation(claim, evidence, verdict_label)
     #     # justification_label = "See retrieved docs."
     #     justification_part = f"""<span>Justification: {justification_label}</span>"""
 
@@ -406,11 +416,8 @@ def QAprediction(claim, evidence, sources):
     # justification_label = "See retrieved docs."
     justification_part = f"""<span>Justification: {justification_label}</span>"""
 
-
     verdict_part = f"""Verdict: <span>{verdict_label}.</span><br>"""
-
     content_parts = "".join([evidence_title, evidence_part, prediction_title, verdict_part, justification_part])
-    # content_parts = "".join([evidence_title, evidence_part, verdict_title, verdict_part, justification_title, justification_part])
 
     return content_parts, [verdict_label, justification_label]
 
@@ -418,8 +425,8 @@ def QAprediction(claim, evidence, sources):
 # ----------GoogleAPIretriever---------
 def generate_reference_corpus(reference_file):
     with open(reference_file) as f:
-        j = json.load(f)
-        train_examples = j
+        #
+        train_examples = json.load(f)
 
     all_data_corpus = []
     tokenized_corpus = []
@@ -456,16 +463,16 @@ def docs2prompt(top_docs):
 
 def prompt_question_generation(test_claim, speaker="they", topk=10):
     #
-    reference_file = "averitec_code/data/train.json"
+    reference_file = "averitec/data/train.json"
     tokenized_corpus, all_data_corpus = generate_reference_corpus(reference_file)
     bm25 = BM25Okapi(tokenized_corpus)
 
     # Define the bloom model:
     accelerator = Accelerator()
-    accel_device = accelerator.device
-    device = "cuda:0" if torch.cuda.is_available() else "cpu"
-    tokenizer = BloomTokenizerFast.from_pretrained("bigscience/bloom-7b1")
-    model = BloomForCausalLM.from_pretrained("bigscience/bloom-7b1", torch_dtype=torch.bfloat16).to(device)
+    # accel_device = accelerator.device
+    # device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    # tokenizer = BloomTokenizerFast.from_pretrained("bigscience/bloom-7b1")
+    # model = BloomForCausalLM.from_pretrained("bigscience/bloom-7b1", torch_dtype=torch.bfloat16).to(device)
 
     # --------------------------------------------------
     # test claim
@@ -478,11 +485,11 @@ def prompt_question_generation(test_claim, speaker="they", topk=10):
              "\". Criticism includes questions like: "
     sentences = [prompt]
 
-    inputs = tokenizer(sentences, padding=True, return_tensors="pt").to(device)
-    outputs = model.generate(inputs["input_ids"], max_length=2000, num_beams=2, no_repeat_ngram_size=2,
+    inputs = qg_tokenizer(sentences, padding=True, return_tensors="pt").to(device)
+    outputs = qg_model.generate(inputs["input_ids"], max_length=2000, num_beams=2, no_repeat_ngram_size=2,
                              early_stopping=True)
 
-    tgt_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
+    tgt_text = qg_tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
     in_len = len(sentences[0])
     questions_str = tgt_text[in_len:].split("\n")[0]
 
@@ -592,7 +599,7 @@ def get_google_search_results(api_key, search_engine_id, google_search, sort_dat
     return search_results
 
 
-def averitec_search(claim, generate_question, speaker="they", check_date="2024-01-01", n_pages=1):  # n_pages=3
+def averitec_search(claim, generate_question, speaker="they", check_date="2024-07-01", n_pages=1):  # n_pages=3
     # default config
     api_key = os.environ["GOOGLE_API_KEY"]
     search_engine_id = os.environ["GOOGLE_SEARCH_ENGINE_ID"]
@@ -612,13 +619,14 @@ def averitec_search(claim, generate_question, speaker="they", check_date="2024-0
     ]
 
     # save to folder
-    store_folder = "averitec_code/store/retrieved_docs"
+    store_folder = "averitec/data/store/retrieved_docs"
     #
     index = 0
     questions = [q["question"] for q in generate_question]
 
     # check the date of the claim
-    sort_date = check_claim_date(check_date)  # check_date="2022-01-01"
+    current_date = datetime.now().strftime("%Y-%m-%d")
+    sort_date = check_claim_date(current_date)  # check_date="2022-01-01"
 
     #
     search_strings = []
@@ -643,6 +651,7 @@ def averitec_search(claim, generate_question, speaker="they", check_date="2024-0
         for page_num in range(n_pages):
             search_results = get_google_search_results(api_key, search_engine_id, google_search, sort_date,
                                                        this_search_string, page=page_num)
+            search_results = search_results[:5]
 
             for result in search_results:
                 link = str(result["link"])
@@ -733,21 +742,21 @@ def generate_step2_reference_corpus(reference_file):
 
 def decorate_with_questions(claim, retrieve_evidence, top_k=10):  # top_k=100
     #
-    reference_file = "averitec_code/data/train.json"
+    reference_file = "averitec/data/train.json"
     tokenized_corpus, prompt_corpus = generate_step2_reference_corpus(reference_file)
     prompt_bm25 = BM25Okapi(tokenized_corpus)
 
     # Define the bloom model:
     accelerator = Accelerator()
     accel_device = accelerator.device
-    device = "cuda:0" if torch.cuda.is_available() else "cpu"
-    tokenizer = BloomTokenizerFast.from_pretrained("bigscience/bloom-7b1")
-    model = BloomForCausalLM.from_pretrained(
-        "bigscience/bloom-7b1",
-        device_map="auto",
-        torch_dtype=torch.bfloat16,
-        offload_folder="./offload"
-    )
+    # device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    # tokenizer = BloomTokenizerFast.from_pretrained("bigscience/bloom-7b1")
+    # model = BloomForCausalLM.from_pretrained(
+    #     "bigscience/bloom-7b1",
+    #     device_map="auto",
+    #     torch_dtype=torch.bfloat16,
+    #     offload_folder="./offload"
+    # )
 
     #
     tokenized_corpus = []
@@ -795,11 +804,11 @@ def decorate_with_questions(claim, retrieve_evidence, top_k=10):  # top_k=100
         prompt = "\n\n".join(prompt_docs + [claim_prompt])
         sentences = [prompt]
 
-        inputs = tokenizer(sentences, padding=True, return_tensors="pt").to(device)
-        outputs = model.generate(inputs["input_ids"], max_length=5000, num_beams=2, no_repeat_ngram_size=2,
+        inputs = qg_tokenizer(sentences, padding=True, return_tensors="pt").to(device)
+        outputs = qg_model.generate(inputs["input_ids"], max_length=5000, num_beams=2, no_repeat_ngram_size=2,
                                  early_stopping=True)
 
-        tgt_text = tokenizer.batch_decode(outputs[:, inputs["input_ids"].shape[-1]:], skip_special_tokens=True)[0]
+        tgt_text = qg_tokenizer.batch_decode(outputs[:, inputs["input_ids"].shape[-1]:], skip_special_tokens=True)[0]
         # We are not allowed to generate more than 250 characters:
         tgt_text = tgt_text[:250]
 
@@ -815,13 +824,13 @@ def triple_to_string(x):
 
 def rerank_questions(claim, bm25_qas, topk=3):
     #
-    tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
-    bert_model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2,
-                                                               problem_type="single_label_classification")  # Must specify single_label for some reason
-    best_checkpoint = "averitec_code/pretrained_models/bert_dual_encoder.ckpt"
-    device = "cuda:0" if torch.cuda.is_available() else "cpu"
-    trained_model = DualEncoderModule.load_from_checkpoint(best_checkpoint, tokenizer=tokenizer, model=bert_model).to(
-        device)
+    # tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+    # bert_model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2,
+    #                                                            problem_type="single_label_classification")  # Must specify single_label for some reason
+    # best_checkpoint = "averitec/pretrained_models/bert_dual_encoder.ckpt"
+    # device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    # trained_model = DualEncoderModule.load_from_checkpoint(best_checkpoint, tokenizer=tokenizer, model=bert_model).to(
+    #     device)
 
     #
     strs_to_score = []
@@ -834,13 +843,13 @@ def rerank_questions(claim, bm25_qas, topk=3):
         values.append([question, answer, source])
 
     if len(bm25_qas) > 0:
-        encoded_dict = tokenizer(strs_to_score, max_length=512, padding="longest", truncation=True,
+        encoded_dict = rerank_tokenizer(strs_to_score, max_length=512, padding="longest", truncation=True,
                                  return_tensors="pt").to(device)
 
         input_ids = encoded_dict['input_ids']
         attention_masks = encoded_dict['attention_mask']
 
-        scores = torch.softmax(trained_model(input_ids, attention_mask=attention_masks).logits, axis=-1)[:, 1]
+        scores = torch.softmax(rerank_trained_model(input_ids, attention_mask=attention_masks).logits, axis=-1)[:, 1]
 
         top_n = torch.argsort(scores, descending=True)[:topk]
         pass_through = [{"question": values[i][0], "answers": values[i][1], "source_url": values[i][2]} for i in top_n]
@@ -852,20 +861,16 @@ def rerank_questions(claim, bm25_qas, topk=3):
     return top3_qa_pairs
 
 
-def GoogleAPIretriever(query):
+def Googleretriever(query, sources):
     # ----- Generate QA pairs using AVeriTeC
-    top3_qa_pairs_path = "averitec_code/top3_qa_pairs1.json"
-    if not os.path.exists(top3_qa_pairs_path):
-        # step 1: generate questions for the query/claim using Bloom
-        generate_question = prompt_question_generation(query)
-        # step 2: retrieve evidence for the generated questions using Google API
-        retrieve_evidence = averitec_search(query, generate_question)
-        # step 3: generate QA pairs for each retrieved document
-        bm25_qa_pairs = decorate_with_questions(query, retrieve_evidence)
-        # step 4: rerank QA pairs
-        top3_qa_pairs = rerank_questions(query, bm25_qa_pairs)
-    else:
-        top3_qa_pairs = json.load(open(top3_qa_pairs_path, 'r'))
+    # step 1: generate questions for the query/claim using Bloom
+    generate_question = prompt_question_generation(query)
+    # step 2: retrieve evidence for the generated questions using Google API
+    retrieve_evidence = averitec_search(query, generate_question)
+    # step 3: generate QA pairs for each retrieved document
+    bm25_qa_pairs = decorate_with_questions(query, retrieve_evidence)
+    # step 4: rerank QA pairs
+    top3_qa_pairs = rerank_questions(query, bm25_qa_pairs)
 
     # Add score to metadata
     results = []
@@ -877,12 +882,14 @@ def GoogleAPIretriever(query):
         metadata['cached_source_url'] = qa['source_url']
         metadata['short_name'] = "Evidence {}".format(i + 1)
         metadata['page_number'] = ""
+        metadata['title'] = qa['question']
+        metadata['evidence'] = qa['answers']
         metadata['query'] = qa['question']
         metadata['answer'] = qa['answers']
         metadata['page_content'] = "<b>Question</b>: " + qa['question'] + "<br>" + "<b>Answer</b>: " + qa['answers']
         page_content = f"""{metadata['page_content']}"""
-        results.append((metadata, page_content))
 
+        results.append(Docs(metadata, page_content))
     return results
 
 
@@ -1181,11 +1188,8 @@ def log_on_azure(file, logs, azure_share_client):
 
 def chat(claim, history, sources):
     evidence = []
-    # if 'Google' in sources:
-    #     evidence = GoogleAPIretriever(query)
-
-    # if 'WikiPediaDumps' in sources:
-    #     evidence = WikipediaDumpsretriever(query)
+    if 'Google' in sources:
+        evidence = Googleretriever(claim, sources)
 
     if 'WikiPedia' in sources:
         evidence = Wikipediaretriever(claim, sources)
@@ -1212,7 +1216,8 @@ def chat(claim, history, sources):
     for evi in evidence:
         title_str = evi.metadata['title']
         evi_str = evi.metadata['evidence']
-        evi_list.append([title_str, evi_str])
+        url_str = evi.metadata['url']
+        evi_list.append([title_str, evi_str, url_str])
 
     try:
         # Log answer on Azure Blob Storage
@@ -1226,7 +1231,6 @@ def chat(claim, history, sources):
                 "claim": claim,
                 "sources": sources,
                 "evidence": evi_list,
-                "url": url_of_evidence,
                 "answer": answer_output,
                 "time": timestamp,
             }
@@ -1254,7 +1258,7 @@ def main():
                     chatbot = gr.Chatbot(
                         value=[(None, init_prompt)],
                         show_copy_button=True, show_label=False, elem_id="chatbot", layout="panel",
-                        avatar_images=(None, "assets/averitec.png")
+                        avatar_images = (None, "assets/averitec.png")
                     )  # avatar_images=(None, "https://i.ibb.co/YNyd5W2/logo4.png"),
 
                     with gr.Row(elem_id="input-message"):