Update main.py

main.py

@@ -4,20 +4,16 @@ os.environ["HF_HOME"] = "/.cache"
 import re
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-model_dir_small = 'edithram23/Redaction'
-tokenizer_small = AutoTokenizer.from_pretrained(model_dir_small)
-model_small = AutoModelForSeq2SeqLM.from_pretrained(model_dir_small)
-
 
 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=64):
+def mask_generation(text,model=model_large,tokenizer=tokenizer_large):
     import re
-    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)
+    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)
     decoded_output = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
     predicted_title = decoded_output.strip()
     pattern = r'\[.*?\]'
@@ -36,10 +32,7 @@ async def hello():
 
 @app.post("/mask")
 async def mask_input(query):
-    if(len(query)<90):
-            output = mask_generation(query)
-    else:
-        output = mask_generation(query,model_large,tokenizer_large,512,len(query))
+    output = mask_generation(query)
     return {"data" : output}
 
 if __name__ == '__main__':