update

app.py

@@ -43,7 +43,7 @@ try:
 except Exception as e:
     pass
 
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
+# os.environ["TOKENIZERS_PARALLELISM"] = "false"
 account_url = os.environ["AZURE_ACCOUNT_URL"]
 credential = {
     "account_key":  os.environ['AZURE_ACCOUNT_KEY'],
@@ -93,30 +93,38 @@ LABEL = [
     "Conflicting Evidence/Cherrypicking",
 ]
 
-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",
-                                                                  tokenizer=veracity_tokenizer, model=bert_model).to(device)
-# Justification
-justification_tokenizer = BartTokenizer.from_pretrained('facebook/bart-large', add_prefix_space=True)
-bart_model = BartForConditionalGeneration.from_pretrained("facebook/bart-large")
-best_checkpoint = 'averitec/pretrained_models/bart_justifications_verdict-epoch=13-val_loss=2.03-val_meteor=0.28.ckpt'
-justification_model = JustificationGenerationModule.load_from_checkpoint(best_checkpoint, tokenizer=justification_tokenizer, model=bart_model).to(device)
-# ---------------------------------------------------------------------------
+if torch.cuda.is_available():
+    # # device
+    # device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+    # question generation
+    qg_tokenizer = BloomTokenizerFast.from_pretrained("bigscience/bloom-7b1")
+    qg_model = BloomForCausalLM.from_pretrained("bigscience/bloom-7b1", torch_dtype=torch.bfloat16)
+    # qg_model = BloomForCausalLM.from_pretrained("bigscience/bloom-7b1", 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)
+    # 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", tokenizer=veracity_tokenizer, model=bert_model)
+    # veracity_model = SequenceClassificationModule.load_from_checkpoint("averitec/pretrained_models/bert_veracity.ckpt", tokenizer=veracity_tokenizer, model=bert_model).to(device)
+
+
+    # Justification
+    justification_tokenizer = BartTokenizer.from_pretrained('facebook/bart-large', add_prefix_space=True)
+    bart_model = BartForConditionalGeneration.from_pretrained("facebook/bart-large")
+    best_checkpoint = 'averitec/pretrained_models/bart_justifications_verdict-epoch=13-val_loss=2.03-val_meteor=0.28.ckpt'
+    justification_model = JustificationGenerationModule.load_from_checkpoint(best_checkpoint, tokenizer=justification_tokenizer, model=bart_model)
+    # justification_model = JustificationGenerationModule.load_from_checkpoint(best_checkpoint, tokenizer=justification_tokenizer, model=bart_model).to(device)
+
 
 # Set up Gradio Theme
 theme = gr.themes.Base(
@@ -124,9 +132,9 @@ theme = gr.themes.Base(
     secondary_hue="red",
     font=[gr.themes.GoogleFont("Poppins"), "ui-sans-serif", "system-ui", "sans-serif"],
 )
-
 # ---------- Setting ----------
 
+
 class Docs:
     def __init__(self, metadata=dict(), page_content=""):
         self.metadata = metadata
@@ -184,6 +192,7 @@ class SequenceClassificationDataLoader(pl.LightningDataModule):
 
         return input_ids, attention_masks
 
+
     def quadruple_to_string(self, claim, question, answer, bool_explanation=""):
         if bool_explanation is not None and len(bool_explanation) > 0:
             bool_explanation = ", because " + bool_explanation.lower().strip()
@@ -200,91 +209,8 @@ class SequenceClassificationDataLoader(pl.LightningDataModule):
         )
 
 
-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,
-                                                               problem_type="single_label_classification")
-
-    device = "cuda:0" if torch.cuda.is_available() else "cpu"
-    trained_model = SequenceClassificationModule.load_from_checkpoint("averitec/pretrained_models/bert_veracity.ckpt",
-                                                                      tokenizer=tokenizer, model=bert_model).to(device)
-
-    dataLoader = SequenceClassificationDataLoader(
-        tokenizer=tokenizer,
-        data_file="this_is_discontinued",
-        batch_size=32,
-        add_extra_nee=False,
-    )
-
-    evidence_strings = []
-    for evidence in qa_evidence:
-        evidence_strings.append(
-            dataLoader.quadruple_to_string(claim, evidence.metadata["query"], evidence.metadata["answer"], ""))
-
-    if len(evidence_strings) == 0:  # If we found no evidence e.g. because google returned 0 pages, just output NEI.
-        pred_label = "Not Enough Evidence"
-        return pred_label
-
-    tokenized_strings, attention_mask = dataLoader.tokenize_strings(evidence_strings)
-    example_support = torch.argmax(
-        trained_model(tokenized_strings.to(device), attention_mask=attention_mask.to(device)).logits, axis=1)
-
-    has_unanswerable = False
-    has_true = False
-    has_false = False
-
-    for v in example_support:
-        if v == 0:
-            has_true = True
-        if v == 1:
-            has_false = True
-        if v in (2, 3,):  # TODO another hack -- we cant have different labels for train and test so we do this
-            has_unanswerable = True
-
-    if has_unanswerable:
-        answer = 2
-    elif has_true and not has_false:
-        answer = 0
-    elif not has_true and has_false:
-        answer = 1
-    else:
-        answer = 3
-
-    pred_label = LABEL[answer]
-
-    return pred_label
-
-
-def fever_veracity_prediction(claim, evidence):
-    tokenizer = RobertaTokenizer.from_pretrained('Dzeniks/roberta-fact-check')
-    model = RobertaForSequenceClassification.from_pretrained('Dzeniks/roberta-fact-check')
-
-    evidence_string = ""
-    for evi in evidence:
-        evidence_string += evi.metadata['title'] + evi.metadata['evidence'] + ' '
-
-    input_sequence = tokenizer.encode_plus(claim, evidence_string, return_tensors="pt")
-    with torch.no_grad():
-        prediction = model(**input_sequence)
-
-    label = torch.argmax(prediction[0]).item()
-    pred_label = LABEL[label]
-
-    return pred_label
-
-
 @spaces.GPU
 def 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,
-    #                                                            problem_type="single_label_classification")
-    #
-    # device = "cuda:0" if torch.cuda.is_available() else "cpu"
-    # trained_model = SequenceClassificationModule.load_from_checkpoint("averitec/pretrained_models/bert_veracity.ckpt",
-    #                                                                   tokenizer=tokenizer, model=bert_model).to(device)
-
     dataLoader = SequenceClassificationDataLoader(
         tokenizer=veracity_tokenizer,
         data_file="this_is_discontinued",
@@ -302,7 +228,8 @@ def veracity_prediction(claim, qa_evidence):
         return pred_label
 
     tokenized_strings, attention_mask = dataLoader.tokenize_strings(evidence_strings)
-    example_support = torch.argmax(veracity_model(tokenized_strings.to(device), attention_mask=attention_mask.to(device)).logits, axis=1)
+    example_support = torch.argmax(veracity_model(tokenized_strings.to(veracity_model.device), attention_mask=attention_mask.to(veracity_model.device)).logits, axis=1)
+    # example_support = torch.argmax(veracity_model(tokenized_strings.to(device), attention_mask=attention_mask.to(device)).logits, axis=1)
 
     has_unanswerable = False
     has_true = False
@@ -330,6 +257,7 @@ def veracity_prediction(claim, qa_evidence):
     return pred_label
 
 
+@spaces.GPU
 def extract_claim_str(claim, qa_evidence, verdict_label):
     claim_str = "[CLAIM] " + claim + " [EVIDENCE] "
 
@@ -359,43 +287,43 @@ def extract_claim_str(claim, qa_evidence, verdict_label):
     return claim_str
 
 
-def google_justification_generation(claim, qa_evidence, verdict_label):
+@spaces.GPU
+def justification_generation(claim, qa_evidence, verdict_label):
     #
-    claim_str = extract_claim_str(claim, qa_evidence, verdict_label)
-    claim_str.strip()
+    # claim_str = extract_claim_str(claim, qa_evidence, verdict_label)
+    claim_str = "[CLAIM] " + claim + " [EVIDENCE] "
 
-    device = "cuda:0" if torch.cuda.is_available() else "cpu"
-    tokenizer = BartTokenizer.from_pretrained('facebook/bart-large', add_prefix_space=True)
-    bart_model = BartForConditionalGeneration.from_pretrained("facebook/bart-large")
+    for evi in qa_evidence:
+        q_text = evi.metadata['query'].strip()
 
-    best_checkpoint = 'averitec/pretrained_models/bart_justifications_verdict-epoch=13-val_loss=2.03-val_meteor=0.28.ckpt'
-    trained_model = JustificationGenerationModule.load_from_checkpoint(best_checkpoint, tokenizer=tokenizer,
-                                                                       model=bart_model).to(device)
+        if len(q_text) == 0:
+            continue
 
-    pred_justification = trained_model.generate(claim_str, device=device)
+        if not q_text[-1] == "?":
+            q_text += "?"
 
-    return pred_justification.strip()
+        answer_strings = []
+        answer_strings.append(evi.metadata['answer'])
 
+        claim_str += q_text
+        for a_text in answer_strings:
+            if a_text:
+                if not a_text[-1] == ".":
+                    a_text += "."
+                claim_str += " " + a_text.strip()
 
-@spaces.GPU
-def justification_generation(claim, qa_evidence, verdict_label):
-    #
-    claim_str = extract_claim_str(claim, qa_evidence, verdict_label)
-    claim_str.strip()
+        claim_str += " "
 
-    # device = "cuda:0" if torch.cuda.is_available() else "cpu"
-    # tokenizer = BartTokenizer.from_pretrained('facebook/bart-large', add_prefix_space=True)
-    # bart_model = BartForConditionalGeneration.from_pretrained("facebook/bart-large")
+    claim_str += " [VERDICT] " + verdict_label
     #
-    # best_checkpoint = 'averitec/pretrained_models/bart_justifications_verdict-epoch=13-val_loss=2.03-val_meteor=0.28.ckpt'
-    # trained_model = JustificationGenerationModule.load_from_checkpoint(best_checkpoint, tokenizer=tokenizer,
-    #                                                                    model=bart_model).to(device)
+    claim_str.strip()
 
-    pred_justification = justification_model.generate(claim_str, device=device)
+    pred_justification = justification_model.generate(claim_str, device=justification_model.device)
+    # pred_justification = justification_model.generate(claim_str, device=device)
 
     return pred_justification.strip()
 
-
+@spaces.GPU
 def QAprediction(claim, evidence, sources):
     parts = []
     #
@@ -493,9 +421,9 @@ def prompt_question_generation(test_claim, speaker="they", topk=10):
              "\". Criticism includes questions like: "
     sentences = [prompt]
 
-    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)
+    inputs = qg_tokenizer(sentences, padding=True, return_tensors="pt").to(qg_model.device)
+    # 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 = qg_tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
     in_len = len(sentences[0])
@@ -629,96 +557,7 @@ def get_google_search_results(api_key, search_engine_id, google_search, sort_dat
     return search_results
 
 
-def averitec_search_michael(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"]
-
-    blacklist = [
-        "jstor.org",  # Blacklisted because their pdfs are not labelled as such, and clog up the download
-        "facebook.com",  # Blacklisted because only post titles can be scraped, but the scraper doesn't know this,
-        "ftp.cs.princeton.edu",  # Blacklisted because it hosts many large NLP corpora that keep showing up
-        "nlp.cs.princeton.edu",
-        "huggingface.co"
-    ]
-
-    blacklist_files = [  # Blacklisted some NLP nonsense that crashes my machine with OOM errors
-        "/glove.",
-        "ftp://ftp.cs.princeton.edu/pub/cs226/autocomplete/words-333333.txt",
-        "https://web.mit.edu/adamrose/Public/googlelist",
-    ]
-
-    # save to folder
-    store_folder = "averitec/data/store/retrieved_docs"
-    #
-    index = 0
-    questions = [q["question"] for q in generate_question]
-
-    # check the date of the claim
-    current_date = datetime.now().strftime("%Y-%m-%d")
-    sort_date = check_claim_date(current_date)  # check_date="2022-01-01"
-
-    #
-    search_strings = []
-    search_types = []
-
-    search_string_2 = string_to_search_query(claim, None)
-    search_strings += [search_string_2, claim, ]
-    search_types += ["claim", "claim-noformat", ]
-
-    search_strings += questions
-    search_types += ["question" for _ in questions]
-
-    # start to search
-    search_results = []
-    visited = {}
-    store_counter = 0
-    worker_stack = list(range(10))
-
-    retrieve_evidence = []
-
-    for this_search_string, this_search_type in zip(search_strings, search_types):
-        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)
-
-            for result in search_results:
-                link = str(result["link"])
-                domain = get_domain_name(link)
-
-                if domain in blacklist:
-                    continue
-                broken = False
-                for b_file in blacklist_files:
-                    if b_file in link:
-                        broken = True
-                if broken:
-                    continue
-                if link.endswith(".pdf") or link.endswith(".doc"):
-                    continue
-
-                if link in visited:
-                    store_file_path = visited[link]
-                else:
-                    store_counter += 1
-                    store_file_path = store_folder + "/search_result_" + str(index) + "_" + str(
-                        store_counter) + ".store"
-                    visited[link] = store_file_path
-
-                    while len(worker_stack) == 0:  # Wait for a worker to become available. Check every second.
-                        sleep(1)
-
-                    worker = worker_stack.pop()
-
-                    t = threading.Thread(target=get_and_store, args=(link, store_file_path, worker, worker_stack))
-                    t.start()
-
-                line = [str(index), claim, link, str(page_num), this_search_string, this_search_type, store_file_path]
-                retrieve_evidence.append(line)
-
-    return retrieve_evidence
-
-
+@spaces.GPU
 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"]
@@ -848,7 +687,7 @@ def generate_step2_reference_corpus(reference_file):
 
     return tokenized_corpus, prompt_corpus
 
-
+@spaces.GPU
 def decorate_with_questions(claim, retrieve_evidence, top_k=5):  # top_k=10, 100
     #
     reference_file = "averitec/data/train.json"
@@ -916,9 +755,9 @@ def decorate_with_questions(claim, retrieve_evidence, top_k=5):  # top_k=10, 100
         prompt = "\n\n".join(prompt_docs + [claim_prompt])
         sentences = [prompt]
 
-        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)
+        inputs = qg_tokenizer(sentences, padding=True, return_tensors="pt").to(qg_model.device)
+        # 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 = 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:
@@ -1034,8 +873,8 @@ def rerank_questions(claim, bm25_qas, topk=3):
         values.append([question, answer, source])
 
     if len(bm25_qas) > 0:
-        encoded_dict = rerank_tokenizer(strs_to_score, max_length=512, padding="longest", truncation=True,
-                                 return_tensors="pt").to(device)
+        encoded_dict = rerank_tokenizer(strs_to_score, max_length=512, padding="longest", truncation=True, return_tensors="pt").to(rerank_trained_model.device)
+        # 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']
@@ -1052,6 +891,7 @@ def rerank_questions(claim, bm25_qas, topk=3):
     return top3_qa_pairs
 
 
+@spaces.GPU
 def Googleretriever(query, sources):
     # ----- Generate QA pairs using AVeriTeC
     # step 1: generate questions for the query/claim using Bloom
@@ -1207,6 +1047,7 @@ def WikipediaDumpsretriever(claim):
 
     return results
 
+
 # ----------WikipediaAPIretriever---------
 def clean_str(p):
   return p.encode().decode("unicode-escape").encode("latin1").decode("utf-8")
@@ -1556,6 +1397,7 @@ def main():
         dropdown_samples.change(change_sample_questions, dropdown_samples, samples)
         demo.queue()
 
+    # demo.launch()
     demo.launch(share=True)
 
 

requirements.txt

@@ -1,11 +1,12 @@
 gradio
-nltk==3.8.1
+nltk
 rank_bm25
 accelerate
 trafilatura
-spacy==3.7.5
+spacy
 pytorch_lightning
 transformers==4.29.2
+SentencePiece
 datasets
 leven
 scikit-learn