Update app.py

app.py

@@ -5,13 +5,20 @@ from transformers import AutoModelForSequenceClassification
 tokenizer = BertTokenizer.from_pretrained('huolongguo10/check_sec')
 model = AutoModelForSequenceClassification.from_pretrained('huolongguo10/check_sec', num_labels=2)
 import torch
-def check(text):
-    inputs = tokenizer(text, return_tensors="pt")
+def check_each(text):
+    inputs = tokenizer(text, return_tensors="pt",max_length=512)
     with torch.no_grad():
         logits = model(**inputs).logits
     predicted_class_id = logits.argmax().item()
     print(f'{logits.argmax().item()}:{text}')
     return 'secure' if predicted_class_id==0 else 'insecure'
+def check(text):
+    t=text
+    while len(t)>512:
+        t=t[0:511]
+        if check_each(t)=='insecure':
+            return 'insecure'
+    return check_each(t)
 with gr.Blocks() as demo:
     text = gr.Textbox(label="Text")
     output = gr.Textbox(label="Output Box")