Update main.py

main.py

@@ -6,12 +6,12 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
 model_dir_small = 'edithram23/Redaction'
 tokenizer_small = AutoTokenizer.from_pretrained(model_dir)
-model_small = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
+model_small = AutoModelForSeq2SeqLM.from_pretrained(model_dir_small)
 
 
 model_dir_large = 'edithram23/Redaction_Personal_info_v1'
 tokenizer_large = AutoTokenizer.from_pretrained(model_dir)
-model_large = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
+model_large = AutoModelForSeq2SeqLM.from_pretrained(model_dir_large)
 
 def mask_generation(text,model=model_small,tokenizer=tokenizer_small):
     import re