Update main.py

main.py

@@ -13,7 +13,7 @@ def mask_generation(text,model=model_large,tokenizer=tokenizer_large):
     import re
     inputs = ["Mask Generation: " + text+'.']
     inputs = tokenizer(inputs, max_length=512, truncation=True, return_tensors="pt")
-    output = model.generate(**inputs, num_beams=8, do_sample=True, max_length=512)
+    output = model.generate(**inputs, num_beams=8, do_sample=True, max_length=len(text))
     decoded_output = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
     predicted_title = decoded_output.strip()
     pattern = r'\[.*?\]'