Update README.md

README.md

@@ -26,11 +26,12 @@ trained_tokenizer_path = 'ZhangCheng/T5-Base-Fine-Tuned-for-Question-Generation'
 
 class QuestionGeneration:
 
-    def __init__(self, model_dir=None):
+    def __init__(self):
         self.model = T5ForConditionalGeneration.from_pretrained(trained_model_path)
         self.tokenizer = T5Tokenizer.from_pretrained(trained_tokenizer_path)
         self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         self.model = self.model.to(self.device)
+        self.model.eval()
 
     def generate(self, answer:str, context:str):
         input_text = '<answer> %s <context> %s ' % (answer, context)
@@ -40,7 +41,6 @@ class QuestionGeneration:
         )
         input_ids = encoding['input_ids'].to(self.device)
         attention_mask = encoding['attention_mask'].to(self.device)
-        self.model.eval()
         beam_outputs = self.model.generate(
             input_ids = input_ids,
             attention_mask = attention_mask