commented_small model

app.py

@@ -15,24 +15,24 @@ 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=256, 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'\[.*?\]'
-    redacted_text = re.sub(pattern, '[redacted]', predicted_title)
-    return redacted_text
+# 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=256, 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'\[.*?\]'
+#     redacted_text = re.sub(pattern, '[redacted]', predicted_title)
+#     return redacted_text
 
 def mask_generation(text, model=model_large, tokenizer=tokenizer_large):
-    if len(text) < 200:
+    if len(text) < 90:
         text = text + '.'
-        return small(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))