import torch
from transformers import AutoTokenizer, AutoModelForQuestionAnswering

# Load the pre-trained model and tokenizer
model_name = "microsoft/DialoGPT-large"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForQuestionAnswering.from_pretrained(model_name)

# Load the WikiNews dataset
with open("wikinews.txt", "r", encoding="utf-8") as f:
    dataset = f.readlines()

# Define a function to perform MRC on a given context and question
def ask_question(context, question):
    # Tokenize the input
    inputs = tokenizer(question, context, return_tensors="pt")

    # Perform MRC using the model
    start_logits, end_logits = model(**inputs).values()

    # Find the start and end indices of the answer in the context
    start_index = torch.argmax(start_logits)
    end_index = torch.argmax(end_logits) + 1

    # Convert the token indices to actual text
    answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][start_index:end_index]))

    return answer

# Loop through the dataset and ask a sample question for each article
for article in dataset:
    # Extract the title and content of the article
    title, content = article.strip().split("\n")
    print(f"Title: {title}")

    # Ask a sample question about the article
    question = "What is this article about?"
    answer = ask_question(content, question)
    print(f"Answer: {answer}\n")
