from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline
import numpy as np
import re

class QuestionGeneration(object):

    def __init__(self):
        self.tokenizer_t5 = AutoTokenizer.from_pretrained("question-generation-squad-t5-large")
        #self.model_t5 = AutoModelForSeq2SeqLM.from_pretrained("question-generation-squad-t5-large")
        self.model_t5 = AutoModelForSeq2SeqLM.from_pretrained("question-generation-squad-t5-large").cuda()

        self.tokenizer_bart = AutoTokenizer.from_pretrained("question-generation-squad-bart-large")
        #self.model_bart = AutoModelForSeq2SeqLM.from_pretrained("question-generation-squad-bart-large")
        self.model_bart = AutoModelForSeq2SeqLM.from_pretrained("question-generation-squad-bart-large").cuda()

        model_name = "a-ware/roberta-large-squadv2"
        #self.nlp = pipeline('question-answering', model=model_name, tokenizer=model_name)
        self.nlp = pipeline('question-answering', model=model_name, tokenizer=model_name, device=0)

    def generate_questions(self, text):
        paragraphs = ["generate questions: " + paragraph for paragraph in re.split(r'\s*?\n',text) if paragraph!='']
        question_score = []
        for paragraph in paragraphs:
            #input_ids = self.tokenizer_t5(paragraph, return_tensors="pt").input_ids
            input_ids = self.tokenizer_t5(paragraph, return_tensors="pt").input_ids.cuda()
            outputs = self.model_t5.generate(input_ids,
                max_length=50, 
                num_beams=5, 
                no_repeat_ngram_size=2, 
                early_stopping=True,
                output_scores= True, return_dict_in_generate=True
            )
            question_t5= self.tokenizer_t5.decode(outputs.sequences[0], skip_special_tokens=True)
            #score_t5 = np.exp(outputs.sequences_scores[0].numpy())
            score_t5 = np.exp(outputs.sequences_scores[0].cpu().numpy())
            #print(question_t5,score_t5)
            question_score.append((question_t5,score_t5))
        question_t5,score_t5 = sorted(question_score, key=lambda x: x[1], reverse=True)[0]

        question_score = []
        QA_input = {'question': question_t5 ,'context': text}
        res = self.nlp(QA_input)
        answer_t5 = res['answer']
        answer_t5_score = res['score']

        for paragraph in paragraphs:
            #input_ids = self.tokenizer_bart(paragraph, return_tensors="pt").input_ids
            input_ids = self.tokenizer_bart(paragraph, return_tensors="pt").input_ids.cuda()
            outputs = self.model_bart.generate(input_ids,
                max_length=50, 
                num_beams=5, 
                no_repeat_ngram_size=2, 
                num_return_sequences=5, 
                early_stopping=True,
                output_scores= True, return_dict_in_generate=True
            )
            question_bart= self.tokenizer_bart.decode(outputs.sequences[0], skip_special_tokens=True)
            #score_bart = np.exp(outputs.sequences_scores[0].numpy())
            score_bart = np.exp(outputs.sequences_scores[0].cpu().numpy())
            i=1
            while question_bart == question_t5 and i < 5: ## Generate different questions from the one generated by T5 model
                question_bart= self.tokenizer_bart.decode(outputs.sequences[i], skip_special_tokens=True)
                #score_bart = np.exp(outputs.sequences_scores[i].numpy())
                score_bart = np.exp(outputs.sequences_scores[i].cpu().numpy())
                i+=1
            #print(question_bart,score_bart)
            question_score.append((question_bart,score_bart))
        question_bart,score_bart = sorted(question_score, key=lambda x: x[1], reverse=True)[0]

        QA_input = {'question': question_bart ,'context': text}
        res = self.nlp(QA_input)
        answer_bart = res['answer']
        answer_bart_score = res['score']

        return {
                "t5":{
                    "question":question_t5, 
                    "answer": answer_t5, 
                    "score":score_t5.astype(float), 
                    "answer_score": answer_t5_score
                    }, 
                "bart":{
                    "question":question_bart, 
                    "answer": answer_bart, 
                    "score":score_bart.astype(float), 
                    "answer_score": answer_bart_score
                    }
                }

