adjust

app.py

@@ -18,7 +18,6 @@ def inference(input_text):
     inputs = tokenizer.batch_encode_plus(
                 [input_text],
                 max_length=512,
-                pad_to_max_length=True,
                 truncation=True,
                 padding="max_length",
                 return_tensors="pt",
@@ -27,7 +26,7 @@ def inference(input_text):
     with torch.no_grad():
         logits = model(**inputs).logits
     
-    predicted_class_id = logits.argmax().item()
+    predicted_class_id = logits.argmax(dim=-1).item()
     output = model.config.id2label[predicted_class_id]
     return output
 
@@ -36,8 +35,8 @@ demo = gr.Interface(
     inputs=gr.Textbox(label="Input Text", scale=2, container=False),
     outputs=gr.Textbox(label="Output Label"),
     examples = [
-        ["My last two weather pics from the storm on August 2nd. People packed up real fast after the temp dropped and winds picked up.", 1],
-        ["Lying Clinton sinking! Donald Trump singing: Let's Make America Great Again!", 0],
+        ["My last two weather pics from the storm on August 2nd. "],
+        ["Lying Clinton sinking! Donald Trump singing."],
         ],
     title="Tutorial: BERT-based Text Classificatioin",
     )