Update main.py

main.py

@@ -13,7 +13,7 @@ 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=len(text)):
+def mask_generation(text,model=model_small,tokenizer=tokenizer_small,max_len=128,max_len2):
     import re
     inputs = ["Mask Generation: " + text]
     inputs = tokenizer(inputs, max_length=max_len, truncation=True, return_tensors="pt")
@@ -39,7 +39,7 @@ async def mask_input(query):
     if(len(query)<90):
             output = mask_generation(query)
     else:
-        output = mask_generation(query,model_large,tokenizer_large,512,512)
+        output = mask_generation(query,model_large,tokenizer_large,512,len(query))
     return {"data" : output}
 
 if __name__ == '__main__':