Update main.py

main.py

@@ -4,15 +4,15 @@ os.environ["HF_HOME"] = "/.cache"
 import re
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-model_dir = 'edithram23/Redaction'
+model_dir = 'edithram23/Redaction_Personal_info_v1'
 tokenizer = AutoTokenizer.from_pretrained(model_dir)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
 
 def mask_generation(text):
     import re
     inputs = ["Mask Generation: " + text]
-    inputs = tokenizer(inputs, max_length=128, truncation=True, return_tensors="pt")
-    output = model.generate(**inputs, num_beams=8, do_sample=True, max_length=len(i))
+    inputs = tokenizer(inputs, max_length=512, truncation=True, return_tensors="pt")
+    output = model.generate(**inputs, num_beams=8, do_sample=True, max_length=512)
     decoded_output = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
     predicted_title = decoded_output.strip()
     pattern = r'\[.*?\]'