Update main.py

main.py

@@ -4,11 +4,16 @@ os.environ["HF_HOME"] = "/.cache"
 import re
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-model_dir = 'edithram23/Redaction_Personal_info_v1'
-tokenizer = AutoTokenizer.from_pretrained(model_dir)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
+model_dir_small = 'edithram23/Redaction'
+tokenizer_small = AutoTokenizer.from_pretrained(model_dir)
+model_small = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
 
-def mask_generation(text):
+
+model_dir_large = 'edithram23/Redaction_Personal_info_v1'
+tokenizer_large = AutoTokenizer.from_pretrained(model_dir)
+model_large = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
+
+def mask_generation(text,model=model_small,tokenizer=tokenizer_small):
     import re
     inputs = ["Mask Generation: " + text]
     inputs = tokenizer(inputs, max_length=512, truncation=True, return_tensors="pt")
@@ -31,7 +36,10 @@ async def hello():
 
 @app.post("/mask")
 async def mask_input(query):
-    output = mask_generation(query)
+    if(len(query)<90):
+            output = mask_generation(query)
+    else:
+        output = mask_generation(query,model_large,tokenizer_large)
     return {"data" : output}
 
 if __name__ == '__main__':