Update main.py

main.py

@@ -13,8 +13,9 @@ model_dir_large = 'edithram23/Redaction_Personal_info_v1'
 tokenizer_large = AutoTokenizer.from_pretrained(model_dir_large)
 model_large = AutoModelForSeq2SeqLM.from_pretrained(model_dir_large)
 
-def mask_generation(text,model=model_small,tokenizer=tokenizer_small,max_len=128,max_len2):
+def mask_generation(text,model=model_small,tokenizer=tokenizer_small,max_len=128):
     import re
+    max_len2=len(text)
     inputs = ["Mask Generation: " + text]
     inputs = tokenizer(inputs, max_length=max_len, truncation=True, return_tensors="pt")
     output = model.generate(**inputs, num_beams=8, do_sample=True, max_length=max_len2)