Extractive Question Answering With HuggingFace Using PyTorch and W&B

import json
import pandas as pd
import torch
from transformers import (
	BertTokenizerFast,
        BertForQuestionAnswering,
        TrainingArguments,
        Trainer,
 )
from sklearn.model_selection import train_test_split
from datasets import Dataset, DatasetDict

data_path = "/kaggle/input/covidqa-dataset/COVID-QA.json"
with open(data_path, "r") as f:
  data = json.load(f)

questions = []
answers = []
contexts = []

for entry in data['data']:
 for paragraph in entry['paragraphs']:
 context = paragraph['context']
  for qa in paragraph['qas']:
    questions.append(qa['question'])
    answers.append(qa['answers'][0]['text'])
    contexts.append(context)



df = pd.DataFrame({
  'question': questions,
  'answer': answers,
  'context': contexts
})

dataset = Dataset.from_pandas(df)
tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")

def tokenize(batch):
    tokenized_batch = tokenizer(batch["question"], batch["context"],
    max_length=512,
    padding="max_length",
    truncation=True,
 .  return_offsets_mapping=True,
    return_token_type_ids=True)

    answer_starts = []
    answer_ends = []

    for i, context in enumerate(batch["context"]):              
       answer_start = context.find(batch["answer"][i])   
       answer_end = answer_start + len(batch["answer"][i])   
       answer_starts.append(answer_start)
       answer_ends.append(answer_end)

   tokenized_batch["answer_start"] = answer_starts
   tokenized_batch["answer_end"] = answer_ends

   return tokenized_batch

tokenized_dataset = dataset.map(tokenize, batched=True)

def prepare_train_features(example):
    start_position = example["input_ids"].index(tokenizer.cls_token_id)
    end_position = example["input_ids"].index(tokenizer.sep_token_id)

    found_start = False
    found_end = False
    for i, (offset_start, offset_end) in enumerate(example["offset_mapping"]):
        if not found_start and offset_start == example["answer_start"]:
            start_position = i
            found_start = True
        if not found_end and offset_end == example["answer_end"]:
            end_position = i
            found_end = True
        if found_start and found_end:
            break

    if not found_start or not found_end:
        start_position = -1
        end_position = -1

    example["start_positions"] = start_position
    example["end_positions"] = end_position
    return example

prepared_dataset = tokenized_dataset.map(prepare_train_features, batched=False)

def filter_invalid_examples(example):
    return example["start_positions"] != -1 and example["end_positions"] != -1

filtered_dataset = prepared_dataset.filter(filter_invalid_examples, batched=False)

train_indices, eval_indices = train_test_split(list(range(len(filtered_dataset))), test_size=0.1, random_state=42)
train_dataset = filtered_dataset.select(train_indices)
eval_dataset = filtered_dataset.select(eval_indices)

def convert_to_tensors(example):
    example["input_ids"] = torch.tensor(example["input_ids"], dtype=torch.long)
    example["attention_mask"] = torch.tensor(example["attention_mask"], dtype=torch.long)
    return example

train_dataset = train_dataset.map(convert_to_tensors)
eval_dataset = eval_dataset.map(convert_to_tensors)

dataset_dict = DatasetDict({"train": train_dataset, "eval": eval_dataset})

model = BertForQuestionAnswering.from_pretrained("bert-base-uncased")

training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=64,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir='./logs',
    logging_steps=10,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    save_total_limit=2,
    fp16=True,
    load_best_model_at_end=True,
    report_to="wandb",  # Enable logging to W&B
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset_dict["train"],
    eval_dataset=dataset_dict["eval"],
)

trainer.train()

model.save_pretrained("trained_model")
model = BertForQuestionAnswering.from_pretrained("trained_model")

def get_answer(question, context):
    inputs = tokenizer.encode_plus(question, context, return_tensors="pt")
    start_logits, end_logits = model(**inputs).values()

    start_index_and_logits = torch.argmax(start_logits, dim=1).item(), start_logits[0].max().item()
    end_index_and_logits = torch.argmax(end_logits, dim=1).item(), end_logits[0].max().item()

    if end_index_and_logits[0] >= start_index_and_logits[0]:
        start_index, end_index = start_index_and_logits[0], end_index_and_logits[0]
    else:
        if start_index_and_logits[1] > end_index_and_logits[1]:
            start_index, end_index = start_index_and_logits[0], start_index_and_logits[0]
        else:
            start_index, end_index = end_index_and_logits[0], end_index_and_logits[0]

    answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][start_index:end_index+1]))
    return answer

question1 = "What is the incubation period of the novel coronavirus?"
context1 = "The novel coronavirus, also known as COVID-19, has an incubation period ranging from 1 to 14 days, with the majority of cases showing symptoms around 5 days after exposure."
answer1 = get_answer(question1, context1)
print("Answer 1:", answer1)

question2 = "What are the common symptoms of COVID-19?"
context2 = "COVID-19 symptoms can vary widely and may include fever, cough, shortness of breath, fatigue, body aches, and loss of taste or smell. Some people may also experience gastrointestinal symptoms like nausea, vomiting, and diarrhea."
answer2 = get_answer(question2, context2)
print("Answer 2:", answer2)

question3 = "How is the COVID-19 virus primarily transmitted?"
context3 = "COVID-19 is primarily transmitted through respiratory droplets produced when an infected person coughs, sneezes, or talks. These droplets can land in the mouths or noses of people who are nearby or be inhaled into their lungs. The virus can also spread through close personal contact, such as touching or shaking hands, or by touching a surface or object that has the virus on it and then touching their own mouth, nose, or eyes."
answer3 = get_answer(question3, context3)
print("Answer 3:", answer3)

question4 = "Is a vaccine for COVID-19 available?"
context4 = "As of September 2021, several COVID-19 vaccines have been developed and are being administered worldwide. These include the Pfizer-BioNTech, Moderna, AstraZeneca, and Johnson & Johnson vaccines. The availability and distribution of these vaccines vary depending on the country and region."
answer4 = get_answer(question4, context4)
print("Answer 4:", answer4)

question5 = "What is the role of ACE2 in COVID-19 infection?"
context5 = "ACE2, or angiotensin-converting enzyme 2, is a protein found on the surface of various human cells, including those in the lungs, heart, kidneys, and intestines. The novel coronavirus, SARS-CoV-2, responsible for COVID-19, uses ACE2 as a receptor to enter and infect human cells. The spike protein on the virus's surface binds to ACE2, allowing the virus to enter the cell and replicate."
answer5 = get_answer(question5, context5)
print("Answer 5:", answer5)

question6 = "Can COVID-19 be transmitted through food?"
context6 = "According to the World Health Organization (WHO) and the U.S. Food and Drug Administration (FDA), there is currently no evidence to suggest that COVID-19 can be transmitted through food or food packaging. The primary mode of transmission is through respiratory droplets from person to person. However, it is still essential to practice good food hygiene and wash your hands before preparing or consuming food."
answer6 = get_answer(question6, context6)
print("Answer 6:", answer6)