Update app.py

app.py

@@ -134,6 +134,17 @@ with torch.no_grad():
 tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])  # Convert input ids back to tokens
 predictions = torch.argmax(logits, dim=2)
 
+# Define labels
+id2label = {
+    0: "No binding site",
+    1: "Binding site"
+}
+
+# Print the predicted labels for each token
+for token, prediction in zip(tokens, predictions[0].numpy()):
+    if token not in ['<pad>', '<cls>', '<eos>']:
+        print((token, id2label[prediction]))
+        
 # debug result
 dubug_result = predictions  #class_weights