Update generator.py

generator.py

@@ -46,19 +46,23 @@ hfmodel, hftokenizer, tok, model = load_model()
 nlp = pipeline("e2e-qg")
 def run_model(input_string, **generator_args):
   generator_args = {
-  "max_length": 256,
+  "max_model_length": 256,
   "num_beams": 4,
-  # "length_penalty": 1.5,
+  "length_penalty": 1.5,
   "no_repeat_ngram_size": 2,
   "early_stopping": False,
   }
   # tokenizer = att.from_pretrained("ThomasSimonini/t5-end2end-question-generation")
-  output = nlp(input_string)
-  # input_string = "generate questions: " + input_string + " </s>"
-  # input_ids = hftokenizer.encode(input_string, return_tensors="pt")
-  # res = hfmodel.generate(input_ids, **generator_args)
-  # output = hftokenizer.batch_decode(res, skip_special_tokens=True)
-  # output = [item.split("<sep>") for item in output]
+  # output = nlp(input_string)
+    
+  input_string = "generate questions: " + input_string + " </s>"
+    
+  inputs = hftokenizer.encode(input_string, return_tensors="pt")
+    
+  res = model.generate(input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'], **generate_kwargs)
+  output = hftokenizer.decode(res[0], skip_special_tokens=True)
+  output = [item.split("<sep>") for item in output]
+  output = [o.strip() for o in output[:-1]]
   return output