Removed bugs

Questgen/main.py

@@ -1,8 +1,9 @@
-import numpy as np # linear algebra
-import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)
+import numpy as np 
+import pandas as pd
 import time
 import torch
 from transformers import T5ForConditionalGeneration,T5Tokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import random
 import spacy
 import zipfile
@@ -43,7 +44,7 @@ def __init__(self):
         # model.eval()
         self.device = device
         self.model = model
-        self.nlp = spacy.load('en_core_web_sm', quiet=True)
+        self.nlp = spacy.load('en_core_web_sm')
 
         self.s2v = Sense2Vec().from_disk('s2v_old')
 
@@ -165,7 +166,7 @@ def paraphrase(self,payload):
             num_beams=50,
             num_return_sequences=num,
             no_repeat_ngram_size=2,
-            early_stopping=True
+            early_stopping=True,
             )
 
 #         print ("\nOriginal Question ::")
@@ -246,9 +247,18 @@ def predict_boolq(self,payload):
 
 class AnswerPredictor:
 
-    def __init__(self):
-        self.tokenizer = T5Tokenizer.from_pretrained('t5-large', model_max_length=512)
-        model = T5ForConditionalGeneration.from_pretrained('Parth/boolean')
+    def __init__(self, model_name = "T5"):
+        self.model_name = model_name
+        if self.model_name == "T5":
+            self.tokenizer = T5Tokenizer.from_pretrained('t5-base', model_max_length=512)
+            model = T5ForConditionalGeneration.from_pretrained('Parth/boolean')
+            # print(model, self.tokenizer)
+
+        if self.model_name == "BERT":
+            from transformers import AutoModelForQuestionAnswering, AutoTokenizer
+            model = AutoModelForQuestionAnswering.from_pretrained("Falconsai/question_answering")
+            self.tokenizer = AutoTokenizer.from_pretrained("Falconsai/question_answering", model_max_length=512)
+
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         model.to(device)
         # model.eval()
@@ -269,20 +279,28 @@ def greedy_decoding (inp_ids,attn_mask,model,tokenizer):
 
     def predict_answer(self,payload):
         answers = []
-        inp = {
-                "input_text": payload.get("input_text"),
-                "input_question" : payload.get("input_question")
-            }
+        context = payload["input_text"]
+
         for ques in payload.get("input_question"):
-
-            context = inp["input_text"]
             question = ques
-            input = "question: %s <s> context: %s </s>" % (question,context)
-
-            encoding = self.tokenizer.encode_plus(input, return_tensors="pt")
-            input_ids, attention_masks = encoding["input_ids"].to(self.device), encoding["attention_mask"].to(self.device)
-            greedy_output = self.model.generate(input_ids=input_ids, attention_mask=attention_masks, max_length=256)
-            Question =  self.tokenizer.decode(greedy_output[0], skip_special_tokens=True,clean_up_tokenization_spaces=True)
-            answers.append(Question.strip().capitalize())
+            input =  "question: %s <s> context: %s </s>" % (question,context)
+
+            inputs = self.tokenizer(question, context, return_tensors="pt")
+            if self.model_name == "T5":
+                encoding = self.tokenizer.encode_plus(input, return_tensors="pt")
+                input_ids, attention_masks = encoding["input_ids"].to(self.device), encoding["attention_mask"].to(self.device)
+                greedy_output = self.model.generate(input_ids=input_ids, attention_mask=attention_masks, max_length=512)
+                Answer =  self.tokenizer.decode(greedy_output[0], skip_special_tokens=True,clean_up_tokenization_spaces=True)
+                output = Answer.strip().capitalize()
+                answers.append(output)
+            break
+            if self.model_name == "BERT":
+                with torch.no_grad():
+                    outputs = self.model(**inputs)
+                answer_start_index = outputs.start_logits.argmax()
+                answer_end_index = outputs.end_logits.argmax()
+                predict_answer_tokens = inputs.input_ids[0, answer_start_index : answer_end_index + 1]
+                answer = self.tokenizer.decode(predict_answer_tokens)
+                answers.append(answer.strip().capitalize())
 
         return answers

README.md

@@ -188,8 +188,10 @@ pprint (output)
 </details>
 
 ### 2.5 Question Answering (Simple)
+BERT model is also added for Question Answering for faster inference and shorter answers. Default is T5 Large. 
+A list of questions can be passed to receive a list of answers.
 ```
-answer = main.AnswerPredictor()
+answer = main.AnswerPredictor(model_name="BERT")
 payload3 = {
     "input_text" : '''Sachin Ramesh Tendulkar is a former international cricketer from 
               India and a former captain of the Indian national team. He is widely regarded 
@@ -210,6 +212,7 @@ Sachin ramesh tendulkar is a former international cricketer from india and a for
 </details>
 
 ### 2.6 Question Answering (Boolean)
+
 ```
 payload4 = {
     "input_text" : '''Sachin Ramesh Tendulkar is a former international cricketer from 
@@ -233,6 +236,7 @@ Yes, sachin tendulkar is a former cricketer.
 
 For maintaining meaningfulness in Questions, Questgen uses Three T5 models. One for Boolean Question generation, one for MCQs, FAQs, Paraphrasing and one for answer generation.
 
+
 ### Online Demo website.
 https://questgen.ai/
 

setup.py

@@ -14,7 +14,7 @@
             'strsim==0.0.3',
             'six==1.16.0',
             'networkx==3.1',
-            'numpy==1.22.4',
+            'numpy',
             'scipy==1.10.1',
             'scikit-learn==1.2.2',
             'unidecode==1.3',
@@ -23,7 +23,7 @@
             'pytz==2022.7.1',
             'python-dateutil==2.8.2',
             'flashtext==2.7',
-            'pandas==1.5.3',
+            'pandas',
             'sentencepiece==0.1.99'
       ],
       package_data={'Questgen': ['questgen.py', 'mcq.py', 'train_gpu.py', 'encoding.py']}