Update app.py

app.py

@@ -295,6 +295,10 @@ class_weights = compute_class_weight(class_weight='balanced', classes=classes, y
 accelerator = Accelerator()
 class_weights = torch.tensor(class_weights, dtype=torch.float32).to(accelerator.device)
 
+# Define labels and model
+id2label = {0: "No binding site", 1: "Binding site"}
+label2id = {v: k for k, v in id2label.items()}
+
 '''
 # inference
 # Path to the saved LoRA model
@@ -323,14 +327,13 @@ 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>']: