from transformers import (
    T5ForConditionalGeneration,
    T5Tokenizer
)

import pandas as pd
import numpy as np

tokenizer = T5Tokenizer.from_pretrained("ashishkat/questionAnswer")
model = T5ForConditionalGeneration.from_pretrained("ashishkat/questionAnswer", return_dict=True)


def generate_answer(question, context):
    """Function gives the answer to the  question asked, given context

    question(str) : question asked by user
    context(str): Paragraph given by used 

    Returns:
        string: Answer to respective question asked
    """

    ## tokenizeing question + context at a same time 
    ## max length is 512, greater are removed, less are padded
    source_encoding = tokenizer(
            question,
            context,
            max_length = 512,  
            padding="max_length",
            truncation="only_second",
            return_attention_mask = True,
            return_tensors="pt",
            add_special_tokens=True
            )
    ## generating answer from model
    generate_ids = model.generate(
           input_ids = source_encoding["input_ids"],
           attention_mask = source_encoding["attention_mask"],
           max_length = 30,
           use_cache=True,
       )

    ## decoding the tokenized prediction
    pred = [
               tokenizer.decode(ids, skip_special_tokens=True) for ids in generate_ids
       ]

    return " ".join(pred)  ## returns the predicted string as answer




context = "The Normans (Norman: Nourmands; French: Normands; Latin: Normanni) were the people who \
        in the 10th and 11th centuries gave their name to Normandy, a region in France. They were\
        descended from Norse (\"Norman\" comes from \"Norseman\") raiders and pirates from Denmark,\
        Iceland and Norway who, under their leader Rollo, agreed to swear fealty to King Charles III\
        of West Francia. Through generations of assimilation and mixing with the native Frankish and\
        Roman-Gaulish populations, their descendants would gradually merge with the Carolingian-based\
        cultures of West Francia. The distinct cultural and ethnic identity of the Normans emerged initially\
        in the first half of the 10th century, and it continued to evolve over the succeeding centuries."

question = "From which countries did the Norse originate?"
# question = "Who was the Norse leader?" "rollo"
# actual_answer = "Denmark, Iceland and Norway"
answer = generate_answer(context, question)
print(answer)