Update app.py

app.py

@@ -15,9 +15,26 @@ model_dir_large = 'edithram23/Redaction_Personal_info_v1'
 tokenizer_large = AutoTokenizer.from_pretrained(model_dir_large)
 model_large = AutoModelForSeq2SeqLM.from_pretrained(model_dir_large)
 
+model_dir_small = 'edithram23/Redaction'
+tokenizer_small = AutoTokenizer.from_pretrained(model_dir_small)
+model_small = AutoModelForSeq2SeqLM.from_pretrained(model_dir_small)
+
+def small(text,model=model_small,tokenizer=tokenizer_small):
+    inputs = ["Mask Generation: " + text.lower()+'.']
+    inputs = tokenizer(inputs, max_length=512, truncation=True, return_tensors="pt")
+    output = model.generate(**inputs, num_beams=8, do_sample=True, max_length=len(text))
+    decoded_output = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
+    predicted_title = decoded_output.strip()
+    pattern = r'\[.*?\]'
+    # Replace all occurrences of the pattern with [redacted]
+    redacted_text = re.sub(pattern, '[redacted]', predicted_title)
+    return redacted_text
+
+
 def mask_generation(text,model=model_large,tokenizer=tokenizer_large):
-    if(len(text)<30):
+    if(len(text)<90):
         text = text+'.'
+        return small(text)
     inputs = ["Mask Generation: " + text.lower()+'.']
     inputs = tokenizer(inputs, max_length=512, truncation=True, return_tensors="pt")
     output = model.generate(**inputs, num_beams=8, do_sample=True, max_length=len(text))