# Пример загрузки HotpotQA
dataset = load_dataset("hotpot_qa", "distractor")

# Пример загрузки Natural Questions
dataset = load_dataset("nq")

# Пример загрузки Quoref
dataset = load_dataset("quoref")

from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForQuestionAnswering, pipeline

# Загрузка набора данных HotpotQA
dataset = load_dataset("hotpot_qa", "distractor")

# Выбор первого примера из набора данных
example = dataset['train'][0]
context = example['context']
question = example['question']
answer = example['answer']

# Вывод примера из набора данных
print(f"Question: {question}")
print(f"Context: {context}")
print(f"Answer: {answer}")

# Имя модели mT5
model_name = "google/mT5-small"

# Загрузка токенизатора и модели
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForQuestionAnswering.from_pretrained(model_name)

# Создание пайплайна для задачи вопрос-ответ
qa_pipeline = pipeline("question-answering", model=model, tokenizer=tokenizer)

# Выполнение предсказания
result = qa_pipeline(question=question, context=context)

# Вывод результата
print(f"Predicted Answer: {result['answer']}")


# Примеры вопросов и контекста
examples = [
    {
        "context": "The capital of France is Paris.",
        "question": "What is the capital of France?",
        "answer": "Paris"
    },
    {
        "context": "The tallest mountain in the world is Mount Everest.",
        "question": "What is the tallest mountain in the world?",
        "answer": "Mount Everest"
    },
    {
        "context": "Python is a programming language that lets you work quickly and integrate systems more effectively.",
        "question": "What is Python?",
        "answer": "A programming language"
    },
    {
        "context": "Reader PocketBook",
        "question": "How on PocketBook?",
    }
]

# Выполнение предсказания для каждого примера
for example in examples:
    result = qa_pipeline(question=example["question"], context=example["context"])
    print(f"Question: {example['question']}")
    print(f"Context: {example['context']}")
    print(f"Answer: {example['answer']}")
    print(f"Predicted Answer: {result['answer']}")
    print()