from transformers import BertTokenizer, BertForQuestionAnswering
import torch

# Load pre-trained model and tokenizer
model_name = "google/bert_uncased_L-12_H-768_A-12"
tokenizer = BertTokenizer.from_pretrained(model_name)
model = BertForQuestionAnswering.from_pretrained(model_name)

def answer_question(question, context):
    # Tokenize input question and context
    inputs = tokenizer.encode_plus(question, context, add_special_tokens=True, return_tensors="pt")
    
    input_ids = inputs["input_ids"].tolist()[0]

    # Get the model's prediction
    outputs = model(**inputs)
    answer_start_scores = outputs.start_logits
    answer_end_scores = outputs.end_logits

    # Find the tokens with the highest `start` and `end` scores
    answer_start = torch.argmax(answer_start_scores)
    answer_end = torch.argmax(answer_end_scores) + 1

    # Convert tokens to the answer string
    answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(input_ids[answer_start:answer_end]))

    return answer

# Example usage
if __name__ == "__main__":
    context = """
    In the ancient land of India, amidst the vibrant city of Mathura, the divine being Krishna emerged, a manifestation of the supreme deity Vishnu. His birth, foretold by celestial signs, marked a pivotal moment in history, drawing the attention of both mortals and gods. The tyrant King Kamsa, driven by fear and prophecy, sought to thwart destiny by imprisoning Krishna's parents, Devaki and Vasudeva, and eliminating each of their newborns. However, Krishna's arrival defied all expectations. Guided by divine intervention, Vasudeva carried the infant Krishna across the treacherous river Yamuna to the safety of Vrindavan, where he would be nurtured away from Kamsa's grasp. In the pastoral setting of Vrindavan, Krishna's childhood unfolded amidst the pastoral beauty, surrounded by cowherds, gopis, and the enchanting melodies of his flute. His playful nature endeared him to all, as he engaged in divine acts known as leelas, showcasing his extraordinary abilities. Legends abound of his triumph over the serpent Kaliya, purifying the poisoned waters of the Yamuna, and his mischievous escapades with butter-stealing antics that charmed the hearts of the gopis. As Krishna matured, his purpose became clear—to challenge the oppressive rule of Kamsa and restore righteousness. Returning to Mathura, he confronted Kamsa in a series of daring feats, ultimately vanquishing the tyrant and freeing the people from his tyranny. Yet, Krishna's journey was far from over. His life was a tapestry of adventures, from battling demons to guiding the righteous in the pursuit of dharma. One of the most profound episodes in Krishna's life unfolded on the battlefield of Kurukshetra, where he served as Arjuna's charioteer and imparted the timeless wisdom of the Bhagavad Gita. In this sacred dialogue, Krishna expounded on duty, righteousness, and devotion, offering guidance that transcends time and resonates with seekers seeking enlightenment. Krishna's divine leelas, immortalized in myth and scripture, continue to captivate hearts and minds, inspiring devotion and love.
    """

    question = "Who is krishna ?"

    # Get the answer
    answer = answer_question(question, context)
    print(f"Question: {question}")
    print(f"Answer: {answer}")